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

Example 1 with ExecutionStateTracker

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

the class BeamFnStatusClientTest method testActiveBundleState.

@Test
public void testActiveBundleState() {
    ProcessBundleHandler handler = mock(ProcessBundleHandler.class);
    BundleProcessorCache processorCache = mock(BundleProcessorCache.class);
    Map<String, BundleProcessor> bundleProcessorMap = new HashMap<>();
    for (int i = 0; i < 11; i++) {
        BundleProcessor processor = mock(BundleProcessor.class);
        ExecutionStateTracker executionStateTracker = mock(ExecutionStateTracker.class);
        when(processor.getStateTracker()).thenReturn(executionStateTracker);
        when(executionStateTracker.getMillisSinceLastTransition()).thenReturn(Integer.toUnsignedLong((10 - i) * 1000));
        when(executionStateTracker.getTrackedThread()).thenReturn(Thread.currentThread());
        String instruction = Integer.toString(i);
        when(processorCache.find(instruction)).thenReturn(processor);
        bundleProcessorMap.put(instruction, processor);
    }
    when(handler.getBundleProcessorCache()).thenReturn(processorCache);
    when(processorCache.getActiveBundleProcessors()).thenReturn(bundleProcessorMap);
    ManagedChannelFactory channelFactory = ManagedChannelFactory.createInProcess();
    BeamFnStatusClient client = new BeamFnStatusClient(apiServiceDescriptor, channelFactory::forDescriptor, handler.getBundleProcessorCache(), PipelineOptionsFactory.create(), Caches.noop());
    StringJoiner joiner = new StringJoiner("\n");
    joiner.add(client.getActiveProcessBundleState());
    String actualState = joiner.toString();
    List<String> expectedInstructions = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
        expectedInstructions.add(String.format("Instruction %d", i));
    }
    assertThat(actualState, stringContainsInOrder(expectedInstructions));
    assertThat(actualState, not(containsString("Instruction 10")));
}

Example 2 with ExecutionStateTracker

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

the class MapTaskExecutorTest 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 (MapTaskExecutor executor = new MapTaskExecutor(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 3 with ExecutionStateTracker

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

the class MapTaskExecutorTest method testNoOperation.

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

Example 4 with ExecutionStateTracker

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

the class MapTaskExecutorTest 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 (MapTaskExecutor executor = new MapTaskExecutor(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 5 with ExecutionStateTracker

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

the class MapTaskExecutorTest method testAbort.

@Test
public void testAbort() throws Exception {
    // Operation must be an instance of ReadOperation or ReceivingOperation per preconditions
    // in MapTaskExecutor.
    ReadOperation o1 = Mockito.mock(ReadOperation.class);
    ReadOperation o2 = Mockito.mock(ReadOperation.class);
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    MapTaskExecutor executor = new MapTaskExecutor(Arrays.<Operation>asList(o1, o2), counterSet, stateTracker);
    Mockito.doAnswer(invocation -> {
        executor.abort();
        return null;
    }).when(o1).finish();
    executor.execute();
    Mockito.verify(o1, atLeastOnce()).abortReadLoop();
    Mockito.verify(o2, atLeastOnce()).abortReadLoop();
}