Examples with ProcessBundleHandler org.apache.beam.fn.harness.control.ProcessBundleHandler used on opensource projects

Example 1 with ProcessBundleHandler

use of org.apache.beam.fn.harness.control.ProcessBundleHandler in project beam by apache.

the class FnHarness method main.

public static void main(PipelineOptions options, BeamFnApi.ApiServiceDescriptor loggingApiServiceDescriptor, BeamFnApi.ApiServiceDescriptor controlApiServiceDescriptor) throws Exception {
    ManagedChannelFactory channelFactory = ManagedChannelFactory.from(options);
    StreamObserverFactory streamObserverFactory = StreamObserverFactory.fromOptions(options);
    PrintStream originalErrStream = System.err;
    try (BeamFnLoggingClient logging = new BeamFnLoggingClient(options, loggingApiServiceDescriptor, channelFactory::forDescriptor, streamObserverFactory::from)) {
        LOG.info("Fn Harness started");
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        RegisterHandler fnApiRegistry = new RegisterHandler();
        BeamFnDataGrpcClient beamFnDataMultiplexer = new BeamFnDataGrpcClient(options, channelFactory::forDescriptor, streamObserverFactory::from);
        ProcessBundleHandler processBundleHandler = new ProcessBundleHandler(options, fnApiRegistry::getById, beamFnDataMultiplexer);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, fnApiRegistry::register);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, processBundleHandler::processBundle);
        BeamFnControlClient control = new BeamFnControlClient(controlApiServiceDescriptor, channelFactory::forDescriptor, streamObserverFactory::from, handlers);
        LOG.info("Entering instruction processing loop");
        control.processInstructionRequests(options.as(GcsOptions.class).getExecutorService());
    } catch (Throwable t) {
        t.printStackTrace(originalErrStream);
    } finally {
        originalErrStream.println("Shutting SDK harness down.");
    }
}

Example 2 with ProcessBundleHandler

use of org.apache.beam.fn.harness.control.ProcessBundleHandler in project beam by apache.

the class BeamFnStatusClientTest method testActiveBundleState.

@Test
public void testActiveBundleState() {
    ProcessBundleHandler handler = mock(ProcessBundleHandler.class);
    BundleProcessorCache processorCache = mock(BundleProcessorCache.class);
    Map<String, BundleProcessor> bundleProcessorMap = new HashMap<>();
    for (int i = 0; i < 11; i++) {
        BundleProcessor processor = mock(BundleProcessor.class);
        ExecutionStateTracker executionStateTracker = mock(ExecutionStateTracker.class);
        when(processor.getStateTracker()).thenReturn(executionStateTracker);
        when(executionStateTracker.getMillisSinceLastTransition()).thenReturn(Integer.toUnsignedLong((10 - i) * 1000));
        when(executionStateTracker.getTrackedThread()).thenReturn(Thread.currentThread());
        String instruction = Integer.toString(i);
        when(processorCache.find(instruction)).thenReturn(processor);
        bundleProcessorMap.put(instruction, processor);
    }
    when(handler.getBundleProcessorCache()).thenReturn(processorCache);
    when(processorCache.getActiveBundleProcessors()).thenReturn(bundleProcessorMap);
    ManagedChannelFactory channelFactory = ManagedChannelFactory.createInProcess();
    BeamFnStatusClient client = new BeamFnStatusClient(apiServiceDescriptor, channelFactory::forDescriptor, handler.getBundleProcessorCache(), PipelineOptionsFactory.create(), Caches.noop());
    StringJoiner joiner = new StringJoiner("\n");
    joiner.add(client.getActiveProcessBundleState());
    String actualState = joiner.toString();
    List<String> expectedInstructions = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        expectedInstructions.add(String.format("Instruction %d", i));
    }
    assertThat(actualState, stringContainsInOrder(expectedInstructions));
    assertThat(actualState, not(containsString("Instruction 10")));
}

Example 3 with ProcessBundleHandler

use of org.apache.beam.fn.harness.control.ProcessBundleHandler 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();
    }
}