Examples with SimpleExecutionState org.apache.beam.runners.core.metrics.SimpleExecutionState used on opensource projects

Example 1 with SimpleExecutionState

use of org.apache.beam.runners.core.metrics.SimpleExecutionState in project beam by apache.

the class PCollectionConsumerRegistry method register.

/**
 * Register the specified consumer to handle the elements in the pCollection associated with
 * pCollectionId. All consumers must be registered before extracting the combined consumer by
 * calling getMultiplexingConsumer(), or an exception will be thrown.
 *
 * <p>This will cause both Element Count and Process Bundle Execution time metrics to be
 * collected.
 *
 * @param pCollectionId
 * @param pTransformId
 * @param consumer
 * @param <T> the element type of the PCollection
 * @throws RuntimeException if {@code register()} is called after {@code
 *     getMultiplexingConsumer()} is called.
 */
public <T> void register(String pCollectionId, String pTransformId, FnDataReceiver<WindowedValue<T>> consumer, Coder<T> valueCoder) {
    // if there are multiple consumers.
    if (pCollectionIdsToWrappedConsumer.containsKey(pCollectionId)) {
        throw new RuntimeException("New consumers for a pCollectionId cannot be register()-d after " + "calling getMultiplexingConsumer.");
    }
    HashMap<String, String> labelsMetadata = new HashMap<>();
    labelsMetadata.put(MonitoringInfoConstants.Labels.PTRANSFORM, pTransformId);
    SimpleExecutionState state = new SimpleExecutionState(ExecutionStateTracker.PROCESS_STATE_NAME, MonitoringInfoConstants.Urns.PROCESS_BUNDLE_MSECS, labelsMetadata);
    executionStates.register(state);
    pCollectionIdsToConsumers.put(pCollectionId, ConsumerAndMetadata.forConsumer(consumer, pTransformId, state, valueCoder, metricsContainerRegistry.getContainer(pTransformId)));
}

Example 2 with SimpleExecutionState

use of org.apache.beam.runners.core.metrics.SimpleExecutionState 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);
}