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

Example 26 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest method testExecuteMapTaskExecutor.

@Test
public void testExecuteMapTaskExecutor() throws Exception {
    Operation o1 = Mockito.mock(Operation.class);
    Operation o2 = Mockito.mock(Operation.class);
    Operation o3 = Mockito.mock(Operation.class);
    List<Operation> operations = Arrays.asList(new Operation[] { o1, o2, o3 });
    ExecutionStateTracker stateTracker = Mockito.mock(ExecutionStateTracker.class);
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        executor.execute();
    }
    InOrder inOrder = Mockito.inOrder(stateTracker, o1, o2, o3);
    inOrder.verify(o3).start();
    inOrder.verify(o2).start();
    inOrder.verify(o1).start();
    inOrder.verify(o1).finish();
    inOrder.verify(o2).finish();
    inOrder.verify(o3).finish();
}

Example 27 with ExecutionStateTracker

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

the class DataflowSideInputReadCounterTest method testAddBytesReadSkipsAddingCountersIfCurrentCounterIsNull.

@Test
public void testAddBytesReadSkipsAddingCountersIfCurrentCounterIsNull() {
    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);
    DataflowSideInputReadCounter testObject = new DataflowSideInputReadCounter(mockedExecutionContext, mockedOperationContext, siIndexId);
    try {
        testObject.addBytesRead(10);
    } catch (Exception e) {
        fail("Supposedly, we tried to add bytes to counter and that shouldn't happen. Ex.: " + e.toString());
    }
}

Example 28 with ExecutionStateTracker

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

the class DataflowSideInputReadCounterTest method testEnterDoesntEnterStateIfCalledFromDifferentThread.

@Test
public void testEnterDoesntEnterStateIfCalledFromDifferentThread() {
    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);
    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, never()).enterState(any());
}

Example 29 with ExecutionStateTracker

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

the class BatchModeExecutionContextTest method stateSamplingInBatch.

@Test(timeout = 2000)
public void stateSamplingInBatch() {
    // Test that when writing on one thread and reading from another, updates always eventually
    // reach the reading thread.
    BatchModeExecutionState state = new BatchModeExecutionState(NameContextsForTests.nameContextForTest(), "testState", null, /* requestingStepName */
    null, /* inputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP);
    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
        long value = 0;
        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(value));
            value = newValue;
        }
    } finally {
        sampler.stop();
    }
}

Example 30 with ExecutionStateTracker

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

the class MapTaskExecutorTest method testExceptionInAbortSuppressed.

@Test
public void testExceptionInAbortSuppressed() throws Exception {
    Operation o1 = Mockito.mock(Operation.class);
    Operation o2 = Mockito.mock(Operation.class);
    Operation o3 = Mockito.mock(Operation.class);
    Operation o4 = Mockito.mock(Operation.class);
    Mockito.doThrow(new Exception("in finish")).when(o2).finish();
    Mockito.doThrow(new Exception("suppressed in abort")).when(o3).abort();
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (MapTaskExecutor executor = new MapTaskExecutor(Arrays.<Operation>asList(o1, o2, o3, o4), counterSet, stateTracker)) {
        executor.execute();
        fail("Should have thrown");
    } catch (Exception e) {
        InOrder inOrder = Mockito.inOrder(o1, o2, o3, o4);
        inOrder.verify(o4).start();
        inOrder.verify(o3).start();
        inOrder.verify(o2).start();
        inOrder.verify(o1).start();
        inOrder.verify(o1).finish();
        // this fails
        inOrder.verify(o2).finish();
        // Order of abort doesn't matter
        Mockito.verify(o1).abort();
        Mockito.verify(o2).abort();
        // will throw an exception, but we shouldn't fail
        Mockito.verify(o3).abort();
        Mockito.verify(o4).abort();
        Mockito.verifyNoMoreInteractions(o1, o2, o3, o4);
        // Make sure the failure while aborting shows up as a suppressed error
        assertThat(e.getMessage(), equalTo("in finish"));
        assertThat(e.getSuppressed(), arrayWithSize(1));
        assertThat(e.getSuppressed()[0].getMessage(), equalTo("suppressed in abort"));
    }
}