Search in sources :

Example 6 with InboundDataClient

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

the class FnApiWindowMappingFn method loadIfNeeded.

private TargetWindowT loadIfNeeded(BoundedWindow mainWindow) {
    try {
        String processRequestInstructionId = idGenerator.getId();
        InstructionRequest processRequest = InstructionRequest.newBuilder().setInstructionId(processRequestInstructionId).setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId(registerIfRequired())).build();
        ConcurrentLinkedQueue<WindowedValue<KV<byte[], TargetWindowT>>> outputValue = new ConcurrentLinkedQueue<>();
        // Open the inbound consumer
        InboundDataClient waitForInboundTermination = beamFnDataService.receive(LogicalEndpoint.data(processRequestInstructionId, "write"), inboundCoder, outputValue::add);
        CompletionStage<InstructionResponse> processResponse = instructionRequestHandler.handle(processRequest);
        // Open the outbound consumer
        try (CloseableFnDataReceiver<WindowedValue<KV<byte[], BoundedWindow>>> outboundConsumer = beamFnDataService.send(LogicalEndpoint.data(processRequestInstructionId, "read"), outboundCoder)) {
            outboundConsumer.accept(WindowedValue.valueInGlobalWindow(KV.of(EMPTY_ARRAY, mainWindow)));
        }
        // Check to see if processing the request failed.
        MoreFutures.get(processResponse);
        waitForInboundTermination.awaitCompletion();
        WindowedValue<KV<byte[], TargetWindowT>> sideInputWindow = outputValue.poll();
        checkState(sideInputWindow != null, "Expected side input window to have been emitted by SDK harness.");
        checkState(sideInputWindow.getValue() != null, "Side input window emitted by SDK harness was a WindowedValue with no value in it.");
        checkState(sideInputWindow.getValue().getValue() != null, "Side input window emitted by SDK harness was a WindowedValue<KV<...>> with a null V.");
        checkState(outputValue.isEmpty(), "Expected only a single side input window to have been emitted by " + "the SDK harness but also received %s", outputValue);
        return sideInputWindow.getValue().getValue();
    } catch (Throwable e) {
        LOG.error("Unable to map main input window {} to side input window.", mainWindow, e);
        throw new IllegalStateException(e);
    }
}
Also used : InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) KV(org.apache.beam.sdk.values.KV) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) InstructionRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest) WindowedValue(org.apache.beam.sdk.util.WindowedValue) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue)

Example 7 with InboundDataClient

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

the class RemoteGrpcPortReadOperationTest method testStartAndAbort.

@Test
public void testStartAndAbort() throws Exception {
    InboundDataClient inboundDataClient = CompletableFutureInboundDataClient.create();
    when(beamFnDataService.receive(any(), Matchers.<Coder<WindowedValue<String>>>any(), any())).thenReturn(inboundDataClient);
    when(bundleIdSupplier.getId()).thenReturn(BUNDLE_ID);
    operation.start();
    verify(beamFnDataService).receive(eq(LogicalEndpoint.data(BUNDLE_ID, TRANSFORM_ID)), eq(CODER), consumerCaptor.capture());
    assertFalse(inboundDataClient.isDone());
    operation.abort();
    assertTrue(inboundDataClient.isDone());
    verify(bundleIdSupplier, times(1)).getId();
}
Also used : WindowedValue(org.apache.beam.sdk.util.WindowedValue) CompletableFutureInboundDataClient(org.apache.beam.sdk.fn.data.CompletableFutureInboundDataClient) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) Test(org.junit.Test)

Example 8 with InboundDataClient

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

the class RemoteGrpcPortReadOperationTest method testSuccessfulProcessing.

@Test
public void testSuccessfulProcessing() throws Exception {
    InboundDataClient inboundDataClient = CompletableFutureInboundDataClient.create();
    when(beamFnDataService.receive(any(), Matchers.<Coder<WindowedValue<String>>>any(), any())).thenReturn(inboundDataClient);
    when(bundleIdSupplier.getId()).thenReturn(BUNDLE_ID);
    operation.start();
    verify(beamFnDataService).receive(eq(LogicalEndpoint.data(BUNDLE_ID, TRANSFORM_ID)), eq(CODER), consumerCaptor.capture());
    Future<Void> operationFinish = Executors.newSingleThreadExecutor().submit(() -> {
        operation.finish();
        return null;
    });
    consumerCaptor.getValue().accept(valueInGlobalWindow("ABC"));
    consumerCaptor.getValue().accept(valueInGlobalWindow("DEF"));
    consumerCaptor.getValue().accept(valueInGlobalWindow("GHI"));
    // Purposefully sleep to show that the operation is still not done until the finish signal
    // is completed.
    Thread.sleep(100L);
    assertFalse(operationFinish.isDone());
    inboundDataClient.complete();
    operationFinish.get();
    verify(bundleIdSupplier, times(1)).getId();
    assertThat(testReceiver.outputElems, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI")));
    // Ensure that the old bundle id is cleared.
    when(bundleIdSupplier.getId()).thenReturn(BUNDLE_ID_2);
    operation.start();
    verify(beamFnDataService).receive(eq(LogicalEndpoint.data(BUNDLE_ID_2, TRANSFORM_ID)), eq(CODER), consumerCaptor.capture());
}
Also used : WindowedValue(org.apache.beam.sdk.util.WindowedValue) CompletableFutureInboundDataClient(org.apache.beam.sdk.fn.data.CompletableFutureInboundDataClient) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) Test(org.junit.Test)

Example 9 with InboundDataClient

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

the class SdkHarnessClientTest method testBundleFinalizationCallback.

@Test
public void testBundleFinalizationCallback() 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().setRequiresFinalization(true).build();
    String bundleId;
    try (ActiveBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), Collections.emptyMap(), (request) -> {
        throw new UnsupportedOperationException();
    }, mockProgressHandler, mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler)) {
        bundleId = activeBundle.getId();
        processBundleResponseFuture.complete(InstructionResponse.newBuilder().setProcessBundle(response).build());
    }
    verify(mockProgressHandler).onCompleted(response);
    verify(mockFinalizationHandler).requestsFinalization(bundleId);
    verifyZeroInteractions(mockCheckpointHandler, 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)

Example 10 with InboundDataClient

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

the class SdkHarnessClientTest method handleCleanupWithStateWhenAwaitingOnClosingOutputReceivers.

@Test
public void handleCleanupWithStateWhenAwaitingOnClosingOutputReceivers() 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(mockOutputReceiver).awaitCompletion();
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) StateRequestHandler(org.apache.beam.runners.fnexecution.state.StateRequestHandler) 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) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) StateDelegator(org.apache.beam.runners.fnexecution.state.StateDelegator) Test(org.junit.Test)

Aggregations

InboundDataClient (org.apache.beam.sdk.fn.data.InboundDataClient)13 Test (org.junit.Test)12 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)10 InstructionResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse)9 CompletableFuture (java.util.concurrent.CompletableFuture)8 BundleProcessor (org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor)8 CloseableFnDataReceiver (org.apache.beam.sdk.fn.data.CloseableFnDataReceiver)8 ExecutionException (java.util.concurrent.ExecutionException)6 ExpectedException (org.junit.rules.ExpectedException)6 WindowedValue (org.apache.beam.sdk.util.WindowedValue)5 ProcessBundleResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse)4 StateDelegator (org.apache.beam.runners.fnexecution.state.StateDelegator)3 StateRequestHandler (org.apache.beam.runners.fnexecution.state.StateRequestHandler)3 ArrayList (java.util.ArrayList)2 Collection (java.util.Collection)2 CountDownLatch (java.util.concurrent.CountDownLatch)2 ExecutorService (java.util.concurrent.ExecutorService)2 LinkedBlockingQueue (java.util.concurrent.LinkedBlockingQueue)2 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)2 Elements (org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements)2