DataflowWorkUnitClient implements WorkUnitClient { @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); List<String> capabilities = ImmutableList.<String>of( options.getWorkerId(), CAPABILITY_REMOTE_SOURCE, PropertyNames.CUSTOM_SOURCE_FORMAT); Optional<WorkItem> workItem = getWorkItemInternal(workItemTypes, capabilities); if (!workItem.isPresent()) { return Optional.absent(); } if (workItem.isPresent() && workItem.get().getId() == null) { logger.debug("Discarding invalid work item {}", workItem.orNull()); return Optional.absent(); } 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; } DataflowWorkUnitClient(DataflowWorkerHarnessOptions options, Logger logger); @Override Optional<WorkItem> getWorkItem(); @Override Optional<WorkItem> getGlobalStreamingConfigWorkItem(); @Override Optional<WorkItem> getStreamingConfigWorkItem(String computationId); @Override WorkItemServiceState reportWorkItemStatus(WorkItemStatus workItemStatus);  }
@Test public void testCloudServiceCall() throws Exception { WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID); when(request.execute()).thenReturn(generateMockResponse(workItem)); WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); assertEquals(Optional.of(workItem), client.getWorkItem()); LeaseWorkItemRequest actualRequest = Transport.getJsonFactory() .fromString(request.getContentAsString(), LeaseWorkItemRequest.class); assertEquals(WORKER_ID, actualRequest.getWorkerId()); assertEquals( ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities()); assertEquals( ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes()); assertEquals("1234", DataflowWorkerLoggingMDC.getWorkId()); }
@Test public void testCloudServiceCallMapTaskStagePropagation() throws Exception { WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); final String stageName = "test_stage_name"; MapTask mapTask = new MapTask(); mapTask.setStageName(stageName); WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID); workItem.setMapTask(mapTask); when(request.execute()).thenReturn(generateMockResponse(workItem)); assertEquals(Optional.of(workItem), client.getWorkItem()); assertEquals(stageName, DataflowWorkerLoggingMDC.getStageName()); }
@Test public void testCloudServiceCallSeqMapTaskStagePropagation() throws Exception { WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); final String stageName = "test_stage_name"; SeqMapTask seqMapTask = new SeqMapTask(); seqMapTask.setStageName(stageName); WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID); workItem.setSeqMapTask(seqMapTask); when(request.execute()).thenReturn(generateMockResponse(workItem)); assertEquals(Optional.of(workItem), client.getWorkItem()); assertEquals(stageName, DataflowWorkerLoggingMDC.getStageName()); }
@Test public void testCloudServiceCallNoWorkPresent() throws Exception { WorkItem workItem = new WorkItem(); when(request.execute()).thenReturn(generateMockResponse(workItem)); WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); assertEquals(Optional.absent(), client.getWorkItem()); LeaseWorkItemRequest actualRequest = Transport.getJsonFactory() .fromString(request.getContentAsString(), LeaseWorkItemRequest.class); assertEquals(WORKER_ID, actualRequest.getWorkerId()); assertEquals( ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities()); assertEquals( ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes()); }
@Test public void testCloudServiceCallNoWorkId() throws Exception { WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID); workItem.setId(null); when(request.execute()).thenReturn(generateMockResponse(workItem)); WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); assertEquals(Optional.absent(), client.getWorkItem()); LeaseWorkItemRequest actualRequest = Transport.getJsonFactory() .fromString(request.getContentAsString(), LeaseWorkItemRequest.class); assertEquals(WORKER_ID, actualRequest.getWorkerId()); assertEquals( ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities()); assertEquals( ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes()); }
@Test public void testCloudServiceCallNoWorkItem() throws Exception { when(request.execute()).thenReturn(generateMockResponse()); WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); assertEquals(Optional.absent(), client.getWorkItem()); LeaseWorkItemRequest actualRequest = Transport.getJsonFactory() .fromString(request.getContentAsString(), LeaseWorkItemRequest.class); assertEquals(WORKER_ID, actualRequest.getWorkerId()); assertEquals( ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities()); assertEquals( ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes()); }
@Test public void testCloudServiceCallMultipleWorkItems() throws Exception { expectedException.expect(IOException.class); expectedException.expectMessage( "This version of the SDK expects no more than one work item from the service"); WorkItem workItem1 = createWorkItem(PROJECT_ID, JOB_ID); WorkItem workItem2 = createWorkItem(PROJECT_ID, JOB_ID); when(request.execute()).thenReturn(generateMockResponse(workItem1, workItem2)); WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG); client.getWorkItem(); }