Examples with ManagedChannelFactory org.apache.beam.sdk.fn.channel.ManagedChannelFactory used on opensource projects

Example 1 with ManagedChannelFactory

use of org.apache.beam.sdk.fn.channel.ManagedChannelFactory 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 2 with ManagedChannelFactory

use of org.apache.beam.sdk.fn.channel.ManagedChannelFactory in project beam by apache.

the class BeamFnStatusClientTest method testCacheStatsExist.

@Test
public void testCacheStatsExist() {
    ManagedChannelFactory channelFactory = ManagedChannelFactory.createInProcess();
    BeamFnStatusClient client = new BeamFnStatusClient(apiServiceDescriptor, channelFactory::forDescriptor, mock(BundleProcessorCache.class), PipelineOptionsFactory.create(), Caches.fromOptions(PipelineOptionsFactory.fromArgs("--maxCacheMemoryUsageMb=234").create()));
    assertThat(client.getCacheStats(), containsString("used/max 0/234 MB"));
}

Example 3 with ManagedChannelFactory

use of org.apache.beam.sdk.fn.channel.ManagedChannelFactory in project beam by apache.

the class BeamFnStatusClientTest method testWorkerStatusResponse.

@Test
public void testWorkerStatusResponse() throws Exception {
    BlockingQueue<WorkerStatusResponse> values = new LinkedBlockingQueue<>();
    BlockingQueue<StreamObserver<WorkerStatusRequest>> requestObservers = new LinkedBlockingQueue<>();
    StreamObserver<WorkerStatusResponse> inboundServerObserver = TestStreams.withOnNext(values::add).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnWorkerStatusImplBase() {

        @Override
        public StreamObserver<WorkerStatusResponse> workerStatus(StreamObserver<WorkerStatusRequest> responseObserver) {
            Uninterruptibles.putUninterruptibly(requestObservers, responseObserver);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        BundleProcessorCache processorCache = mock(BundleProcessorCache.class);
        when(processorCache.getActiveBundleProcessors()).thenReturn(Collections.emptyMap());
        ManagedChannelFactory channelFactory = ManagedChannelFactory.createInProcess();
        new BeamFnStatusClient(apiServiceDescriptor, channelFactory::forDescriptor, processorCache, PipelineOptionsFactory.create(), Caches.noop());
        StreamObserver<WorkerStatusRequest> requestObserver = requestObservers.take();
        requestObserver.onNext(WorkerStatusRequest.newBuilder().setId("id").build());
        WorkerStatusResponse response = values.take();
        assertThat(response.getStatusInfo(), containsString("No active processing bundles."));
        assertThat(response.getId(), is("id"));
    } finally {
        server.shutdownNow();
    }
}

Example 4 with ManagedChannelFactory

use of org.apache.beam.sdk.fn.channel.ManagedChannelFactory 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.
 *
 * @param id Harness ID
 * @param options The options for this pipeline
 * @param runnerCapabilites
 * @param loggingApiServiceDescriptor
 * @param controlApiServiceDescriptor
 * @param statusApiServiceDescriptor
 * @throws Exception
 */
public static void main(String id, PipelineOptions options, Set<String> runnerCapabilites, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, @Nullable Endpoints.ApiServiceDescriptor statusApiServiceDescriptor) throws Exception {
    ManagedChannelFactory channelFactory;
    if (ExperimentalOptions.hasExperiment(options, "beam_fn_api_epoll")) {
        channelFactory = ManagedChannelFactory.createEpoll();
    } else {
        channelFactory = ManagedChannelFactory.createDefault();
    }
    OutboundObserverFactory outboundObserverFactory = HarnessStreamObserverFactories.fromOptions(options);
    main(id, options, runnerCapabilites, loggingApiServiceDescriptor, controlApiServiceDescriptor, statusApiServiceDescriptor, channelFactory, outboundObserverFactory, Caches.fromOptions(options));
}