@Override
  public Optional<WorkItem> getWorkItem() throws IOException {
    List<String> workItemTypes =
        ImmutableList.of(
            WORK_ITEM_TYPE_MAP_TASK,
            WORK_ITEM_TYPE_SEQ_MAP_TASK,
            WORK_ITEM_TYPE_REMOTE_SOURCE_TASK);
    // All remote sources require the "remote_source" capability. Dataflow's
    // custom sources are further tagged with the format "custom_source".
    List<String> capabilities =
        new ArrayList<String>(
            Arrays.asList(
                options.getWorkerId(),
                CAPABILITY_REMOTE_SOURCE,
                PropertyNames.CUSTOM_SOURCE_FORMAT));
    if (options.getWorkerPool() != null) {
      capabilities.add(options.getWorkerPool());
    }

    Optional<WorkItem> workItem = getWorkItemInternal(workItemTypes, capabilities);
    if (!workItem.isPresent()) {
      // Normal case, this means that the response contained no work, i.e. no work is available
      // at this time.
      return Optional.empty();
    }
    if (workItem.get().getId() == null) {
      logger.debug("Discarding invalid work item {}", workItem.get());
      return Optional.empty();
    }

    WorkItem work = workItem.get();

    final String stage;
    if (work.getMapTask() != null) {
      stage = work.getMapTask().getStageName();
      logger.info("Starting MapTask stage {}", stage);
    } else if (work.getSeqMapTask() != null) {
      stage = work.getSeqMapTask().getStageName();
      logger.info("Starting SeqMapTask stage {}", stage);
    } else if (work.getSourceOperationTask() != null) {
      stage = work.getSourceOperationTask().getStageName();
      logger.info("Starting SourceOperationTask stage {}", stage);
    } else {
      stage = null;
    }
    DataflowWorkerLoggingMDC.setStageName(stage);

    stageStartTime.set(DateTime.now());
    DataflowWorkerLoggingMDC.setWorkId(Long.toString(work.getId()));

    return workItem;
  }
@Test
  public void testCloudServiceCallSeqMapTaskStagePropagation() throws Exception {
    // Publish and acquire a seq map task work item, and verify we're now processing that stage.
    final String stageName = "test_stage_name";
    SeqMapTask seqMapTask = new SeqMapTask();
    seqMapTask.setStageName(stageName);
    WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
    workItem.setSeqMapTask(seqMapTask);

    MockLowLevelHttpResponse response = generateMockResponse(workItem);
    MockLowLevelHttpRequest request = new MockLowLevelHttpRequest().setResponse(response);
    MockHttpTransport transport =
        new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build();
    DataflowWorkerHarnessOptions pipelineOptions = createPipelineOptionsWithTransport(transport);
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);

    assertEquals(Optional.of(workItem), client.getWorkItem());
    assertEquals(stageName, DataflowWorkerLoggingMDC.getStageName());
  }