MapTaskExecutor implements WorkExecutor { @Override public void execute() throws Exception { LOG.debug("Executing map task"); try (Closeable stateCloser = executionStateTracker.activate()) { try { LOG.debug("Starting operations"); ListIterator<Operation> iterator = operations.listIterator(operations.size()); while (iterator.hasPrevious()) { Operation op = iterator.previous(); op.start(); } LOG.debug("Finishing operations"); for (Operation op : operations) { op.finish(); } } catch (Exception | Error exn) { LOG.debug("Aborting operations", exn); for (Operation op : operations) { try { op.abort(); } catch (Exception | Error exn2) { exn.addSuppressed(exn2); if (exn2 instanceof InterruptedException) { Thread.currentThread().interrupt(); } } } throw exn; } } LOG.debug("Map task execution complete"); } MapTaskExecutor(
      List<Operation> operations,
      CounterSet counters,
      ExecutionStateTracker executionStateTracker); @Override CounterSet getOutputCounters(); @Override void execute(); @Override NativeReader.Progress getWorkerProgress(); @Override NativeReader.@Nullable DynamicSplitResult requestCheckpoint(); @Override NativeReader.DynamicSplitResult requestDynamicSplit(
      NativeReader.DynamicSplitRequest splitRequest); ReadOperation getReadOperation(); boolean supportsRestart(); @Override void abort(); @Override List<Integer> reportProducedEmptyOutput(); final List<Operation> operations; }
@Test public void testExecuteMapTaskExecutor() throws Exception { List<String> log = new ArrayList<>(); 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(); }
@Test public void testPerElementProcessingTimeCounters() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); options .as(DataflowPipelineDebugOptions.class) .setExperiments( Lists.newArrayList(DataflowElementExecutionTracker.TIME_PER_ELEMENT_EXPERIMENT)); ExecutionStateSampler stateSampler = ExecutionStateSampler.newForTest(); DataflowExecutionStateTracker stateTracker = new DataflowExecutionStateTracker( stateSampler, new TestDataflowExecutionState( NameContext.forStage("test-stage"), "other", null , null , null , NoopProfileScope.NOOP), counterSet, options, "test-work-item-id"); NameContext parDoName = nameForStep("s1"); ReadOperation read = ReadOperation.forTest( new TestReader("a", "b", "c"), new OutputReceiver(), TestOperationContext.create(counterSet, nameForStep("s0"), null, stateTracker)); ParDoOperation parDo = new ParDoOperation( new NoopParDoFn(), new OutputReceiver[0], TestOperationContext.create(counterSet, parDoName, null, stateTracker)); parDo.attachInput(read, 0); List<Operation> operations = Lists.newArrayList(read, parDo); try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { executor.execute(); } stateSampler.doSampling(100L); CounterName counterName = CounterName.named("per-element-processing-time").withOriginalName(parDoName); Counter<Long, CounterDistribution> counter = (Counter<Long, CounterDistribution>) counterSet.getExistingCounter(counterName); assertThat(counter.getAggregate().getCount(), equalTo(3L)); }
@Test @SuppressWarnings("unchecked") public void testGetMetricContainers() throws Exception { ExecutionStateTracker stateTracker = new DataflowExecutionStateTracker( ExecutionStateSampler.newForTest(), new TestDataflowExecutionState( NameContext.forStage("testStage"), "other", null , null , null , NoopProfileScope.NOOP), new CounterSet(), PipelineOptionsFactory.create(), "test-work-item-id"); final String o1 = "o1"; TestOperationContext context1 = createContext(o1, stateTracker); final String o2 = "o2"; TestOperationContext context2 = createContext(o2, stateTracker); final String o3 = "o3"; TestOperationContext context3 = createContext(o3, stateTracker); List<Operation> operations = Arrays.asList( new Operation(new OutputReceiver[] {}, context1) { @Override public void start() throws Exception { super.start(); try (Closeable scope = context.enterStart()) { Metrics.counter("TestMetric", "MetricCounter").inc(1L); } } }, new Operation(new OutputReceiver[] {}, context2) { @Override public void start() throws Exception { super.start(); try (Closeable scope = context.enterStart()) { Metrics.counter("TestMetric", "MetricCounter").inc(2L); } } }, new Operation(new OutputReceiver[] {}, context3) { @Override public void start() throws Exception { super.start(); try (Closeable scope = context.enterStart()) { Metrics.counter("TestMetric", "MetricCounter").inc(3L); } } }); try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) { executor.execute(); assertThat( context1.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o1, 1L))); assertThat( context2.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o2, 2L))); assertThat( context3.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o3, 3L))); } }