IntrinsicMapTaskExecutorFactory implements DataflowMapTaskExecutorFactory { @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) { checkArgument( !DataflowRunner.hasExperiment( options.as(DataflowPipelineDebugOptions.class), "beam_fn_api"), "experiment beam_fn_api turned on but non-Fn API MapTaskExecutorFactory invoked"); Networks.replaceDirectedNetworkNodes( network, createOutputReceiversTransform(stageName, counterSet)); Networks.replaceDirectedNetworkNodes( network, createOperationTransformForParallelInstructionNodes( stageName, network, options, readerFactory, sinkFactory, executionContext)); 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()); } private  IntrinsicMapTaskExecutorFactory(); static IntrinsicMapTaskExecutorFactory defaultFactory(); @Override 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);  }
@Test public void testCreateMapTaskExecutor() throws Exception { List<ParallelInstruction> instructions = Arrays.asList( createReadInstruction("Read"), createParDoInstruction(0, 0, "DoFn1"), createParDoInstruction(0, 0, "DoFnWithContext"), createFlattenInstruction(1, 0, 2, 0, "Flatten"), createWriteInstruction(3, 0, "Write")); MapTask mapTask = new MapTask(); mapTask.setStageName(STAGE); mapTask.setSystemName("systemName"); mapTask.setInstructions(instructions); mapTask.setFactory(Transport.getJsonFactory()); try (DataflowMapTaskExecutor executor = mapTaskExecutorFactory.create( null , null , null , null, mapTaskToNetwork.apply(mapTask), options, STAGE, readerRegistry, sinkRegistry, BatchModeExecutionContext.forTesting(options, counterSet, "testStage"), counterSet, idGenerator)) { @SuppressWarnings({"rawtypes", "unchecked"}) List<Object> operations = (List) executor.operations; assertThat( operations, hasItems( instanceOf(ReadOperation.class), instanceOf(ParDoOperation.class), instanceOf(ParDoOperation.class), instanceOf(FlattenOperation.class), instanceOf(WriteOperation.class))); ReadOperation readOperation = Iterables.getOnlyElement(Iterables.filter(operations, ReadOperation.class)); assertEquals(2, readOperation.receivers[0].getReceiverCount()); FlattenOperation flattenOperation = Iterables.getOnlyElement(Iterables.filter(operations, FlattenOperation.class)); for (ParDoOperation operation : Iterables.filter(operations, ParDoOperation.class)) { assertSame(flattenOperation, operation.receivers[0].getOnlyReceiver()); } WriteOperation writeOperation = Iterables.getOnlyElement(Iterables.filter(operations, WriteOperation.class)); assertSame(writeOperation, flattenOperation.receivers[0].getOnlyReceiver()); } @SuppressWarnings("unchecked") Counter<Long, ?> otherMsecCounter = (Counter<Long, ?>) counterSet.getExistingCounter("test-other-msecs"); assertNull(otherMsecCounter); counterSet.extractUpdates(false, updateExtractor); verifyOutputCounters( updateExtractor, "read_output_name", "DoFn1_output", "DoFnWithContext_output", "flatten_output_name"); verify(updateExtractor).longSum(eq(named("Read-ByteCount")), anyBoolean(), anyLong()); verify(updateExtractor).longSum(eq(named("Write-ByteCount")), anyBoolean(), anyLong()); verifyNoMoreInteractions(updateExtractor); }
@Test public void testExecutionContextPlumbing() throws Exception { List<ParallelInstruction> instructions = Arrays.asList( createReadInstruction("Read", ReaderFactoryTest.SingletonTestReaderFactory.class), createParDoInstruction(0, 0, "DoFn1", "DoFnUserName"), createParDoInstruction(1, 0, "DoFnWithContext", "DoFnWithContextUserName")); MapTask mapTask = new MapTask(); mapTask.setStageName(STAGE); mapTask.setInstructions(instructions); mapTask.setFactory(Transport.getJsonFactory()); BatchModeExecutionContext context = BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); try (DataflowMapTaskExecutor executor = mapTaskExecutorFactory.create( null , null , null , null, mapTaskToNetwork.apply(mapTask), options, STAGE, readerRegistry, sinkRegistry, context, counterSet, idGenerator)) { executor.execute(); } List<String> stepNames = new ArrayList<>(); for (BatchModeExecutionContext.StepContext stepContext : context.getAllStepContexts()) { stepNames.add(stepContext.getNameContext().systemName()); } assertThat(stepNames, hasItems("DoFn1", "DoFnWithContext")); }