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

Example 16 with BundleProcessorCache

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

the class ProcessBundleHandlerTest method testBundleFinalizationIsPropagated.

@Test
public void testBundleFinalizationIsPropagated() 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);
    FinalizeBundleHandler mockFinalizeBundleHandler = mock(FinalizeBundleHandler.class);
    BundleFinalizer.Callback mockCallback = mock(BundleFinalizer.Callback.class);
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
    mockFinalizeBundleHandler, new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (context) -> {
        BundleFinalizer bundleFinalizer = context.getBundleFinalizer();
        context.addStartBundleFunction(() -> bundleFinalizer.afterBundleCommit(Instant.ofEpochMilli(42L), mockCallback));
        return null;
    }), Caches.noop(), new BundleProcessorCache());
    BeamFnApi.InstructionResponse.Builder response = handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("2L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    assertTrue(response.getProcessBundle().getRequiresFinalization());
    verify(mockFinalizeBundleHandler).registerCallbacks(eq("2L"), argThat((Collection<CallbackRegistration> arg) -> {
        CallbackRegistration registration = Iterables.getOnlyElement(arg);
        assertEquals(Instant.ofEpochMilli(42L), registration.getExpiryTime());
        assertSame(mockCallback, registration.getCallback());
        return true;
    }));
}

Example 17 with BundleProcessorCache

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

the class ProcessBundleHandlerTest method testTrySplitBeforeBundleDoesNotFail.

@Test
public void testTrySplitBeforeBundleDoesNotFail() {
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), null, beamFnDataClient, null, /* beamFnStateClient */
    null, /* finalizeBundleHandler */
    new ShortIdMap(), ImmutableMap.of(), Caches.noop(), new BundleProcessorCache());
    BeamFnApi.InstructionResponse response = handler.trySplit(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundleSplit(BeamFnApi.ProcessBundleSplitRequest.newBuilder().setInstructionId("unknown-id")).build()).build();
    assertNotNull(response.getProcessBundleSplit());
    assertEquals(0, response.getProcessBundleSplit().getChannelSplitsCount());
}

Example 18 with BundleProcessorCache

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