Search in sources :

Example 6 with BundleProcessor

use of org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) ExecutionException(java.util.concurrent.ExecutionException) Test(org.junit.Test)

Example 7 with BundleProcessor

use of org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) StateRequestHandler(org.apache.beam.runners.fnexecution.state.StateRequestHandler) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) StateDelegator(org.apache.beam.runners.fnexecution.state.StateDelegator) Test(org.junit.Test)

Example 8 with BundleProcessor

use of org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor in project beam by apache.

the class SdkHarnessClientTest method testCheckpointHappensAfterAnySplitCalls.

@Test
@SuppressWarnings("FutureReturnValueIgnored")
public void testCheckpointHappensAfterAnySplitCalls() throws Exception {
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    CompletableFuture<InstructionResponse> splitResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenAnswer(invocationOnMock -> {
        switch(invocationOnMock.<BeamFnApi.InstructionRequest>getArgument(0).getRequestCase()) {
            case PROCESS_BUNDLE:
                return processBundleResponseFuture;
            case PROCESS_BUNDLE_SPLIT:
                return splitResponseFuture;
            default:
                throw new IllegalArgumentException("Unexpected request " + invocationOnMock.<BeamFnApi.InstructionRequest>getArgument(0));
        }
    });
    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.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
    BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class);
    BundleSplitHandler mockSplitHandler = mock(BundleSplitHandler.class);
    BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class);
    RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), Collections.emptyMap(), StateRequestHandler.unsupported(), BundleProgressHandler.ignored(), mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler);
    BeamFnApi.ProcessBundleResponse response = ProcessBundleResponse.newBuilder().addResidualRoots(DelayedBundleApplication.newBuilder().setApplication(BundleApplication.newBuilder().setTransformId("test").build()).build()).build();
    BeamFnApi.ProcessBundleSplitResponse splitResponse = ProcessBundleSplitResponse.newBuilder().addChannelSplits(ChannelSplit.newBuilder().setTransformId("test2")).build();
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    // Correlating the request and response is owned by the underlying
    // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
    // the response.
    // 
    // Schedule the split response to come in after the bundle response and after the close call.
    activeBundle.split(0.5);
    executor.schedule(() -> processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build()), 1, TimeUnit.SECONDS);
    executor.schedule(() -> splitResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundleSplit(splitResponse).build()), 2, TimeUnit.SECONDS);
    activeBundle.close();
    InOrder inOrder = Mockito.inOrder(mockCheckpointHandler, mockSplitHandler);
    inOrder.verify(mockSplitHandler).split(eq(splitResponse));
    inOrder.verify(mockCheckpointHandler).onCheckpoint(eq(response));
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) InOrder(org.mockito.InOrder) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ProcessBundleSplitResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleSplitResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) CompletableFuture(java.util.concurrent.CompletableFuture) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) Test(org.junit.Test)

Example 9 with BundleProcessor

use of org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor in project beam by apache.

the class SdkHarnessClientTest method testProgressHandlerOnCompletedHappensAfterOnProgress.

@Test
@SuppressWarnings("FutureReturnValueIgnored")
public void testProgressHandlerOnCompletedHappensAfterOnProgress() throws Exception {
    CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
    CompletableFuture<InstructionResponse> progressResponseFuture = new CompletableFuture<>();
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenAnswer(invocationOnMock -> {
        switch(invocationOnMock.<BeamFnApi.InstructionRequest>getArgument(0).getRequestCase()) {
            case PROCESS_BUNDLE:
                return processBundleResponseFuture;
            case PROCESS_BUNDLE_PROGRESS:
                return progressResponseFuture;
            default:
                throw new IllegalArgumentException("Unexpected request " + invocationOnMock.<BeamFnApi.InstructionRequest>getArgument(0));
        }
    });
    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.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
    BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
    RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), mockProgressHandler);
    BeamFnApi.ProcessBundleResponse response = ProcessBundleResponse.newBuilder().putMonitoringData("test", ByteString.EMPTY).build();
    BeamFnApi.ProcessBundleProgressResponse progressResponse = ProcessBundleProgressResponse.newBuilder().putMonitoringData("test2", ByteString.EMPTY).build();
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    // Correlating the request and response is owned by the underlying
    // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
    // the response.
    // 
    // Schedule the progress response to come in after the bundle response and after the close call.
    activeBundle.requestProgress();
    executor.schedule(() -> processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build()), 1, TimeUnit.SECONDS);
    executor.schedule(() -> progressResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundleProgress(progressResponse).build()), 2, TimeUnit.SECONDS);
    activeBundle.close();
    InOrder inOrder = Mockito.inOrder(mockProgressHandler);
    inOrder.verify(mockProgressHandler).onProgress(eq(progressResponse));
    inOrder.verify(mockProgressHandler).onCompleted(eq(response));
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) InOrder(org.mockito.InOrder) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) CompletableFuture(java.util.concurrent.CompletableFuture) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) ProcessBundleProgressResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse) Test(org.junit.Test)

Example 10 with BundleProcessor

use of org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor 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);
}
Also used : ActiveBundle(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor.ActiveBundle) CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) Test(org.junit.Test)

Aggregations

BundleProcessor (org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor)26 Test (org.junit.Test)26 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)23 CloseableFnDataReceiver (org.apache.beam.sdk.fn.data.CloseableFnDataReceiver)14 CompletableFuture (java.util.concurrent.CompletableFuture)13 InstructionResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse)13 ProcessBundleResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse)12 WindowedValue (org.apache.beam.sdk.util.WindowedValue)11 HashMap (java.util.HashMap)10 ConcurrentHashMap (java.util.concurrent.ConcurrentHashMap)10 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)10 ExecutableStage (org.apache.beam.runners.core.construction.graph.ExecutableStage)10 FusedPipeline (org.apache.beam.runners.core.construction.graph.FusedPipeline)10 ExecutableProcessBundleDescriptor (org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.ExecutableProcessBundleDescriptor)10 Pipeline (org.apache.beam.sdk.Pipeline)10 BigEndianLongCoder (org.apache.beam.sdk.coders.BigEndianLongCoder)10 Coder (org.apache.beam.sdk.coders.Coder)10 KvCoder (org.apache.beam.sdk.coders.KvCoder)10 StringUtf8Coder (org.apache.beam.sdk.coders.StringUtf8Coder)10 Collection (java.util.Collection)9