Examples with BeamFnStateGrpcClientCache org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache used on opensource projects

Example 1 with BeamFnStateGrpcClientCache

use of org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache in project beam by apache.

the class ProcessBundleHandlerTest method testCreatingPTransformExceptionsArePropagated.

@Test
public void testCreatingPTransformExceptionsArePropagated() throws Exception {
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
    null, /* finalizeBundleHandler */
    new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (context) -> {
        throw new IllegalStateException("TestException");
    }), Caches.noop(), new BundleProcessorCache());
    assertThrows("TestException", IllegalStateException.class, () -> handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build()));
}

Example 2 with BeamFnStateGrpcClientCache

use of org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache in project beam by apache.

the class ProcessBundleHandlerTest method testBundleProcessorIsResetWhenAddedBackToCache.

@Test
public void testBundleProcessorIsResetWhenAddedBackToCache() throws Exception {
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
    null, /* finalizeBundleHandler */
    new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (context) -> null), Caches.noop(), new TestBundleProcessorCache());
    assertThat(TestBundleProcessor.resetCnt, equalTo(0));
    handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("998L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    // Check that BundleProcessor is reset when added back to the cache
    assertThat(TestBundleProcessor.resetCnt, equalTo(1));
    // BundleProcessor is added back to the BundleProcessorCache
    assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().size(), equalTo(1));
    assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(1));
    // Add a reset handler that throws to test discarding the bundle processor on reset failure.
    Iterables.getOnlyElement(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L")).getResetFunctions().add(() -> {
        throw new IllegalStateException("ResetFailed");
    });
    handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    // BundleProcessor is discarded instead of being added back to the BundleProcessorCache
    assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(0));
}

Example 3 with BeamFnStateGrpcClientCache

use of org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache in project beam by apache.

the class ProcessBundleHandlerTest method testPendingStateCallsBlockTillCompletion.

@Test
public void testPendingStateCallsBlockTillCompletion() throws Exception {
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).setStateApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    CompletableFuture<StateResponse>[] successfulResponse = new CompletableFuture[1];
    CompletableFuture<StateResponse>[] unsuccessfulResponse = new CompletableFuture[1];
    BeamFnStateGrpcClientCache mockBeamFnStateGrpcClient = Mockito.mock(BeamFnStateGrpcClientCache.class);
    BeamFnStateClient mockBeamFnStateClient = Mockito.mock(BeamFnStateClient.class);
    when(mockBeamFnStateGrpcClient.forApiServiceDescriptor(any())).thenReturn(mockBeamFnStateClient);
    doAnswer(invocation -> {
        StateRequest.Builder stateRequestBuilder = (StateRequest.Builder) invocation.getArguments()[0];
        CompletableFuture<StateResponse> completableFuture = new CompletableFuture<>();
        new Thread(() -> {
            // Simulate sleeping which introduces a race which most of the time requires
            // the ProcessBundleHandler to block.
            Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
            switch(stateRequestBuilder.getInstructionId()) {
                case "SUCCESS":
                    completableFuture.complete(StateResponse.getDefaultInstance());
                    break;
                case "FAIL":
                    completableFuture.completeExceptionally(new RuntimeException("TEST ERROR"));
            }
        }).start();
        return completableFuture;
    }).when(mockBeamFnStateClient).handle(any());
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, mockBeamFnStateGrpcClient, null, /* finalizeBundleHandler */
    new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {

        @Override
        public Object createRunnerForPTransform(Context context) throws IOException {
            BeamFnStateClient beamFnStateClient = context.getBeamFnStateClient();
            context.addStartBundleFunction(() -> doStateCalls(beamFnStateClient));
            return null;
        }

        private void doStateCalls(BeamFnStateClient beamFnStateClient) {
            successfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("SUCCESS"));
            unsuccessfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("FAIL"));
        }
    }), Caches.noop(), new BundleProcessorCache());
    handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    assertTrue(successfulResponse[0].isDone());
    assertTrue(unsuccessfulResponse[0].isDone());
}

Example 4 with BeamFnStateGrpcClientCache

use of org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache in project beam by apache.

the class FnHarness method main.

/**
 * Run a FnHarness with the given id and options that attaches to the specified logging and
 * control API service descriptors using the given channel factory and outbound observer factory.
 *
 * @param id Harness ID
 * @param options The options for this pipeline
 * @param runnerCapabilites
 * @param loggingApiServiceDescriptor
 * @param controlApiServiceDescriptor
 * @param statusApiServiceDescriptor
 * @param channelFactory
 * @param outboundObserverFactory
 * @param processWideCache
 * @throws Exception
 */
public static void main(String id, PipelineOptions options, Set<String> runnerCapabilites, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, Endpoints.ApiServiceDescriptor statusApiServiceDescriptor, ManagedChannelFactory channelFactory, OutboundObserverFactory outboundObserverFactory, Cache<Object, Object> processWideCache) throws Exception {
    channelFactory = channelFactory.withInterceptors(ImmutableList.of(AddHarnessIdInterceptor.create(id)));
    IdGenerator idGenerator = IdGenerators.decrementingLongs();
    ShortIdMap metricsShortIds = new ShortIdMap();
    ExecutorService executorService = options.as(GcsOptions.class).getExecutorService();
    // intercepts logging and sends it to the logging service.
    try (BeamFnLoggingClient logging = new BeamFnLoggingClient(options, loggingApiServiceDescriptor, channelFactory::forDescriptor)) {
        LOG.info("Fn Harness started");
        // Register standard file systems.
        FileSystems.setDefaultPipelineOptions(options);
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        ManagedChannel channel = channelFactory.forDescriptor(controlApiServiceDescriptor);
        BeamFnControlGrpc.BeamFnControlStub controlStub = BeamFnControlGrpc.newStub(channel);
        BeamFnControlGrpc.BeamFnControlBlockingStub blockingControlStub = BeamFnControlGrpc.newBlockingStub(channel);
        BeamFnDataGrpcClient beamFnDataMultiplexer = new BeamFnDataGrpcClient(options, channelFactory::forDescriptor, outboundObserverFactory);
        BeamFnStateGrpcClientCache beamFnStateGrpcClientCache = new BeamFnStateGrpcClientCache(idGenerator, channelFactory, outboundObserverFactory);
        FinalizeBundleHandler finalizeBundleHandler = new FinalizeBundleHandler(options.as(GcsOptions.class).getExecutorService());
        Function<String, BeamFnApi.ProcessBundleDescriptor> getProcessBundleDescriptor = new Function<String, ProcessBundleDescriptor>() {

            private static final String PROCESS_BUNDLE_DESCRIPTORS = "ProcessBundleDescriptors";

            private final Cache<String, BeamFnApi.ProcessBundleDescriptor> cache = Caches.subCache(processWideCache, PROCESS_BUNDLE_DESCRIPTORS);

            @Override
            public BeamFnApi.ProcessBundleDescriptor apply(String id) {
                return cache.computeIfAbsent(id, this::loadDescriptor);
            }

            private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
                return blockingControlStub.getProcessBundleDescriptor(BeamFnApi.GetProcessBundleDescriptorRequest.newBuilder().setProcessBundleDescriptorId(id).build());
            }
        };
        MetricsEnvironment.setProcessWideContainer(MetricsContainerImpl.createProcessWideContainer());
        ProcessBundleHandler processBundleHandler = new ProcessBundleHandler(options, runnerCapabilites, getProcessBundleDescriptor, beamFnDataMultiplexer, beamFnStateGrpcClientCache, finalizeBundleHandler, metricsShortIds, processWideCache);
        BeamFnStatusClient beamFnStatusClient = null;
        if (statusApiServiceDescriptor != null) {
            beamFnStatusClient = new BeamFnStatusClient(statusApiServiceDescriptor, channelFactory::forDescriptor, processBundleHandler.getBundleProcessorCache(), options, processWideCache);
        }
        // TODO(BEAM-9729): Remove once runners no longer send this instruction.
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, request -> BeamFnApi.InstructionResponse.newBuilder().setRegister(BeamFnApi.RegisterResponse.getDefaultInstance()));
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.FINALIZE_BUNDLE, finalizeBundleHandler::finalizeBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, processBundleHandler::processBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_PROGRESS, processBundleHandler::progress);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_SPLIT, processBundleHandler::trySplit);
        handlers.put(InstructionRequest.RequestCase.MONITORING_INFOS, request -> BeamFnApi.InstructionResponse.newBuilder().setMonitoringInfos(BeamFnApi.MonitoringInfosMetadataResponse.newBuilder().putAllMonitoringInfo(StreamSupport.stream(request.getMonitoringInfos().getMonitoringInfoIdList().spliterator(), false).collect(Collectors.toMap(Function.identity(), metricsShortIds::get)))));
        HarnessMonitoringInfosInstructionHandler processWideHandler = new HarnessMonitoringInfosInstructionHandler(metricsShortIds);
        handlers.put(InstructionRequest.RequestCase.HARNESS_MONITORING_INFOS, processWideHandler::harnessMonitoringInfos);
        JvmInitializers.runBeforeProcessing(options);
        String samplingPeriodMills = ExperimentalOptions.getExperimentValue(options, ExperimentalOptions.STATE_SAMPLING_PERIOD_MILLIS);
        if (samplingPeriodMills != null) {
            ExecutionStateSampler.setSamplingPeriod(Integer.parseInt(samplingPeriodMills));
        }
        ExecutionStateSampler.instance().start();
        LOG.info("Entering instruction processing loop");
        // The control client immediately dispatches requests to an executor so we execute on the
        // direct executor. If we created separate channels for different stubs we could use
        // directExecutor() when building the channel.
        BeamFnControlClient control = new BeamFnControlClient(controlStub.withExecutor(MoreExecutors.directExecutor()), outboundObserverFactory, executorService, handlers);
        control.waitForTermination();
        if (beamFnStatusClient != null) {
            beamFnStatusClient.close();
        }
        processBundleHandler.shutdown();
    } finally {
        System.out.println("Shutting SDK harness down.");
        ExecutionStateSampler.instance().stop();
        executorService.shutdown();
    }
}