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

Example 21 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest method testGetOutputCounters.

@Test
@SuppressWarnings("unchecked")
public void testGetOutputCounters() throws Exception {
    List<Operation> operations = Arrays.asList(new Operation[] { createOperation("o1", 1), createOperation("o2", 2), createOperation("o3", 3) });
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        CounterSet counterSet = executor.getOutputCounters();
        CounterUpdateExtractor<?> updateExtractor = Mockito.mock(CounterUpdateExtractor.class);
        counterSet.extractUpdates(false, updateExtractor);
        verify(updateExtractor).longSum(eq(named("test-o1-ElementCount")), anyBoolean(), eq(1L));
        verify(updateExtractor).longSum(eq(named("test-o2-ElementCount")), anyBoolean(), eq(2L));
        verify(updateExtractor).longSum(eq(named("test-o3-ElementCount")), anyBoolean(), eq(3L));
        verifyNoMoreInteractions(updateExtractor);
    }
}

Example 22 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest method testGetProgressAndRequestSplit.

@Test
public void testGetProgressAndRequestSplit() throws Exception {
    TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContextsForTests.nameContextForTest());
    TestReadOperation operation = new TestReadOperation(receiver, createContext("ReadOperation"));
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(Arrays.asList(new Operation[] { operation }), counterSet, stateTracker)) {
        operation.setProgress(approximateProgressAtIndex(1L));
        Assert.assertEquals(positionAtIndex(1L), positionFromProgress(executor.getWorkerProgress()));
        Assert.assertEquals(positionAtIndex(1L), positionFromSplitResult(executor.requestDynamicSplit(splitRequestAtIndex(1L))));
    }
}

Example 23 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest method testNoOperation.

@Test
public void testNoOperation() throws Exception {
    // Test MapTaskExecutor without a single operation.
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(new ArrayList<Operation>(), counterSet, stateTracker)) {
        thrown.expect(IllegalStateException.class);
        thrown.expectMessage("has no operation");
        executor.getReadOperation();
    }
}

Example 24 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest 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 (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        thrown.expect(IllegalStateException.class);
        thrown.expectMessage("is not a ReadOperation");
        executor.getReadOperation();
    }
}

Example 25 with ExecutionStateTracker

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

the class IntrinsicMapTaskExecutorTest method testExceptionInStartAbortsAllOperations.

@Test
public void testExceptionInStartAbortsAllOperations() throws Exception {
    Operation o1 = Mockito.mock(Operation.class);
    Operation o2 = Mockito.mock(Operation.class);
    Operation o3 = Mockito.mock(Operation.class);
    Mockito.doThrow(new Exception("in start")).when(o2).start();
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(Arrays.<Operation>asList(o1, o2, o3), counterSet, stateTracker)) {
        executor.execute();
        fail("Should have thrown");
    } catch (Exception e) {
        InOrder inOrder = Mockito.inOrder(o1, o2, o3);
        inOrder.verify(o3).start();
        inOrder.verify(o2).start();
        // Order of abort doesn't matter
        Mockito.verify(o1).abort();
        Mockito.verify(o2).abort();
        Mockito.verify(o3).abort();
        Mockito.verifyNoMoreInteractions(o1, o2, o3);
    }
}