Examples with ProcessBundleRequest org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest used on opensource projects

Example 1 with ProcessBundleRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest in project beam by apache.

the class SdkHarnessClientTest method verifyCacheTokensAreUsedInNewBundleRequest.

@Test
public void verifyCacheTokensAreUsedInNewBundleRequest() throws InterruptedException {
    when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(CompletableFuture.<InstructionResponse>completedFuture(InstructionResponse.newBuilder().build()));
    ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
    List<RemoteInputDestination> remoteInputs = Collections.singletonList(RemoteInputDestination.of(FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM));
    BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
    when(dataService.send(any(), any())).thenReturn(mock(CloseableFnDataReceiver.class));
    StateRequestHandler stateRequestHandler = Mockito.mock(StateRequestHandler.class);
    List<BeamFnApi.ProcessBundleRequest.CacheToken> cacheTokens = Collections.singletonList(BeamFnApi.ProcessBundleRequest.CacheToken.newBuilder().getDefaultInstanceForType());
    when(stateRequestHandler.getCacheTokens()).thenReturn(cacheTokens);
    processor1.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mock(RemoteOutputReceiver.class)), stateRequestHandler, BundleProgressHandler.ignored());
    // Retrieve the requests made to the FnApiControlClient
    ArgumentCaptor<BeamFnApi.InstructionRequest> reqCaptor = ArgumentCaptor.forClass(BeamFnApi.InstructionRequest.class);
    Mockito.verify(fnApiControlClient, Mockito.times(1)).handle(reqCaptor.capture());
    List<BeamFnApi.InstructionRequest> requests = reqCaptor.getAllValues();
    // Verify that the cache tokens are included in the ProcessBundleRequest
    assertThat(requests.get(0).getRequestCase(), is(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE));
    assertThat(requests.get(0).getProcessBundle().getCacheTokensList(), is(cacheTokens));
}

Example 2 with ProcessBundleRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest 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 3 with ProcessBundleRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest in project beam by apache.

the class RegisterAndProcessBundleOperation method getProcessBundleProgress.

/**
 * Returns the compound metrics recorded, by issuing a request to the SDK harness.
 *
 * <p>This includes key progress indicators as well as user-defined metrics.
 *
 * <p>Use {@link #getInputElementsConsumed} on the future value to extract the elements consumed
 * from the upstream read operation.
 *
 * <p>May be called at any time, including before start() and after finish().
 */
public CompletionStage<BeamFnApi.ProcessBundleProgressResponse> getProcessBundleProgress() {
    // processBundleId may be reset if this bundle finishes asynchronously.
    String processBundleId = this.processBundleId;
    if (processBundleId == null) {
        return CompletableFuture.completedFuture(BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance());
    }
    InstructionRequest processBundleRequest = InstructionRequest.newBuilder().setInstructionId(idGenerator.getId()).setProcessBundleProgress(ProcessBundleProgressRequest.newBuilder().setInstructionId(processBundleId)).build();
    return instructionRequestHandler.handle(processBundleRequest).thenApply(InstructionResponse::getProcessBundleProgress);
}

Example 4 with ProcessBundleRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest in project beam by apache.

the class RegisterAndProcessBundleOperation method start.

@Override
public void start() throws Exception {
    try (Closeable scope = context.enterStart()) {
        super.start();
        // Only register once by using the presence of the future as a signal.
        if (registerFuture == null) {
            InstructionRequest request = InstructionRequest.newBuilder().setInstructionId(idGenerator.getId()).setRegister(registerRequest).build();
            registerFuture = instructionRequestHandler.handle(request);
        }
        checkState(registerRequest.getProcessBundleDescriptorCount() == 1, "Only one bundle registration at a time currently supported.");
        InstructionRequest processBundleRequest = InstructionRequest.newBuilder().setInstructionId(getProcessBundleInstructionId()).setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId(registerRequest.getProcessBundleDescriptor(0).getId())).build();
        deregisterStateHandler = beamFnStateDelegator.registerForProcessBundleInstructionId(getProcessBundleInstructionId(), this::delegateByStateKeyType);
        processBundleResponse = getRegisterResponse(registerFuture).thenCompose(registerResponse -> instructionRequestHandler.handle(processBundleRequest));
    }
}

Example 5 with ProcessBundleRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleRequest 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);
    }
}