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

Example 6 with ExecutionStateTracker

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

the class MapTaskExecutorTest method testNoReadOperation.

@Test
public void testNoReadOperation() throws Exception {
    // Test MapTaskExecutor without ReadOperation.
    List<Operation> operations = Arrays.<Operation>asList(createOperation("o1", 1), createOperation("o2", 2));
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) {
        thrown.expect(IllegalStateException.class);
        thrown.expectMessage("is not a ReadOperation");
        executor.getReadOperation();
    }
}

Example 7 with ExecutionStateTracker

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

the class StreamingModeExecutionContextTest method stateSamplingInStreaming.

@Test(timeout = 2000)
public void stateSamplingInStreaming() {
    // Test that when writing on one thread and reading from another, updates always eventually
    // reach the reading thread.
    StreamingModeExecutionState state = new StreamingModeExecutionState(NameContextsForTests.nameContextForTest(), "testState", null, NoopProfileScope.NOOP, null);
    ExecutionStateSampler sampler = ExecutionStateSampler.newForTest();
    try {
        sampler.start();
        ExecutionStateTracker tracker = new ExecutionStateTracker(sampler);
        Thread executionThread = new Thread();
        executionThread.setName("looping-thread-for-test");
        tracker.activate(executionThread);
        tracker.enterState(state);
        // Wait for the state to be incremented 3 times
        for (int i = 0; i < 3; i++) {
            CounterUpdate update = null;
            while (update == null) {
                update = state.extractUpdate(false);
            }
            long newValue = splitIntToLong(update.getInteger());
            assertThat(newValue, Matchers.greaterThan(0L));
        }
    } finally {
        sampler.stop();
    }
}

Example 8 with ExecutionStateTracker

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

the class SimpleParDoFnTest method testStateTracking.

@Test
public void testStateTracking() throws Exception {
    ExecutionStateTracker tracker = ExecutionStateTracker.newForTest();
    TestOperationContext operationContext = TestOperationContext.create(new CounterSet(), NameContextsForTests.nameContextForTest(), new MetricsContainerImpl(NameContextsForTests.ORIGINAL_NAME), tracker);
    class StateTestingDoFn extends DoFn<Integer, String> {

        private boolean startCalled = false;

        @StartBundle
        public void startBundle() throws Exception {
            startCalled = true;
            assertThat(tracker.getCurrentState(), equalTo(operationContext.getStartState()));
        }

        @ProcessElement
        public void processElement(ProcessContext c) throws Exception {
            assertThat(startCalled, equalTo(true));
            assertThat(tracker.getCurrentState(), equalTo(operationContext.getProcessState()));
        }
    }
    StateTestingDoFn fn = new StateTestingDoFn();
    DoFnInfo<?, ?> fnInfo = DoFnInfo.forFn(fn, WindowingStrategy.globalDefault(), null, /* side input views */
    null, /* input coder */
    MAIN_OUTPUT, DoFnSchemaInformation.create(), Collections.emptyMap());
    ParDoFn userParDoFn = new SimpleParDoFn<>(options, DoFnInstanceManagers.singleInstance(fnInfo), NullSideInputReader.empty(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0, new TupleTag<>("declared"), 1), BatchModeExecutionContext.forTesting(options, operationContext.counterFactory(), "testStage").getStepContext(operationContext), operationContext, DoFnSchemaInformation.create(), Collections.emptyMap(), SimpleDoFnRunnerFactory.INSTANCE);
    // This test ensures proper behavior of the state sampling even with lazy initialization.
    try (Closeable trackerCloser = tracker.activate()) {
        try (Closeable processCloser = operationContext.enterProcess()) {
            userParDoFn.processElement(WindowedValue.valueInGlobalWindow(5));
        }
    }
}

Example 9 with ExecutionStateTracker

use of org.apache.beam.runners.core.metrics.ExecutionStateTracker 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 10 with ExecutionStateTracker

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