Examples with BundleProcessor org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor used on opensource projects

Example 21 with BundleProcessor

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

the class SdkHarnessClientTest method testRegisterCachesBundleProcessors.

@Test
public void testRegisterCachesBundleProcessors() throws Exception {
    ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
    ProcessBundleDescriptor descriptor2 = ProcessBundleDescriptor.newBuilder().setId("descriptor2").build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
    BundleProcessor processor2 = sdkHarnessClient.getProcessor(descriptor2, remoteInputs);
    assertNotSame(processor1, processor2);
    // Ensure that caching works.
    assertSame(processor1, sdkHarnessClient.getProcessor(descriptor1, remoteInputs));
}

Example 22 with BundleProcessor

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

the class SdkHarnessClientTest method testNewBundleNoDataDoesNotCrash.

@Test
public void testNewBundleNoDataDoesNotCrash() throws Exception {
    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.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
    try (RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), BundleProgressHandler.ignored())) {
        // 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 = ProcessBundleResponse.getDefaultInstance();
        processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
    }
}

Example 23 with BundleProcessor

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

the class SdkHarnessClientTest method handleCleanupWithStateWhenProcessingBundleFails.

@Test
public void handleCleanupWithStateWhenProcessingBundleFails() 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);
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, mockProgressHandler)) {
            processBundleResponseFuture.completeExceptionally(testException);
        }
        fail("Exception expected");
    } catch (ExecutionException e) {
        assertEquals(testException, e.getCause());
        verify(mockStateRegistration).abort();
        verify(mockOutputReceiver).cancel();
        verifyNoMoreInteractions(mockStateRegistration, mockOutputReceiver);
    }
}

Example 24 with BundleProcessor

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

the class SdkHarnessClientTest method testNewBundleAndProcessElements.

@Test
public void testNewBundleAndProcessElements() throws Exception {
    SdkHarnessClient client = harness.client();
    BundleProcessor processor = client.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM)));
    Collection<WindowedValue<String>> outputs = new ArrayList<>();
    try (RemoteBundle activeBundle = processor.newBundle(Collections.singletonMap(SDK_GRPC_WRITE_TRANSFORM, RemoteOutputReceiver.of(FullWindowedValueCoder.of(LengthPrefixCoder.of(StringUtf8Coder.of()), Coder.INSTANCE), outputs::add)), BundleProgressHandler.ignored())) {
        FnDataReceiver<WindowedValue<?>> bundleInputReceiver = Iterables.getOnlyElement(activeBundle.getInputReceivers().values());
        bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("foo"));
        bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("bar"));
        bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("baz"));
    }
    // The bundle can be a simple function of some sort, but needs to be complete.
    assertThat(outputs, containsInAnyOrder(WindowedValue.valueInGlobalWindow("spam"), WindowedValue.valueInGlobalWindow("ham"), WindowedValue.valueInGlobalWindow("eggs")));
}

Example 25 with BundleProcessor

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

the class SdkHarnessClientTest method testProgressAndSplitCallsAreIgnoredWhenBundleIsComplete.

@Test
public void testProgressAndSplitCallsAreIgnoredWhenBundleIsComplete() throws Exception {
    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.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
    RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), BundleProgressHandler.ignored());
    // Correlating the request and response is owned by the underlying
    // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
    // the response.
    BeamFnApi.ProcessBundleResponse response = ProcessBundleResponse.getDefaultInstance();
    processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
    activeBundle.close();
    verify(fnApiControlClient).registerProcessBundleDescriptor(any(ProcessBundleDescriptor.class));
    verify(fnApiControlClient).handle(any(BeamFnApi.InstructionRequest.class));
    activeBundle.requestProgress();
    activeBundle.split(0);
    verifyNoMoreInteractions(fnApiControlClient);
}