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

Example 6 with InstructionResponse

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

Example 7 with InstructionResponse

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

Example 8 with InstructionResponse

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

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

the class FnHarnessTest method testLaunchFnHarnessAndTeardownCleanly.

@Test
// failure will cause test to timeout.
@SuppressWarnings("FutureReturnValueIgnored")
public void testLaunchFnHarnessAndTeardownCleanly() throws Exception {
    Function<String, String> environmentVariableMock = mock(Function.class);
    PipelineOptions options = PipelineOptionsFactory.create();
    when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id");
    when(environmentVariableMock.apply("PIPELINE_OPTIONS")).thenReturn(PipelineOptionsTranslation.toJson(options));
    List<BeamFnApi.LogEntry> logEntries = new ArrayList<>();
    List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class);
    BeamFnLoggingGrpc.BeamFnLoggingImplBase loggingService = new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<LogControl> responseObserver) {
            return TestStreams.withOnNext((BeamFnApi.LogEntry.List entries) -> logEntries.addAll(entries.getLogEntriesList())).withOnCompleted(responseObserver::onCompleted).build();
        }
    };
    BeamFnControlGrpc.BeamFnControlImplBase controlService = new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<InstructionResponse> control(StreamObserver<InstructionRequest> responseObserver) {
            CountDownLatch waitForResponses = new CountDownLatch(1);
            options.as(GcsOptions.class).getExecutorService().submit(() -> {
                responseObserver.onNext(INSTRUCTION_REQUEST);
                Uninterruptibles.awaitUninterruptibly(waitForResponses);
                responseObserver.onCompleted();
            });
            return TestStreams.withOnNext((InstructionResponse t) -> {
                instructionResponses.add(t);
                waitForResponses.countDown();
            }).withOnCompleted(waitForResponses::countDown).build();
        }
    };
    Server loggingServer = ServerBuilder.forPort(0).addService(loggingService).build();
    loggingServer.start();
    try {
        Server controlServer = ServerBuilder.forPort(0).addService(controlService).build();
        controlServer.start();
        try {
            Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + loggingServer.getPort()).build();
            Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + controlServer.getPort()).build();
            when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(loggingDescriptor));
            when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(controlDescriptor));
            FnHarness.main(environmentVariableMock);
        } finally {
            controlServer.shutdownNow();
        }
    } finally {
        loggingServer.shutdownNow();
    }
    // Verify that we first run onStartup functions before even reading the environment, and that
    // we then call beforeProcessing functions before executing instructions.
    InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses);
    inOrder.verify(onStartupMock).run();
    inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any());
    inOrder.verify(beforeProcessingMock).accept(any());
    inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE);
}

Example 10 with InstructionResponse

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

the class BeamFnMapTaskExecutorTest method testExtractCounterUpdatesReturnsValidProgressTrackerCounterUpdatesIfPresent.

@Test(timeout = ReadOperation.DEFAULT_PROGRESS_UPDATE_PERIOD_MS * 60)
public void testExtractCounterUpdatesReturnsValidProgressTrackerCounterUpdatesIfPresent() throws Exception {
    final CountDownLatch progressSentLatch = new CountDownLatch(1);
    final CountDownLatch processBundleLatch = new CountDownLatch(1);
    final int expectedCounterValue = 5;
    final MonitoringInfo expectedMonitoringInfo = MonitoringInfo.newBuilder().setUrn(MonitoringInfoConstants.Urns.USER_SUM_INT64).putLabels(MonitoringInfoConstants.Labels.NAME, "ExpectedCounter").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyString").setType(MonitoringInfoConstants.TypeUrns.SUM_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "ExpectedPTransform").setPayload(encodeInt64Counter(expectedCounterValue)).build();
    InstructionRequestHandler instructionRequestHandler = new InstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            switch(request.getRequestCase()) {
                case REGISTER:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                case PROCESS_BUNDLE:
                    return MoreFutures.supplyAsync(() -> {
                        processBundleLatch.await();
                        return responseFor(request).setProcessBundle(BeamFnApi.ProcessBundleResponse.newBuilder().addMonitoringInfos(expectedMonitoringInfo)).build();
                    });
                case PROCESS_BUNDLE_PROGRESS:
                    progressSentLatch.countDown();
                    return CompletableFuture.completedFuture(responseFor(request).setProcessBundleProgress(BeamFnApi.ProcessBundleProgressResponse.newBuilder().addMonitoringInfos(expectedMonitoringInfo)).build());
                default:
                    throw new RuntimeException("Reached unexpected code path");
            }
        }

        @Override
        public void registerProcessBundleDescriptor(ProcessBundleDescriptor descriptor) {
        }

        @Override
        public void close() {
        }
    };
    Map<String, DataflowStepContext> stepContextMap = new HashMap<>();
    stepContextMap.put("ExpectedPTransform", generateDataflowStepContext("Expected"));
    RegisterAndProcessBundleOperation processOperation = new RegisterAndProcessBundleOperation(IdGenerators.decrementingLongs(), instructionRequestHandler, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), stepContextMap, ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    BeamFnMapTaskExecutor mapTaskExecutor = BeamFnMapTaskExecutor.forOperations(ImmutableList.of(readOperation, grpcPortWriteOperation, processOperation), executionStateTracker);
    // Launch the BeamFnMapTaskExecutor and wait until we are sure there has been one
    // tentative update
    CompletionStage<Void> doneFuture = MoreFutures.runAsync(mapTaskExecutor::execute);
    progressSentLatch.await();
    Iterable<CounterUpdate> metricsCounterUpdates = Collections.emptyList();
    while (Iterables.size(metricsCounterUpdates) == 0) {
        Thread.sleep(ReadOperation.DEFAULT_PROGRESS_UPDATE_PERIOD_MS);
        metricsCounterUpdates = mapTaskExecutor.extractMetricUpdates();
    }
    // Get the final metrics
    processBundleLatch.countDown();
    MoreFutures.get(doneFuture);
    metricsCounterUpdates = mapTaskExecutor.extractMetricUpdates();
    assertThat(Iterables.size(metricsCounterUpdates), equalTo(1));
    CounterUpdate resultCounter = metricsCounterUpdates.iterator().next();
    assertTrue(new CounterHamcrestMatchers.CounterUpdateIntegerValueMatcher(expectedCounterValue).matches(resultCounter));
    assertEquals("ExpectedCounter", resultCounter.getStructuredNameAndMetadata().getName().getName());
}