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

Example 1 with BundleProcessor

use of org.apache.beam.fn.harness.control.ProcessBundleHandler.BundleProcessor 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 2 with BundleProcessor

use of org.apache.beam.fn.harness.control.ProcessBundleHandler.BundleProcessor 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 BundleProcessor

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

the class ProcessBundleHandlerTest method testBundleProcessorReset.

@Test
public void testBundleProcessorReset() throws Exception {
    PTransformFunctionRegistry startFunctionRegistry = mock(PTransformFunctionRegistry.class);
    PTransformFunctionRegistry finishFunctionRegistry = mock(PTransformFunctionRegistry.class);
    BundleSplitListener.InMemory splitListener = mock(BundleSplitListener.InMemory.class);
    Collection<CallbackRegistration> bundleFinalizationCallbacks = mock(Collection.class);
    PCollectionConsumerRegistry pCollectionConsumerRegistry = mock(PCollectionConsumerRegistry.class);
    MetricsContainerStepMap metricsContainerRegistry = mock(MetricsContainerStepMap.class);
    ExecutionStateTracker stateTracker = mock(ExecutionStateTracker.class);
    ProcessBundleHandler.HandleStateCallsForBundle beamFnStateClient = mock(ProcessBundleHandler.HandleStateCallsForBundle.class);
    ThrowingRunnable resetFunction = mock(ThrowingRunnable.class);
    Cache<Object, Object> processWideCache = Caches.eternal();
    BundleProcessor bundleProcessor = BundleProcessor.create(processWideCache, ProcessBundleDescriptor.getDefaultInstance(), startFunctionRegistry, finishFunctionRegistry, Collections.singletonList(resetFunction), new ArrayList<>(), new ArrayList<>(), splitListener, pCollectionConsumerRegistry, metricsContainerRegistry, stateTracker, beamFnStateClient, bundleFinalizationCallbacks, new HashSet<>());
    bundleProcessor.finish();
    CacheToken cacheToken = CacheToken.newBuilder().setSideInput(CacheToken.SideInput.newBuilder().setTransformId("transformId")).build();
    bundleProcessor.setupForProcessBundleRequest(processBundleRequestFor("instructionId", "descriptorId", cacheToken));
    assertEquals("instructionId", bundleProcessor.getInstructionId());
    assertThat(bundleProcessor.getCacheTokens(), containsInAnyOrder(cacheToken));
    Cache<Object, Object> bundleCache = bundleProcessor.getBundleCache();
    bundleCache.put("A", "B");
    assertEquals("B", bundleCache.peek("A"));
    bundleProcessor.reset();
    assertNull(bundleProcessor.getInstructionId());
    assertNull(bundleProcessor.getCacheTokens());
    assertNull(bundleCache.peek("A"));
    verify(startFunctionRegistry, times(1)).reset();
    verify(finishFunctionRegistry, times(1)).reset();
    verify(splitListener, times(1)).clear();
    verify(pCollectionConsumerRegistry, times(1)).reset();
    verify(metricsContainerRegistry, times(1)).reset();
    verify(stateTracker, times(1)).reset();
    verify(bundleFinalizationCallbacks, times(1)).clear();
    verify(resetFunction, times(1)).run();
    // Ensure that the next setup produces the expected state.
    bundleProcessor.setupForProcessBundleRequest(processBundleRequestFor("instructionId2", "descriptorId2"));
    assertNotSame(bundleCache, bundleProcessor.getBundleCache());
    assertEquals("instructionId2", bundleProcessor.getInstructionId());
    assertThat(bundleProcessor.getCacheTokens(), is(emptyIterable()));
}

Example 4 with BundleProcessor

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

the class ProcessBundleHandlerTest method testBundleProcessorIsFoundWhenActive.

@Test
public void testBundleProcessorIsFoundWhenActive() {
    BundleProcessor bundleProcessor = mock(BundleProcessor.class);
    when(bundleProcessor.getInstructionId()).thenReturn("known");
    BundleProcessorCache cache = new BundleProcessorCache();
    // Check that an unknown bundle processor is not found
    assertNull(cache.find("unknown"));
    // Once it is active, ensure the bundle processor is found
    cache.get(processBundleRequestFor("known", "descriptorId"), () -> bundleProcessor);
    assertSame(bundleProcessor, cache.find("known"));
    // After it is released, ensure the bundle processor is no longer found
    cache.release("descriptorId", bundleProcessor);
    assertNull(cache.find("known"));
    // Once it is active, ensure the bundle processor is found
    cache.get(processBundleRequestFor("known", "descriptorId"), () -> bundleProcessor);
    assertSame(bundleProcessor, cache.find("known"));
    // After it is discarded, ensure the bundle processor is no longer found
    cache.discard(bundleProcessor);
    verify(bundleProcessor).discard();
    assertNull(cache.find("known"));
}

Example 5 with BundleProcessor

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

the class BeamFnStatusClient method getActiveProcessBundleState.

@VisibleForTesting
String getActiveProcessBundleState() {
    StringJoiner activeBundlesState = new StringJoiner("\n");
    activeBundlesState.add("========== ACTIVE PROCESSING BUNDLES ==========");
    if (processBundleCache.getActiveBundleProcessors().isEmpty()) {
        activeBundlesState.add("No active processing bundles.");
    } else {
        List<BundleState> bundleStates = new ArrayList<>();
        processBundleCache.getActiveBundleProcessors().keySet().stream().forEach(instruction -> {
            BundleProcessor bundleProcessor = processBundleCache.find(instruction);
            if (bundleProcessor != null) {
                ExecutionStateTracker executionStateTracker = bundleProcessor.getStateTracker();
                Thread trackedTread = executionStateTracker.getTrackedThread();
                if (trackedTread != null) {
                    bundleStates.add(new BundleState(instruction, trackedTread.getName(), executionStateTracker.getMillisSinceLastTransition()));
                }
            }
        });
        bundleStates.stream().sorted(Comparator.comparing(BundleState::getTimeSinceTransition).reversed()).limit(// only keep top 10
        10).forEachOrdered(bundleState -> {
            activeBundlesState.add(String.format("---- Instruction %s ----", bundleState.getInstruction()));
            activeBundlesState.add(String.format("Tracked thread: %s", bundleState.getTrackedThreadName()));
            activeBundlesState.add(String.format("Time since transition: %.2f seconds%n", bundleState.getTimeSinceTransition() / 1000.0));
        });
    }
    return activeBundlesState.toString();
}