public TaskContext(WorkUnitState workUnitState) { this.taskState = new TaskState(workUnitState); this.taskMetrics = TaskMetrics.get(this.taskState); this.taskState.setProp(Instrumented.METRIC_CONTEXT_NAME_KEY, this.taskMetrics.getName()); }
@Override public void forWorkUnit(WorkUnit workUnit) { jobState.incrementTaskCount(); jobState.addTaskState(new TaskState(new WorkUnitState(workUnit, jobState))); } });
private void getTaskStateWithCommonAndSpecWuProps(int numTaskStates, DataInput in) throws IOException { Properties commonWuProps = new Properties(); for (int i = 0; i < numTaskStates; i++) { TaskState taskState = new TaskState(); taskState.readFields(in); if (i == 0) { commonWuProps.putAll(taskState.getWorkunit().getProperties()); } else { Properties newCommonWuProps = new Properties(); newCommonWuProps .putAll(Maps.difference(commonWuProps, taskState.getWorkunit().getProperties()).entriesInCommon()); commonWuProps = newCommonWuProps; } this.taskStates.put(taskState.getTaskId().intern(), taskState); } ImmutableProperties immutableCommonProperties = new ImmutableProperties(commonWuProps); for (TaskState taskState : this.taskStates.values()) { Properties newSpecProps = new Properties(); newSpecProps.putAll( Maps.difference(immutableCommonProperties, taskState.getWorkunit().getProperties()).entriesOnlyOnRight()); taskState.setWuProperties(immutableCommonProperties, newSpecProps); } }
@Override public boolean apply(WorkUnit workUnit) { if (workUnit instanceof MultiWorkUnit) { Preconditions.checkArgument(!workUnit.contains(ConfigurationKeys.WORK_UNIT_SKIP_KEY), "Error: MultiWorkUnit cannot be skipped"); for (WorkUnit wu : ((MultiWorkUnit) workUnit).getWorkUnits()) { Preconditions.checkArgument(!wu.contains(ConfigurationKeys.WORK_UNIT_SKIP_KEY), "Error: MultiWorkUnit cannot contain skipped WorkUnit"); } } if (workUnit.getPropAsBoolean(ConfigurationKeys.WORK_UNIT_SKIP_KEY, false)) { WorkUnitState workUnitState = new WorkUnitState(workUnit, this.jobState); workUnitState.setWorkingState(WorkUnitState.WorkingState.SKIPPED); this.jobState.addSkippedTaskState(new TaskState(workUnitState)); return false; } return true; } }
@BeforeClass public void setUp() { WorkUnitState workUnitState = new WorkUnitState(); workUnitState.setProp(ConfigurationKeys.JOB_ID_KEY, "Job-1"); workUnitState.setProp(ConfigurationKeys.TASK_ID_KEY, "Task-1"); this.taskState = new TaskState(workUnitState); }
@Test public void testPutIntoTaskStateStore() throws IOException { TaskState taskState1 = new TaskState(); taskState1.setJobId(JOB_ID); taskState1.setTaskId(TASK_ID_0); this.taskStateStore.put(JOB_ID, TASK_ID_0 + AbstractJobLauncher.TASK_STATE_STORE_TABLE_SUFFIX, taskState1); TaskState taskState2 = new TaskState(); taskState2.setJobId(JOB_ID); taskState2.setTaskId(TASK_ID_1); this.taskStateStore.put(JOB_ID, TASK_ID_1 + AbstractJobLauncher.TASK_STATE_STORE_TABLE_SUFFIX, taskState2); }
TaskState getEmptyTestTaskState(String taskId) { // Create a TaskState WorkUnit workUnit = WorkUnit.create( new Extract(Extract.TableType.SNAPSHOT_ONLY, this.getClass().getName(), this.getClass().getSimpleName())); workUnit.setProp(ConfigurationKeys.TASK_KEY_KEY, "taskKey"); TaskState taskState = new TaskState(new WorkUnitState(workUnit)); taskState.setProp(ConfigurationKeys.METRICS_ENABLED_KEY, Boolean.toString(false)); taskState.setTaskId(taskId); taskState.setJobId("1234"); return taskState; }
@Test(dependsOnMethods = {"testSetAndGet"}) public void testSerDe() throws IOException { Closer closer = Closer.create(); try { ByteArrayOutputStream baos = closer.register(new ByteArrayOutputStream()); DataOutputStream dos = closer.register(new DataOutputStream(baos)); this.taskState.write(dos); ByteArrayInputStream bais = closer.register((new ByteArrayInputStream(baos.toByteArray()))); DataInputStream dis = closer.register((new DataInputStream(bais))); TaskState newTaskState = new TaskState(); newTaskState.readFields(dis); Assert.assertEquals(newTaskState.getJobId(), "Job-1"); Assert.assertEquals(newTaskState.getTaskId(), "Task-1"); Assert.assertEquals(this.taskState.getHighWaterMark(), 2000); Assert.assertEquals(newTaskState.getStartTime(), this.startTime); Assert.assertEquals(newTaskState.getEndTime(), this.startTime + 1000); Assert.assertEquals(newTaskState.getTaskDuration(), 1000); Assert.assertEquals(newTaskState.getWorkingState(), WorkUnitState.WorkingState.COMMITTED); Assert.assertEquals(newTaskState.getProp("foo"), "bar"); } catch (Throwable t) { throw closer.rethrow(t); } finally { closer.close(); } }
TaskState getEmptyTestTaskState(String taskId) { // Create a TaskState WorkUnit workUnit = WorkUnit.create( new Extract(Extract.TableType.SNAPSHOT_ONLY, this.getClass().getName(), this.getClass().getSimpleName())); workUnit.setProp(ConfigurationKeys.TASK_KEY_KEY, "taskKey"); TaskState taskState = new TaskState(new WorkUnitState(workUnit)); taskState.setProp(ConfigurationKeys.METRICS_ENABLED_KEY, Boolean.toString(false)); taskState.setTaskId(taskId); taskState.setJobId("1234"); return taskState; }
private TaskState getStreamingTaskState() { WorkUnitState workUnitState = new WorkUnitState(WorkUnit.create( new Extract(Extract.TableType.SNAPSHOT_ONLY, this.getClass().getName(), this.getClass().getSimpleName()))); workUnitState.setProp(ConfigurationKeys.TASK_KEY_KEY, "1234"); TaskState taskState = new TaskState(workUnitState); taskState.setProp(ConfigurationKeys.METRICS_ENABLED_KEY, Boolean.toString(false)); taskState.setProp(TaskConfigurationKeys.TASK_EXECUTION_MODE, ExecutionModel.STREAMING.name()); taskState.setJobId("1234"); taskState.setTaskId("testContinuousTaskId"); return taskState; }
@Test public void testProcessedCount() throws Exception { JobContext jobContext = Mockito.mock(JobContext.class); JobState jobState = new JobState("jobName", "1234"); TaskState taskState1 = new TaskState(); TaskState taskState2 = new TaskState(); taskState1.setTaskId("1"); taskState1.setProp(ConfigurationKeys.WRITER_RECORDS_WRITTEN, "1"); taskState2.setTaskId("2"); taskState2.setProp(ConfigurationKeys.WRITER_RECORDS_WRITTEN, "22"); jobState.addTaskState(taskState1); jobState.addTaskState(taskState2); Mockito.when(jobContext.getJobState()).thenReturn(jobState); ClusterEventMetadataGenerator metadataGenerator = new ClusterEventMetadataGenerator(); Map<String, String> metadata; // processed count is not in job cancel event metadata = metadataGenerator.getMetadata(jobContext, EventName.JOB_CANCEL); Assert.assertEquals(metadata.get("processedCount"), null); // processed count is in job complete event metadata = metadataGenerator.getMetadata(jobContext, EventName.getEnumFromEventId("JobCompleteTimer")); Assert.assertEquals(metadata.get("processedCount"), "23"); }
@Test public void testPersistWatermarkStateToZk() throws IOException { CheckpointableWatermark watermark = new DefaultCheckpointableWatermark("source", new LongWatermark(startTime)); TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setProp(ConfigurationKeys.JOB_NAME_KEY, "JobName-" + startTime); // watermark storage configuration taskState.setProp(StateStoreBasedWatermarkStorage.WATERMARK_STORAGE_TYPE_KEY, "zk"); taskState.setProp(StateStoreBasedWatermarkStorage.WATERMARK_STORAGE_CONFIG_PREFIX + ZkStateStoreConfigurationKeys.STATE_STORE_ZK_CONNECT_STRING_KEY, testingServer.getConnectString()); StateStoreBasedWatermarkStorage watermarkStorage = new StateStoreBasedWatermarkStorage(taskState); watermarkStorage.commitWatermarks(ImmutableList.of(watermark)); Map<String, CheckpointableWatermark> watermarkMap = watermarkStorage.getCommittedWatermarks(DefaultCheckpointableWatermark.class, ImmutableList.of("source")); Assert.assertEquals(watermarkMap.size(), 1); Assert.assertEquals(((LongWatermark) watermarkMap.get("source").getWatermark()).getValue(), startTime); }
@Test public void testTaskGetMetrics() { String jobId = "job_456"; String taskId = "task_456"; String jobName = "jobName"; JobState jobState = new JobState(jobName, jobId); JobMetrics jobMetrics = JobMetrics.get(jobState); State props = new State(); props.setProp(ConfigurationKeys.JOB_ID_KEY, jobId); props.setProp(ConfigurationKeys.TASK_ID_KEY, taskId); SourceState sourceState = new SourceState(props, new ArrayList<WorkUnitState>()); WorkUnit workUnit = new WorkUnit(sourceState, null); WorkUnitState workUnitState = new WorkUnitState(workUnit); TaskState taskState = new TaskState(workUnitState); TaskMetrics taskMetrics = new TaskMetrics(taskState); Assert.assertNotNull(taskMetrics.getMetricContext()); Assert.assertTrue(taskMetrics.getMetricContext().getParent().isPresent()); Assert.assertEquals(taskMetrics.getMetricContext().getParent().get(), jobMetrics.getMetricContext()); }
public void persistDatasetState(List<Dataset> datasets, List<LongWatermark> watermarks, String jobName) throws IOException { Preconditions.checkArgument(datasets.size() >= 2); for (int i = 0; i < datasets.size(); i++) { String datasetUrn = datasets.get(i).getUrn(); JobState.DatasetState datasetState = new JobState.DatasetState(jobName, TEST_JOB_ID); datasetState.setDatasetUrn(datasetUrn); datasetState.setState(JobState.RunningState.COMMITTED); datasetState.setId(datasetUrn); datasetState.setStartTime(this.startTime); datasetState.setEndTime(this.startTime + 1000); datasetState.setDuration(1000); TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i); taskState.setId(TEST_TASK_ID_PREFIX + i); taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED); if (i < datasets.size() - 1) { taskState.setActualHighWatermark(watermarks.get(i)); } datasetState.addTaskState(taskState); this.fsDatasetStateStore.persistDatasetState(datasetUrn, datasetState); } }
@Test(dependsOnMethods = "testGetJobState") public void testPersistDatasetState() throws IOException { JobState.DatasetState datasetState = new JobState.DatasetState(TEST_JOB_NAME, TEST_JOB_ID); datasetState.setDatasetUrn(TEST_DATASET_URN); datasetState.setState(JobState.RunningState.COMMITTED); datasetState.setId(TEST_DATASET_URN); datasetState.setStartTime(this.startTime); datasetState.setEndTime(this.startTime + 1000); datasetState.setDuration(1000); for (int i = 0; i < 3; i++) { TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i); taskState.setId(TEST_TASK_ID_PREFIX + i); taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED); datasetState.addTaskState(taskState); } this.fsDatasetStateStore.persistDatasetState(TEST_DATASET_URN, datasetState); }
@Test public void testPersistJobState() throws IOException { JobState jobState = new JobState(TEST_JOB_NAME, TEST_JOB_ID); jobState.setId(TEST_JOB_ID); jobState.setProp("foo", "bar"); jobState.setState(JobState.RunningState.COMMITTED); jobState.setStartTime(this.startTime); jobState.setEndTime(this.startTime + 1000); jobState.setDuration(1000); for (int i = 0; i < 3; i++) { TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i); taskState.setId(TEST_TASK_ID_PREFIX + i); taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED); jobState.addTaskState(taskState); } this.fsJobStateStore.put(TEST_JOB_NAME, FsDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + FsDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, jobState); }
@Test public void testPersistJobState() throws IOException { JobState jobState = new JobState(TEST_JOB_NAME, TEST_JOB_ID); jobState.setId(TEST_JOB_ID); jobState.setProp("foo", "bar"); jobState.setState(JobState.RunningState.COMMITTED); jobState.setStartTime(this.startTime); jobState.setEndTime(this.startTime + 1000); jobState.setDuration(1000); for (int i = 0; i < 3; i++) { TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i); taskState.setId(TEST_TASK_ID_PREFIX + i); taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED); jobState.addTaskState(taskState); } zkJobStateStore.put(TEST_JOB_NAME, ZkDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + ZkDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, jobState); // second job name for testing getting store names in a later test case jobState.setJobName(TEST_JOB_NAME2); zkJobStateStore.put(TEST_JOB_NAME2, ZkDatasetStateStore.CURRENT_DATASET_STATE_FILE_SUFFIX + ZkDatasetStateStore.DATASET_STATE_STORE_TABLE_SUFFIX, jobState); }
TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i);
@Test(dependsOnMethods = "testGetJobState") public void testPersistDatasetState() throws IOException { JobState.DatasetState datasetState = new JobState.DatasetState(TEST_JOB_NAME, TEST_JOB_ID); datasetState.setDatasetUrn(TEST_DATASET_URN); datasetState.setState(JobState.RunningState.COMMITTED); datasetState.setId(TEST_DATASET_URN); datasetState.setStartTime(this.startTime); datasetState.setEndTime(this.startTime + 1000); datasetState.setDuration(1000); for (int i = 0; i < 3; i++) { TaskState taskState = new TaskState(); taskState.setJobId(TEST_JOB_ID); taskState.setTaskId(TEST_TASK_ID_PREFIX + i); taskState.setId(TEST_TASK_ID_PREFIX + i); taskState.setWorkingState(WorkUnitState.WorkingState.COMMITTED); datasetState.addTaskState(taskState); } zkDatasetStateStore.persistDatasetState(TEST_DATASET_URN, datasetState); // persist a second dataset state to test that retrieval of multiple dataset states works datasetState.setDatasetUrn(TEST_DATASET_URN2); datasetState.setId(TEST_DATASET_URN2); datasetState.setDuration(2000); zkDatasetStateStore.persistDatasetState(TEST_DATASET_URN2, datasetState); // second job name for testing getting store names in a later test case datasetState.setJobName(TEST_JOB_NAME2); zkDatasetStateStore.persistDatasetState(TEST_DATASET_URN2, datasetState); }