Examples with ThrowingRunnable org.apache.beam.sdk.function.ThrowingRunnable used on opensource projects

Example 1 with ThrowingRunnable

use of org.apache.beam.sdk.function.ThrowingRunnable in project beam by apache.

the class PTransformFunctionRegistryTest method functionsAreInvokedIndirectlyAfterRegisteringAndInvoking.

@Test
public void functionsAreInvokedIndirectlyAfterRegisteringAndInvoking() throws Exception {
    PTransformFunctionRegistry testObject = new PTransformFunctionRegistry(mock(MetricsContainerStepMap.class), mock(ExecutionStateTracker.class), "start");
    ThrowingRunnable runnableA = mock(ThrowingRunnable.class);
    ThrowingRunnable runnableB = mock(ThrowingRunnable.class);
    testObject.register("pTransformA", runnableA);
    testObject.register("pTransformB", runnableB);
    for (ThrowingRunnable func : testObject.getFunctions()) {
        func.run();
    }
    verify(runnableA, times(1)).run();
    verify(runnableB, times(1)).run();
}

Example 2 with ThrowingRunnable

use of org.apache.beam.sdk.function.ThrowingRunnable 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 3 with ThrowingRunnable

use of org.apache.beam.sdk.function.ThrowingRunnable in project beam by apache.

the class FnApiStateAccessor method finalizeState.

public void finalizeState() {
    // Persist all dirty state cells
    try {
        for (ThrowingRunnable runnable : stateFinalizers) {
            runnable.run();
        }
    } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw new IllegalStateException(e);
    } catch (Exception e) {
        throw new IllegalStateException(e);
    }
    stateFinalizers.clear();
    stateKeyObjectCache.clear();
}

Example 4 with ThrowingRunnable

use of org.apache.beam.sdk.function.ThrowingRunnable in project beam by apache.

the class ProcessBundleHandler method processBundle.

/**
 * Processes a bundle, running the start(), process(), and finish() functions. This function is
 * required to be reentrant.
 */
public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request) throws Exception {
    BeamFnApi.ProcessBundleResponse.Builder response = BeamFnApi.ProcessBundleResponse.newBuilder();
    BundleProcessor bundleProcessor = bundleProcessorCache.get(request, () -> {
        try {
            return createBundleProcessor(request.getProcessBundle().getProcessBundleDescriptorId(), request.getProcessBundle());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    });
    try {
        PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry();
        PTransformFunctionRegistry finishFunctionRegistry = bundleProcessor.getFinishFunctionRegistry();
        ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker();
        try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) {
            try (Closeable closeTracker = stateTracker.activate()) {
                // Already in reverse topological order so we don't need to do anything.
                for (ThrowingRunnable startFunction : startFunctionRegistry.getFunctions()) {
                    LOG.debug("Starting function {}", startFunction);
                    startFunction.run();
                }
                if (request.getProcessBundle().hasElements()) {
                    boolean inputFinished = bundleProcessor.getInboundObserver().multiplexElements(request.getProcessBundle().getElements());
                    if (!inputFinished) {
                        throw new RuntimeException("Elements embedded in ProcessBundleRequest do not contain stream terminators for " + "all data and timer inputs. Unterminated endpoints: " + bundleProcessor.getInboundObserver().getUnfinishedEndpoints());
                    }
                } else if (!bundleProcessor.getInboundEndpointApiServiceDescriptors().isEmpty()) {
                    BeamFnDataInboundObserver2 observer = bundleProcessor.getInboundObserver();
                    beamFnDataClient.registerReceiver(request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors(), observer);
                    observer.awaitCompletion();
                    beamFnDataClient.unregisterReceiver(request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors());
                }
                // Need to reverse this since we want to call finish in topological order.
                for (ThrowingRunnable finishFunction : Lists.reverse(finishFunctionRegistry.getFunctions())) {
                    LOG.debug("Finishing function {}", finishFunction);
                    finishFunction.run();
                }
            }
            // If bundleProcessor has not flushed any elements, embed them in response.
            embedOutboundElementsIfApplicable(response, bundleProcessor);
            // Add all checkpointed residuals to the response.
            response.addAllResidualRoots(bundleProcessor.getSplitListener().getResidualRoots());
            // Add all metrics to the response.
            Map<String, ByteString> monitoringData = monitoringData(bundleProcessor);
            if (runnerAcceptsShortIds) {
                response.putAllMonitoringData(monitoringData);
            } else {
                for (Map.Entry<String, ByteString> metric : monitoringData.entrySet()) {
                    response.addMonitoringInfos(shortIds.get(metric.getKey()).toBuilder().setPayload(metric.getValue()));
                }
            }
            if (!bundleProcessor.getBundleFinalizationCallbackRegistrations().isEmpty()) {
                finalizeBundleHandler.registerCallbacks(bundleProcessor.getInstructionId(), ImmutableList.copyOf(bundleProcessor.getBundleFinalizationCallbackRegistrations()));
                response.setRequiresFinalization(true);
            }
        }
        // Mark the bundle processor as re-usable.
        bundleProcessorCache.release(request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor);
        return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response);
    } catch (Exception e) {
        // Make sure we clean-up from the active set of bundle processors.
        bundleProcessorCache.discard(bundleProcessor);
        throw e;
    }
}

Example 5 with ThrowingRunnable

use of org.apache.beam.sdk.function.ThrowingRunnable in project beam by apache.

the class PTransformFunctionRegistry method register.

/**
 * Register the runnable to process the specific pTransformId and track its execution time.
 *
 * @param pTransformId
 * @param runnable
 */
public void register(String pTransformId, ThrowingRunnable runnable) {
    HashMap<String, String> labelsMetadata = new HashMap<String, String>();
    labelsMetadata.put(MonitoringInfoConstants.Labels.PTRANSFORM, pTransformId);
    String executionTimeUrn = "";
    if (executionStateName.equals(ExecutionStateTracker.START_STATE_NAME)) {
        executionTimeUrn = MonitoringInfoConstants.Urns.START_BUNDLE_MSECS;
    } else if (executionStateName.equals(ExecutionStateTracker.FINISH_STATE_NAME)) {
        executionTimeUrn = MonitoringInfoConstants.Urns.FINISH_BUNDLE_MSECS;
    }
    SimpleExecutionState state = new SimpleExecutionState(this.executionStateName, executionTimeUrn, labelsMetadata);
    executionStates.register(state);
    ThrowingRunnable wrapped = () -> {
        MetricsContainerImpl container = metricsContainerRegistry.getContainer(pTransformId);
        try (Closeable metricCloseable = MetricsEnvironment.scopedMetricsContainer(container)) {
            try (Closeable trackerCloseable = this.stateTracker.enterState(state)) {
                runnable.run();
            }
        }
    };
    runnables.add(wrapped);
}