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

Example 1 with DataflowExecutionState

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

the class ShuffleReadCounter method checkState.

@SuppressWarnings("ReferenceEquality")
@SuppressFBWarnings("ES_COMPARING_STRINGS_WITH_EQ")
private void checkState() {
    if (this.experimentEnabled) {
        ExecutionStateTracker.ExecutionState currentState = ExecutionStateTracker.getCurrentExecutionState();
        String currentStateName = null;
        if (currentState instanceof DataflowExecutionState) {
            currentStateName = ((DataflowExecutionState) currentState).getStepName().originalName();
        }
        if (this.currentCounter != null && currentStateName == this.currentCounter.getName().originalRequestingStepName()) {
            // If the step name of the state has not changed do not do another lookup.
            return;
        }
        CounterName name = ShuffleReadCounter.generateCounterName(this.originalShuffleStepName, currentStateName);
        this.currentCounter = this.counterSet.longSum(name);
    } else {
        this.currentCounter = this.legacyPerOperationPerDatasetBytesCounter;
    }
}

Example 2 with DataflowExecutionState

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

the class GroupingShuffleEntryIteratorTest 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 3 with DataflowExecutionState

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

the class IsmSideInputReaderTest method testIterableSideInputReadCounter.

@Test
public void testIterableSideInputReadCounter() throws Exception {
    // These are the expected msec and byte counters:
    CounterUpdate expectedSideInputMsecUpdate = new CounterUpdate().setStructuredNameAndMetadata(new CounterStructuredNameAndMetadata().setMetadata(new CounterMetadata().setKind(Kind.SUM.toString())).setName(new CounterStructuredName().setOrigin("SYSTEM").setName("read-sideinput-msecs").setOriginalStepName("originalName").setExecutionStepName("stageName").setOriginalRequestingStepName("originalName2").setInputIndex(1))).setCumulative(true).setInteger(new SplitInt64().setHighBits(0).setLowBits(0L));
    CounterName expectedCounterName = CounterName.named("read-sideinput-byte-count").withOriginalName(operationContext.nameContext()).withOrigin("SYSTEM").withOriginalRequestingStepName("originalName2").withInputIndex(1);
    // Test startup:
    Coder<WindowedValue<Long>> valueCoder = WindowedValue.getFullCoder(VarLongCoder.of(), GLOBAL_WINDOW_CODER);
    IsmRecordCoder<WindowedValue<Long>> ismCoder = IsmRecordCoder.of(1, 0, ImmutableList.of(GLOBAL_WINDOW_CODER, BigEndianLongCoder.of()), valueCoder);
    // Create a new state, which represents a step that receives the side input.
    DataflowExecutionState state2 = executionContext.getExecutionStateRegistry().getState(NameContext.create("stageName", "originalName2", "systemName2", "userName2"), "process", null, NoopProfileScope.NOOP);
    final List<KV<Long, WindowedValue<Long>>> firstElements = Arrays.asList(KV.of(0L, valueInGlobalWindow(0L)));
    final List<KV<Long, WindowedValue<Long>>> secondElements = new ArrayList<>();
    for (long i = 0; i < 100; i++) {
        secondElements.add(KV.of(i, valueInGlobalWindow(i * 10)));
    }
    final PCollectionView<Iterable<Long>> view = Pipeline.create().apply(Create.empty(VarLongCoder.of())).apply(View.asIterable());
    Source sourceA = initInputFile(fromKvsForList(firstElements), ismCoder);
    Source sourceB = initInputFile(fromKvsForList(secondElements), ismCoder);
    try (Closeable state2Closeable = executionContext.getExecutionStateTracker().enterState(state2)) {
        final IsmSideInputReader reader = serialSideInputReader(view.getTagInternal().getId(), sourceA, sourceB);
        // Store a strong reference to the returned value so that the logical reference
        // cache is not cleared for this test.
        Iterable<Long> value = reader.get(view, GlobalWindow.INSTANCE);
        verifyIterable(toValueList(concat(firstElements, secondElements)), value);
        // Assert that the same value reference was returned showing that it was cached.
        assertSame(reader.get(view, GlobalWindow.INSTANCE), value);
        Iterable<CounterUpdate> counterUpdates = executionContext.getExecutionStateRegistry().extractUpdates(true);
        assertThat(counterUpdates, hasItem(expectedSideInputMsecUpdate));
        Counter<?, ?> expectedCounter = counterFactory.getExistingCounter(expectedCounterName);
        assertNotNull(expectedCounter);
    }
}

Example 4 with DataflowExecutionState

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

the class DataflowSideInputReadCounterTest method testEnterEntersStateIfCalledFromTrackedThread.

@Test
public void testEnterEntersStateIfCalledFromTrackedThread() {
    DataflowExecutionContext mockedExecutionContext = mock(DataflowExecutionContext.class);
    DataflowOperationContext mockedOperationContext = mock(DataflowOperationContext.class);
    final int siIndexId = 3;
    ExecutionStateTracker mockedExecutionStateTracker = mock(ExecutionStateTracker.class);
    when(mockedExecutionContext.getExecutionStateTracker()).thenReturn(mockedExecutionStateTracker);
    Thread mockedThreadObject = Thread.currentThread();
    when(mockedExecutionStateTracker.getTrackedThread()).thenReturn(mockedThreadObject);
    DataflowExecutionState mockedExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateTracker.getCurrentState()).thenReturn(mockedExecutionState);
    NameContext mockedNameContext = mock(NameContext.class);
    when(mockedExecutionState.getStepName()).thenReturn(mockedNameContext);
    when(mockedNameContext.originalName()).thenReturn("DummyName");
    NameContext mockedDeclaringNameContext = mock(NameContext.class);
    when(mockedOperationContext.nameContext()).thenReturn(mockedDeclaringNameContext);
    when(mockedDeclaringNameContext.originalName()).thenReturn("DummyDeclaringName");
    CounterFactory mockedCounterFactory = mock(CounterFactory.class);
    when(mockedExecutionContext.getCounterFactory()).thenReturn(mockedCounterFactory);
    Counter<Long, Long> mockedCounter = mock(Counter.class);
    when(mockedCounterFactory.longSum(any())).thenReturn(mockedCounter);
    DataflowExecutionStateRegistry mockedExecutionStateRegistry = mock(DataflowExecutionStateRegistry.class);
    when(mockedExecutionContext.getExecutionStateRegistry()).thenReturn(mockedExecutionStateRegistry);
    DataflowExecutionState mockedCounterExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateRegistry.getIOState(any(), any(), any(), any(), any(), any())).thenReturn(mockedCounterExecutionState);
    DataflowSideInputReadCounter testObject = new DataflowSideInputReadCounter(mockedExecutionContext, mockedOperationContext, siIndexId);
    testObject.enter();
    verify(mockedExecutionStateTracker).enterState(mockedCounterExecutionState);
}

Example 5 with DataflowExecutionState

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

the class DataflowSideInputReadCounterTest method testAddBytesReadUpdatesCounter.

@Test
public void testAddBytesReadUpdatesCounter() {
    DataflowExecutionContext mockedExecutionContext = mock(DataflowExecutionContext.class);
    DataflowOperationContext mockedOperationContext = mock(DataflowOperationContext.class);
    final int siIndexId = 3;
    ExecutionStateTracker mockedExecutionStateTracker = mock(ExecutionStateTracker.class);
    when(mockedExecutionContext.getExecutionStateTracker()).thenReturn(mockedExecutionStateTracker);
    Thread mockedThreadObject = mock(Thread.class);
    when(mockedExecutionStateTracker.getTrackedThread()).thenReturn(mockedThreadObject);
    DataflowExecutionState mockedExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateTracker.getCurrentState()).thenReturn(mockedExecutionState);
    NameContext mockedNameContext = mock(NameContext.class);
    when(mockedExecutionState.getStepName()).thenReturn(mockedNameContext);
    when(mockedNameContext.originalName()).thenReturn("DummyName");
    NameContext mockedDeclaringNameContext = mock(NameContext.class);
    when(mockedOperationContext.nameContext()).thenReturn(mockedDeclaringNameContext);
    when(mockedDeclaringNameContext.originalName()).thenReturn("DummyDeclaringName");
    CounterFactory mockedCounterFactory = mock(CounterFactory.class);
    when(mockedExecutionContext.getCounterFactory()).thenReturn(mockedCounterFactory);
    Counter<Long, Long> mockedCounter = mock(Counter.class);
    when(mockedCounterFactory.longSum(any())).thenReturn(mockedCounter);
    when(mockedExecutionContext.getExecutionStateRegistry()).thenReturn(mock(DataflowExecutionStateRegistry.class));
    DataflowSideInputReadCounter testObject = new DataflowSideInputReadCounter(mockedExecutionContext, mockedOperationContext, siIndexId);
    testObject.addBytesRead(10l);
    verify(mockedCounter).addValue(10l);
}