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();
}
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());
}
}
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());
}
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();
}
}
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"));
}
}
Aggregations