Examples with Operation org.apache.beam.runners.dataflow.worker.util.common.worker.Operation used on opensource projects

Example 6 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation in project beam by apache.

the class IntrinsicMapTaskExecutorTest 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 (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(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"));
    }
}

Example 7 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation in project beam by apache.

the class IntrinsicMapTaskExecutorFactory method create.

/**
 * Creates a new {@link DataflowMapTaskExecutor} from the given {@link MapTask} definition using
 * the provided {@link ReaderFactory}.
 */
@Override
public DataflowMapTaskExecutor create(InstructionRequestHandler instructionRequestHandler, GrpcFnServer<GrpcDataService> grpcDataFnServer, Endpoints.ApiServiceDescriptor dataApiServiceDescriptor, GrpcFnServer<GrpcStateService> grpcStateFnServer, MutableNetwork<Node, Edge> network, PipelineOptions options, String stageName, ReaderFactory readerFactory, SinkFactory sinkFactory, DataflowExecutionContext<?> executionContext, CounterSet counterSet, IdGenerator idGenerator) {
    // TODO: remove this once we trust the code paths
    checkArgument(!DataflowRunner.hasExperiment(options.as(DataflowPipelineDebugOptions.class), "beam_fn_api"), "experiment beam_fn_api turned on but non-Fn API MapTaskExecutorFactory invoked");
    // Swap out all the InstructionOutput nodes with OutputReceiver nodes
    Networks.replaceDirectedNetworkNodes(network, createOutputReceiversTransform(stageName, counterSet));
    // Swap out all the ParallelInstruction nodes with Operation nodes
    Networks.replaceDirectedNetworkNodes(network, createOperationTransformForParallelInstructionNodes(stageName, network, options, readerFactory, sinkFactory, executionContext));
    // Collect all the operations within the network and attach all the operations as receivers
    // to preceding output receivers.
    List<Operation> topoSortedOperations = new ArrayList<>();
    for (OperationNode node : Iterables.filter(Networks.topologicalOrder(network), OperationNode.class)) {
        topoSortedOperations.add(node.getOperation());
        for (Node predecessor : Iterables.filter(network.predecessors(node), OutputReceiverNode.class)) {
            ((OutputReceiverNode) predecessor).getOutputReceiver().addOutput((Receiver) node.getOperation());
        }
    }
    if (LOG.isDebugEnabled()) {
        LOG.info("Map task network: {}", Networks.toDot(network));
    }
    return IntrinsicMapTaskExecutor.withSharedCounterSet(topoSortedOperations, counterSet, executionContext.getExecutionStateTracker());
}

Example 8 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation 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 9 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation 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 10 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation 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();
    }
}