Examples with InboundDataClient org.apache.beam.sdk.fn.data.InboundDataClient used on opensource projects

Example 1 with InboundDataClient

use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.

the class SdkHarnessClientTest method handleCleanupWhenAwaitingOnClosingOutputReceivers.

@Test
public void handleCleanupWhenAwaitingOnClosingOutputReceivers() throws Exception {
    Exception testException = new Exception();
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    doThrow(testException).when(mockOutputReceiver).awaitCompletion();
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
            // Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
            // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
            // the response.
            // 
            // Currently there are no fields so there's nothing to check. This test is formulated
            // to match the pattern it should have if/when the response is meaningful.
            BeamFnApi.ProcessBundleResponse response = BeamFnApi.ProcessBundleResponse.getDefaultInstance();
            processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
        }
        fail("Exception expected");
    } catch (Exception e) {
        assertEquals(testException, e);
    }
}

Example 2 with InboundDataClient

use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.

the class SdkHarnessClientTest method handleCleanupWhenProcessingBundleFails.

@Test
public void handleCleanupWhenProcessingBundleFails() throws Exception {
    Exception testException = new Exception();
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
            processBundleResponseFuture.completeExceptionally(testException);
        }
        fail("Exception expected");
    } catch (ExecutionException e) {
        assertEquals(testException, e.getCause());
        verify(mockOutputReceiver).cancel();
        verifyNoMoreInteractions(mockOutputReceiver);
    }
}

Example 3 with InboundDataClient

use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.

the class SdkHarnessClientTest method handleCleanupWithStateWhenInputSenderFails.

@Test
public void handleCleanupWithStateWhenInputSenderFails() throws Exception {
    Exception testException = new Exception();
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    StateDelegator mockStateDelegator = mock(StateDelegator.class);
    StateDelegator.Registration mockStateRegistration = mock(StateDelegator.Registration.class);
    when(mockStateDelegator.registerForProcessBundleInstructionId(any(), any())).thenReturn(mockStateRegistration);
    StateRequestHandler mockStateHandler = mock(StateRequestHandler.class);
    when(mockStateHandler.getCacheTokens()).thenReturn(Collections.emptyList());
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)), mockStateDelegator);
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    doThrow(testException).when(mockInputSender).close();
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, mockProgressHandler)) {
        // We shouldn't be required to complete the process bundle response future.
        }
        fail("Exception expected");
    } catch (Exception e) {
        assertEquals(testException, e);
        verify(mockStateRegistration).abort();
        verify(mockOutputReceiver).cancel();
        verifyNoMoreInteractions(mockStateRegistration, mockOutputReceiver);
    }
}

Example 4 with InboundDataClient

use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.

the class SdkHarnessClientTest method testBundleCheckpointCallback.

@Test
public void testBundleCheckpointCallback() throws Exception {
    InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
    CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
    FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
    BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
    when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
    when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    BundleSplitHandler mockSplitHandler = mock(BundleSplitHandler.class);
    BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class);
    BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class);
    ProcessBundleResponse response = ProcessBundleResponse.newBuilder().addResidualRoots(DelayedBundleApplication.getDefaultInstance()).build();
    try (ActiveBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), Collections.emptyMap(), (request) -> {
        throw new UnsupportedOperationException();
    }, mockProgressHandler, mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler)) {
        processBundleResponseFuture.complete(InstructionResponse.newBuilder().setProcessBundle(response).build());
    }
    verify(mockProgressHandler).onCompleted(response);
    verify(mockCheckpointHandler).onCheckpoint(response);
    verifyZeroInteractions(mockFinalizationHandler, mockSplitHandler);
}

Example 5 with InboundDataClient

use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.

the class BeamFnDataGrpcServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.

@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
    LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    CountDownLatch waitForInboundElements = new CountDownLatch(1);
    for (int i = 0; i < 3; ++i) {
        String instructionId = Integer.toString(i);
        executorService.submit(() -> {
            ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID))).forDescriptor(service.getApiServiceDescriptor());
            StreamObserver<BeamFnApi.Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
            outboundObserver.onNext(elementsWithData(instructionId));
            waitForInboundElements.await();
            outboundObserver.onCompleted();
            return null;
        });
    }
    List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
    Collection<InboundDataClient> inboundDataClients = new ArrayList<>();
    for (int i = 0; i < 3; ++i) {
        BlockingQueue<WindowedValue<String>> serverInboundValue = new LinkedBlockingQueue<>();
        serverInboundValues.add(serverInboundValue);
        inboundDataClients.add(service.getDataService(WORKER_ID).receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
    }
    // Waiting for the client provides the necessary synchronization for the elements to arrive.
    for (InboundDataClient inboundDataClient : inboundDataClients) {
        inboundDataClient.awaitCompletion();
    }
    waitForInboundElements.countDown();
    for (int i = 0; i < 3; ++i) {
        assertThat(serverInboundValues.get(i), contains(valueInGlobalWindow("A" + i), valueInGlobalWindow("B" + i), valueInGlobalWindow("C" + i)));
    }
    assertThat(clientInboundElements, empty());
}