BatchDataflowWorker implements Closeable { @VisibleForTesting boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) throws IOException { LOG.debug("Executing: {}", workItem); DataflowWorkExecutor worker = null; SdkWorkerHarness sdkWorkerHarness = sdkHarnessRegistry.getAvailableWorkerAndAssignWork(); try { options.setProject(workItem.getProjectId()); final String stageName; if (workItem.getMapTask() != null) { stageName = workItem.getMapTask().getStageName(); } else if (workItem.getSourceOperationTask() != null) { stageName = workItem.getSourceOperationTask().getStageName(); } else { throw new RuntimeException("Unknown kind of work item: " + workItem.toString()); } CounterSet counterSet = new CounterSet(); BatchModeExecutionContext executionContext = BatchModeExecutionContext.create( counterSet, sideInputDataCache, sideInputWeakReferenceCache, readerRegistry, options, stageName, String.valueOf(workItem.getId())); if (workItem.getMapTask() != null) { MutableNetwork<Node, Edge> network = mapTaskToNetwork.apply(workItem.getMapTask()); if (LOG.isDebugEnabled()) { LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(network)); } worker = mapTaskExecutorFactory.create( sdkWorkerHarness.getControlClientHandler(), sdkWorkerHarness.getGrpcDataFnServer(), sdkHarnessRegistry.beamFnDataApiServiceDescriptor(), sdkWorkerHarness.getGrpcStateFnServer(), network, options, stageName, readerRegistry, sinkRegistry, executionContext, counterSet, idGenerator); } else if (workItem.getSourceOperationTask() != null) { worker = SourceOperationExecutorFactory.create( options, workItem.getSourceOperationTask(), counterSet, executionContext, stageName); } else { throw new IllegalStateException("Work Item was neither a MapTask nor a SourceOperation"); } workItemStatusClient.setWorker(worker, executionContext); DataflowWorkProgressUpdater progressUpdater = new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker); executeWork(worker, progressUpdater); workItemStatusClient.reportSuccess(); return true; } catch (Throwable e) { workItemStatusClient.reportError(e); return false; } finally { if (worker != null) { try { worker.close(); } catch (Exception exn) { LOG.warn( "Uncaught exception while closing worker. All work has already committed or " + "been marked for retry.", exn); } } if (sdkWorkerHarness != null) { sdkHarnessRegistry.completeWork(sdkWorkerHarness); } } } protected  BatchDataflowWorker(
      RunnerApi.@Nullable Pipeline pipeline,
      SdkHarnessRegistry sdkHarnessRegistry,
      WorkUnitClient workUnitClient,
      DataflowMapTaskExecutorFactory mapTaskExecutorFactory,
      DataflowWorkerHarnessOptions options); boolean getAndPerformWork(); void startStatusServer(); @Override void close();  }
@Test public void testWhenProcessingWorkUnitFailsWeReportStatus() throws Exception { BatchDataflowWorker worker = new BatchDataflowWorker( null , SdkHarnessRegistries.emptySdkHarnessRegistry(), mockWorkUnitClient, IntrinsicMapTaskExecutorFactory.defaultFactory(), options); long initialReportIndex = 4L; WorkItem workItem = new WorkItem() .setId(1L) .setJobId("Expected to fail the job") .setInitialReportIndex(initialReportIndex); WorkItemStatusClient workItemStatusClient = mock(WorkItemStatusClient.class); worker.doWork(workItem, workItemStatusClient); ArgumentCaptor<Throwable> errorCaptor = ArgumentCaptor.forClass(Throwable.class); verify(workItemStatusClient).reportError(errorCaptor.capture()); Throwable error = errorCaptor.getValue(); assertThat(error, notNullValue()); assertThat(error.getMessage(), equalTo("Unknown kind of work item: " + workItem.toString())); }