Examples with DataflowExecutionState org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState used on opensource projects

Example 6 with DataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState in project beam by apache.

the class IsmReaderTest method setUp.

@Before
public void setUp() {
    cache = CacheBuilder.newBuilder().weigher(Weighers.fixedWeightKeys(1)).maximumWeight(10_000).build();
    executionContext = BatchModeExecutionContext.forTesting(PipelineOptionsFactory.as(DataflowPipelineOptions.class), "testStage");
    DataflowExecutionState state = executionContext.getExecutionStateRegistry().getState(NameContextsForTests.nameContextForTest(), "test", null, NoopProfileScope.NOOP);
    stateCloseable = executionContext.getExecutionStateTracker().enterState(state);
}

Example 7 with DataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState in project beam by apache.

the class DataflowWorkerLoggingHandler method publish.

@Override
public synchronized void publish(LogRecord record) {
    DataflowExecutionState currrentDataflowState = null;
    ExecutionState currrentState = ExecutionStateTracker.getCurrentExecutionState();
    if (currrentState instanceof DataflowExecutionState) {
        currrentDataflowState = (DataflowExecutionState) currrentState;
    }
    // It's okay to pass in the null state, publish() handles and tests this.
    publish(currrentDataflowState, record);
}

Example 8 with DataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState in project beam by apache.

the class DataflowSideInputReadCounter method updateCurrentStateIfOutdated.

private void updateCurrentStateIfOutdated() {
    DataflowExecutionState currentState = (DataflowExecutionState) executionContext.getExecutionStateTracker().getCurrentState();
    if (currentState == null || currentState.getStepName().originalName() == null || Objects.equals(latestConsumingStepName, currentState.getStepName())) {
        // In this case, the step has not changed, and we'll just return.
        return;
    }
    if (!byteCounters.containsKey(currentState.getStepName())) {
        byteCounters.put(currentState.getStepName(), executionContext.getCounterFactory().longSum(CounterName.named("read-sideinput-byte-count").withOriginalName(declaringOperationContext.nameContext()).withOrigin("SYSTEM").withOriginalRequestingStepName(currentState.getStepName().originalName()).withInputIndex(sideInputIndex)));
        executionStates.put(currentState.getStepName(), executionContext.getExecutionStateRegistry().getIOState(declaringOperationContext.nameContext(), "read-sideinput", currentState.getStepName().originalName(), sideInputIndex, currentState.getMetricsContainer(), currentState.getProfileScope()));
    }
    currentCounter = byteCounters.get(currentState.getStepName());
    currentExecutionState = executionStates.get(currentState.getStepName());
    latestConsumingStepName = currentState.getStepName();
}

Example 9 with DataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState in project beam by apache.

the class GroupingShuffleReaderTest method setCurrentExecutionState.

private void setCurrentExecutionState(String mockOriginalName) {
    DataflowExecutionState state = new TestDataflowExecutionState(NameContext.create(MOCK_STAGE_NAME, mockOriginalName, MOCK_SYSTEM_NAME, MOCK_USER_NAME), "activity");
    tracker.enterState(state);
}

Example 10 with DataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.DataflowOperationContext.DataflowExecutionState in project beam by apache.

the class IsmReaderFactoryTest method setUp.

@Before
public void setUp() {
    options = PipelineOptionsFactory.as(DataflowPipelineOptions.class);
    executionContext = BatchModeExecutionContext.forTesting(PipelineOptionsFactory.as(DataflowPipelineOptions.class), NameContextsForTests.nameContextForTest().stageName());
    DataflowExecutionState state = executionContext.getExecutionStateRegistry().getState(NameContextsForTests.nameContextForTest(), "test", null, /*container */
    NoopProfileScope.NOOP);
    operationContext = executionContext.createOperationContext(NameContextsForTests.nameContextForTest());
    stateCloseable = executionContext.getExecutionStateTracker().enterState(state);
}