@Override
public void restartShards(IJobKey jobKey, final Set<Integer> shards, final String requestingUser)
throws ScheduleException {
if (!JobKeys.isValid(jobKey)) {
throw new ScheduleException("Invalid job key: " + jobKey);
}
if (shards.isEmpty()) {
throw new ScheduleException("At least one shard must be specified.");
}
final Query.Builder query = Query.instanceScoped(jobKey, shards).active();
storage.write(
new MutateWork.NoResult<ScheduleException>() {
@Override
protected void execute(MutableStoreProvider storeProvider) throws ScheduleException {
Set<IScheduledTask> matchingTasks = storeProvider.getTaskStore().fetchTasks(query);
if (matchingTasks.size() != shards.size()) {
throw new ScheduleException("Not all requested shards are active.");
}
LOG.info("Restarting shards matching " + query);
for (String taskId : Tasks.ids(matchingTasks)) {
stateManager.changeState(
taskId,
Optional.<ScheduleStatus>absent(),
RESTARTING,
Optional.of("Restarted by " + requestingUser));
}
}
});
}
private static Optional<IScheduledTask> getActiveInstance(
TaskStore taskStore, IJobKey job, int instanceId) {
return Optional.fromNullable(
Iterables.getOnlyElement(
taskStore.fetchTasks(Query.instanceScoped(job, instanceId).active()), null));
}
@Test
public void testGetJobUpdateDiffWithUpdateRemove() throws Exception {
TaskConfig task1 = defaultTask(false).setNumCpus(1.0);
TaskConfig task2 = defaultTask(false).setNumCpus(2.0);
TaskConfig task3 = defaultTask(false).setNumCpus(3.0);
ImmutableSet.Builder<IScheduledTask> tasks = ImmutableSet.builder();
makeTasks(0, 10, task1, tasks);
makeTasks(10, 20, task2, tasks);
makeTasks(20, 30, task3, tasks);
expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
storageUtil.expectTaskFetch(Query.jobScoped(JOB_KEY).active(), tasks.build());
control.replay();
JobUpdateRequest request =
new JobUpdateRequest()
.setTaskConfig(defaultTask(false).setNumCpus(6.0))
.setInstanceCount(20)
.setSettings(new JobUpdateSettings());
GetJobUpdateDiffResult expected =
new GetJobUpdateDiffResult()
.setRemove(ImmutableSet.of(group(task3, new Range(20, 29))))
.setUpdate(
ImmutableSet.of(group(task1, new Range(0, 9)), group(task2, new Range(10, 19))))
.setAdd(ImmutableSet.of())
.setUnchanged(ImmutableSet.of());
Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
}
@Test
public void testGetPendingReasonFailsStatusSet() throws Exception {
Builder query = Query.unscoped().byStatus(ScheduleStatus.ASSIGNED);
control.replay();
assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
}
@Test
public void testGetPendingReasonFailsSlavesSet() throws Exception {
Builder query = Query.unscoped().bySlave("host1");
control.replay();
assertResponse(INVALID_REQUEST, thrift.getPendingReason(query.get()));
}
@Test
public void testGetTasksStatus() throws Exception {
Builder query = Query.unscoped();
Iterable<IScheduledTask> tasks = makeDefaultScheduledTasks(10);
storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(tasks));
control.replay();
ImmutableList<ScheduledTask> expected = IScheduledTask.toBuildersList(tasks);
Response response = assertOkResponse(thrift.getTasksStatus(new TaskQuery()));
assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
}
/**
* Computes totals for each of the {@link MetricType}s.
*
* @return aggregates for each global metric type.
* @throws StorageException if there was a problem fetching tasks from storage.
*/
public List<GlobalMetric> computeConsumptionTotals() throws StorageException {
List<GlobalMetric> counts =
FluentIterable.from(Arrays.asList(MetricType.values()))
.transform(TO_GLOBAL_METRIC)
.toList();
for (ITaskConfig task : getTasks(Query.unscoped().active())) {
for (GlobalMetric count : counts) {
count.accumulate(task);
}
}
return counts;
}
@Override
public synchronized void killTasks(Query.Builder query, final String user)
throws ScheduleException {
checkNotNull(query);
LOG.info("Killing tasks matching " + query);
boolean jobDeleted = false;
if (Query.isOnlyJobScoped(query)) {
// If this looks like a query for all tasks in a job, instruct the scheduler modules to
// delete the job.
IJobKey jobKey = JobKeys.from(query).get();
for (JobManager manager : jobManagers) {
if (manager.deleteJob(jobKey)) {
jobDeleted = true;
}
}
}
// Unless statuses were specifically supplied, only attempt to kill active tasks.
final Query.Builder taskQuery =
query.get().isSetStatuses() ? query.byStatus(ACTIVE_STATES) : query;
int tasksAffected =
storage.write(
new MutateWork.Quiet<Integer>() {
@Override
public Integer apply(MutableStoreProvider storeProvider) {
int total = 0;
for (String taskId :
Tasks.ids(storeProvider.getTaskStore().fetchTasks(taskQuery))) {
boolean changed =
stateManager.changeState(
taskId,
Optional.<ScheduleStatus>absent(),
KILLING,
Optional.of("Killed by " + user));
if (changed) {
total++;
}
}
return total;
}
});
if (!jobDeleted && (tasksAffected == 0)) {
throw new ScheduleException("No jobs to kill");
}
}
@Test
public void testGetPendingReason() throws Exception {
Builder query = Query.unscoped().byJob(JOB_KEY);
Builder filterQuery = Query.unscoped().byJob(JOB_KEY).byStatus(ScheduleStatus.PENDING);
String taskId1 = "task_id_test1";
String taskId2 = "task_id_test2";
ImmutableSet<Veto> result =
ImmutableSet.of(Veto.constraintMismatch("first"), Veto.constraintMismatch("second"));
ITaskConfig taskConfig = ITaskConfig.build(defaultTask(true));
IScheduledTask pendingTask1 =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(
new AssignedTask().setTaskId(taskId1).setTask(taskConfig.newBuilder()))
.setStatus(ScheduleStatus.PENDING));
IScheduledTask pendingTask2 =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(
new AssignedTask().setTaskId(taskId2).setTask(taskConfig.newBuilder()))
.setStatus(ScheduleStatus.PENDING));
storageUtil.expectTaskFetch(filterQuery, pendingTask1, pendingTask2);
expect(nearestFit.getNearestFit(TaskGroupKey.from(taskConfig))).andReturn(result).times(2);
control.replay();
String reason = "Constraint not satisfied: first,Constraint not satisfied: second";
Set<PendingReason> expected =
ImmutableSet.of(
new PendingReason().setTaskId(taskId1).setReason(reason),
new PendingReason().setTaskId(taskId2).setReason(reason));
Response response = assertOkResponse(thrift.getPendingReason(query.get()));
assertEquals(expected, response.getResult().getGetPendingReasonResult().getReasons());
}
@Test
public void testModifySnapshotBeforeCommit() throws Exception {
expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
Snapshot modified = SNAPSHOT1.deepCopy().setTasks(ImmutableSet.of(TASK1));
Capture<MutateWork<Object, Exception>> transaction = createCapture();
expect(primaryStorage.write(capture(transaction))).andReturn(null);
distributedStore.persist(modified);
shutDownNow.execute();
control.replay();
clock.advance(INTERVAL);
storageBackup.createSnapshot();
String backup1 = storageBackup.createBackupName();
recovery.stage(backup1);
assertEquals(
IScheduledTask.setFromBuilders(SNAPSHOT1.getTasks()), recovery.query(Query.unscoped()));
recovery.deleteTasks(Query.taskScoped(Tasks.id(TASK2)));
assertEquals(
IScheduledTask.setFromBuilders(modified.getTasks()), recovery.query(Query.unscoped()));
recovery.commit();
transaction.getValue().apply(storeProvider);
}
@Test
public void testGetTasksWithoutConfigs() throws Exception {
Builder query = Query.unscoped();
storageUtil.expectTaskFetch(query, ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));
control.replay();
ImmutableList<ScheduledTask> expected =
IScheduledTask.toBuildersList(
makeDefaultScheduledTasks(10, defaultTask(true).setExecutorConfig(null)));
Response response = assertOkResponse(thrift.getTasksWithoutConfigs(new TaskQuery()));
assertEquals(expected, response.getResult().getScheduleStatusResult().getTasks());
}
@Test
public void testKillExistingCollisionFailedKill() throws Exception {
IJobConfiguration killExisting =
IJobConfiguration.build(
job.newBuilder().setCronCollisionPolicy(CronCollisionPolicy.KILL_EXISTING));
Capture<Runnable> jobTriggerCapture = expectJobAccepted(killExisting);
expectActiveTaskFetch(TASK);
scheduler.killTasks(Query.jobScoped(killExisting.getKey()).active(), CRON_USER);
expectLastCall().andThrow(new ScheduleException("injected"));
control.replay();
cron.receiveJob(new SanitizedConfiguration(killExisting));
jobTriggerCapture.getValue().run();
}
@Test
public void testEmptyConfigSummary() throws Exception {
IJobKey key = JobKeys.from("test", "test", "test");
storageUtil.expectTaskFetch(Query.jobScoped(key).active(), ImmutableSet.of());
ConfigSummary summary =
new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet());
ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);
control.replay();
Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
assertEquals(expected, response.getResult().getConfigSummaryResult());
}
@Test
public void testGetConfigSummary() throws Exception {
IJobKey key = JobKeys.from("test", "test", "test");
TaskConfig firstGroupTask = defaultTask(true);
TaskConfig secondGroupTask = defaultTask(true).setNumCpus(2);
IScheduledTask first1 =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(0)));
IScheduledTask first2 =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(firstGroupTask).setInstanceId(1)));
IScheduledTask second =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(secondGroupTask).setInstanceId(2)));
storageUtil.expectTaskFetch(Query.jobScoped(key).active(), first1, first2, second);
ConfigGroup group1 =
new ConfigGroup()
.setConfig(firstGroupTask)
.setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(0, 1)))));
ConfigGroup group2 =
new ConfigGroup()
.setConfig(secondGroupTask)
.setInstances(IRange.toBuildersSet(convertRanges(toRanges(ImmutableSet.of(2)))));
ConfigSummary summary =
new ConfigSummary().setKey(key.newBuilder()).setGroups(Sets.newHashSet(group1, group2));
ConfigSummaryResult expected = new ConfigSummaryResult().setSummary(summary);
control.replay();
Response response = assertOkResponse(thrift.getConfigSummary(key.newBuilder()));
assertEquals(
IConfigSummaryResult.build(expected),
IConfigSummaryResult.build(response.getResult().getConfigSummaryResult()));
}
/**
* Displays the aggregate utilization for roles within a metric type.
*
* @param metric Metric id.
* @return HTML-formatted utilization within the metric type.
*/
@GET
@Path("/{metric}")
@Produces(MediaType.TEXT_HTML)
public Response aggregateRoles(@PathParam("metric") final String metric) {
final MetricType type = getTypeByName(metric);
Function<ITaskConfig, Display> toKey =
new Function<ITaskConfig, Display>() {
@Override
public Display apply(ITaskConfig task) {
String role = task.getJob().getRole();
return new Display(role, metric + "/" + role);
}
};
Map<Display, Metric> byRole =
counter.computeAggregates(Query.unscoped().active(), type.filter, toKey);
return Response.ok(fillTemplate(byRole)).build();
}
/**
* Displays the aggregate utilization for jobs within a role.
*
* @param metric Metric id.
* @param role Role for jobs to aggregate.
* @return HTML-formatted utilization within the metric/role.
*/
@GET
@Path("/{metric}/{role}")
@Produces(MediaType.TEXT_HTML)
public Response aggregateJobs(
@PathParam("metric") String metric, @PathParam("role") String role) {
MetricType type = getTypeByName(metric);
Function<ITaskConfig, Display> toKey =
new Function<ITaskConfig, Display>() {
@Override
public Display apply(ITaskConfig task) {
return new Display(task.getJobName(), null);
}
};
Map<Display, Metric> byJob =
counter.computeAggregates(Query.roleScoped(role).active(), type.filter, toKey);
return Response.ok(fillTemplate(byJob)).build();
}
@Test
public void testExecution() {
expect(statsProvider.makeCounter(EXPLICIT_STAT_NAME)).andReturn(explicitRuns);
expect(statsProvider.makeCounter(IMPLICIT_STAT_NAME)).andReturn(implicitRuns);
clock = FakeScheduledExecutor.scheduleAtFixedRateExecutor(executorService, 2, 5);
IScheduledTask task = makeTask("id1", TaskTestUtil.makeConfig(TaskTestUtil.JOB));
storageUtil.expectOperations();
storageUtil
.expectTaskFetch(Query.unscoped().byStatus(Tasks.SLAVE_ASSIGNED_STATES), task)
.times(5);
driver.reconcileTasks(ImmutableSet.of(TASK_TO_PROTO.apply(task)));
expectLastCall().times(5);
driver.reconcileTasks(ImmutableSet.of());
expectLastCall().times(2);
control.replay();
TaskReconciler reconciler =
new TaskReconciler(SETTINGS, storageUtil.storage, driver, executorService, statsProvider);
reconciler.startAsync().awaitRunning();
clock.advance(INITIAL_DELAY);
assertEquals(1L, explicitRuns.get());
assertEquals(0L, implicitRuns.get());
clock.advance(SPREAD);
assertEquals(1L, explicitRuns.get());
assertEquals(1L, implicitRuns.get());
clock.advance(EXPLICIT_SCHEDULE);
assertEquals(2L, explicitRuns.get());
assertEquals(1L, implicitRuns.get());
clock.advance(IMPLICT_SCHEDULE);
assertEquals(5L, explicitRuns.get());
assertEquals(2L, implicitRuns.get());
}
@Test
public void testGetAllJobs() throws Exception {
JobConfiguration cronJobOne =
makeJob()
.setCronSchedule("1 * * * *")
.setKey(JOB_KEY.newBuilder())
.setTaskConfig(nonProductionTask());
JobKey jobKey2 = JOB_KEY.newBuilder().setRole("other_role");
JobConfiguration cronJobTwo =
makeJob().setCronSchedule("2 * * * *").setKey(jobKey2).setTaskConfig(nonProductionTask());
TaskConfig immediateTaskConfig =
defaultTask(false)
.setJob(JOB_KEY.newBuilder().setName("immediate"))
.setJobName("immediate")
.setOwner(ROLE_IDENTITY);
IScheduledTask immediateTask =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(immediateTaskConfig))
.setStatus(ScheduleStatus.ASSIGNED));
JobConfiguration immediateJob =
new JobConfiguration()
.setKey(JOB_KEY.newBuilder().setName("immediate"))
.setOwner(ROLE_IDENTITY)
.setInstanceCount(1)
.setTaskConfig(immediateTaskConfig);
Set<JobConfiguration> crons = ImmutableSet.of(cronJobOne, cronJobTwo);
expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(crons));
storageUtil.expectTaskFetch(Query.unscoped().active(), immediateTask);
control.replay();
Set<JobConfiguration> allJobs =
ImmutableSet.<JobConfiguration>builder().addAll(crons).add(immediateJob).build();
assertEquals(
IJobConfiguration.setFromBuilders(allJobs),
IJobConfiguration.setFromBuilders(
thrift.getJobs(null).getResult().getGetJobsResult().getConfigs()));
}
@Override
public Amount<Long, Time> getReevaluationDelay(
IInstanceKey instance,
IJobUpdateInstructions instructions,
MutableStoreProvider storeProvider,
StateManager stateManager,
JobUpdateStatus status) {
String taskId =
Tasks.id(
Iterables.getOnlyElement(
storeProvider
.getTaskStore()
.fetchTasks(Query.instanceScoped(instance).active())));
LOG.info("Killing " + instance + " while " + status);
stateManager.changeState(
storeProvider,
taskId,
Optional.absent(),
ScheduleStatus.KILLING,
Optional.of("Killed for job update."));
return Amount.of(
(long) instructions.getSettings().getMaxWaitToInstanceRunningMs(), Time.MILLISECONDS);
}
@Test
public void testGetJobSummaryWithoutNextRun() throws Exception {
// 31st of February, there is no such day.
String cronSchedule = "* * 31 2 *";
TaskConfig task =
nonProductionTask()
.setJobName(JOB_KEY.getName())
.setOwner(ROLE_IDENTITY)
.setEnvironment(JOB_KEY.getEnvironment());
JobConfiguration job = makeJob().setCronSchedule(cronSchedule).setTaskConfig(task);
expect(cronPredictor.predictNextRun(CrontabEntry.parse(cronSchedule)))
.andReturn(Optional.absent())
.anyTimes();
storageUtil.expectTaskFetch(Query.roleScoped(ROLE));
Set<JobConfiguration> jobOnly = ImmutableSet.of(job);
expect(storageUtil.jobStore.fetchJobs()).andReturn(IJobConfiguration.setFromBuilders(jobOnly));
control.replay();
JobSummaryResult result = thrift.getJobSummary(ROLE).getResult().getJobSummaryResult();
assertEquals(1, result.getSummaries().size());
assertFalse(result.getSummariesIterator().next().isSetNextCronRunMs());
}
@Test
public void testGetJobUpdateDiffWithUnchanged() throws Exception {
expect(storageUtil.jobStore.fetchJob(JOB_KEY)).andReturn(Optional.absent());
storageUtil.expectTaskFetch(
Query.jobScoped(JOB_KEY).active(), ImmutableSet.copyOf(makeDefaultScheduledTasks(10)));
control.replay();
JobUpdateRequest request =
new JobUpdateRequest()
.setTaskConfig(defaultTask(true))
.setInstanceCount(10)
.setSettings(new JobUpdateSettings());
GetJobUpdateDiffResult expected =
new GetJobUpdateDiffResult()
.setUnchanged(ImmutableSet.of(group(defaultTask(true), new Range(0, 9))))
.setRemove(ImmutableSet.of())
.setUpdate(ImmutableSet.of())
.setAdd(ImmutableSet.of());
Response response = assertOkResponse(thrift.getJobUpdateDiff(request));
assertEquals(expected, response.getResult().getGetJobUpdateDiffResult());
}
@VisibleForTesting
static Query.Builder jobHistoryQuery(IJobKey jobKey) {
return Query.jobScoped(jobKey).byStatus(apiConstants.TERMINAL_STATES);
}
@Test
public void testCreateAndRestoreNewSnapshot() {
ImmutableSet<IScheduledTask> tasks =
ImmutableSet.of(
IScheduledTask.build(new ScheduledTask().setStatus(ScheduleStatus.PENDING)));
Set<QuotaConfiguration> quotas =
ImmutableSet.of(new QuotaConfiguration("steve", ResourceAggregates.none().newBuilder()));
IHostAttributes attribute =
IHostAttributes.build(
new HostAttributes(
"host", ImmutableSet.of(new Attribute("attr", ImmutableSet.of("value")))));
StoredJob job =
new StoredJob(
"jobManager", new JobConfiguration().setKey(new JobKey("owner", "env", "name")));
String frameworkId = "framework_id";
ILock lock =
ILock.build(
new Lock()
.setKey(LockKey.job(JobKeys.from("testRole", "testEnv", "testJob").newBuilder()))
.setToken("lockId")
.setUser("testUser")
.setTimestampMs(12345L));
SchedulerMetadata metadata =
new SchedulerMetadata().setFrameworkId(frameworkId).setVersion(CURRENT_API_VERSION);
storageUtil.expectOperations();
expect(storageUtil.taskStore.fetchTasks(Query.unscoped())).andReturn(tasks);
expect(storageUtil.quotaStore.fetchQuotas())
.andReturn(ImmutableMap.of("steve", ResourceAggregates.none()));
expect(storageUtil.attributeStore.getHostAttributes()).andReturn(ImmutableSet.of(attribute));
expect(storageUtil.jobStore.fetchManagerIds()).andReturn(ImmutableSet.of("jobManager"));
expect(storageUtil.jobStore.fetchJobs("jobManager"))
.andReturn(ImmutableSet.of(IJobConfiguration.build(job.getJobConfiguration())));
expect(storageUtil.schedulerStore.fetchFrameworkId()).andReturn(frameworkId);
expect(storageUtil.lockStore.fetchLocks()).andReturn(ImmutableSet.of(lock));
expectDataWipe();
storageUtil.taskStore.saveTasks(tasks);
storageUtil.quotaStore.saveQuota("steve", ResourceAggregates.none());
storageUtil.attributeStore.saveHostAttributes(attribute);
storageUtil.jobStore.saveAcceptedJob(
job.getJobManagerId(), IJobConfiguration.build(job.getJobConfiguration()));
storageUtil.schedulerStore.saveFrameworkId(frameworkId);
storageUtil.lockStore.saveLock(lock);
control.replay();
Snapshot expected =
new Snapshot()
.setTimestamp(NOW)
.setTasks(IScheduledTask.toBuildersSet(tasks))
.setQuotaConfigurations(quotas)
.setHostAttributes(ImmutableSet.of(attribute.newBuilder()))
.setJobs(ImmutableSet.of(job))
.setSchedulerMetadata(metadata)
.setLocks(ILock.toBuildersSet(ImmutableSet.of(lock)));
assertEquals(expected, snapshotStore.createSnapshot());
snapshotStore.applySnapshot(expected);
}
private IExpectationSetters<?> expectActiveTaskFetch(IScheduledTask... activeTasks) {
return storageUtil.expectTaskFetch(Query.jobScoped(job.getKey()).active(), activeTasks);
}
@Test
public void testGetJobSummary() throws Exception {
long nextCronRunMs = 100;
TaskConfig ownedCronJobTask =
nonProductionTask()
.setJob(JOB_KEY.newBuilder())
.setJobName(JOB_KEY.getName())
.setOwner(ROLE_IDENTITY)
.setEnvironment(JOB_KEY.getEnvironment());
JobConfiguration ownedCronJob =
makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
IScheduledTask ownedCronJobScheduledTask =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
.setStatus(ScheduleStatus.ASSIGNED));
Identity otherOwner = new Identity("other", "other");
JobConfiguration unownedCronJob =
makeJob()
.setOwner(otherOwner)
.setCronSchedule(CRON_SCHEDULE)
.setKey(JOB_KEY.newBuilder().setRole("other"))
.setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
TaskConfig ownedImmediateTaskInfo =
defaultTask(false)
.setJob(JOB_KEY.newBuilder().setName("immediate"))
.setJobName("immediate")
.setOwner(ROLE_IDENTITY);
Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
Set<JobSummary> ownedCronJobSummaryOnly =
ImmutableSet.of(
new JobSummary()
.setJob(ownedCronJob)
.setStats(new JobStats())
.setNextCronRunMs(nextCronRunMs));
Set<JobSummary> ownedCronJobSummaryWithRunningTask =
ImmutableSet.of(
new JobSummary()
.setJob(ownedCronJob)
.setStats(new JobStats().setActiveTaskCount(1))
.setNextCronRunMs(nextCronRunMs));
Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);
IScheduledTask ownedImmediateTask =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
.setStatus(ScheduleStatus.ASSIGNED));
JobConfiguration ownedImmediateJob =
new JobConfiguration()
.setKey(JOB_KEY.newBuilder().setName("immediate"))
.setOwner(ROLE_IDENTITY)
.setInstanceCount(1)
.setTaskConfig(ownedImmediateTaskInfo);
Builder query = Query.roleScoped(ROLE);
Set<JobSummary> ownedImmediateJobSummaryOnly =
ImmutableSet.of(
new JobSummary()
.setJob(ownedImmediateJob)
.setStats(new JobStats().setActiveTaskCount(1)));
expect(cronPredictor.predictNextRun(CrontabEntry.parse(CRON_SCHEDULE)))
.andReturn(Optional.of(new Date(nextCronRunMs)))
.anyTimes();
storageUtil.expectTaskFetch(query);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));
storageUtil.expectTaskFetch(query);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));
storageUtil.expectTaskFetch(query, ownedImmediateTask);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));
storageUtil.expectTaskFetch(query);
expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());
// Handle the case where a cron job has a running task (same JobKey present in both stores).
storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));
control.replay();
assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));
assertEquals(jobSummaryResponse(ownedCronJobSummaryOnly), thrift.getJobSummary(ROLE));
Response jobSummaryResponse = thrift.getJobSummary(ROLE);
assertEquals(
jobSummaryResponse(ownedImmediateJobSummaryOnly),
IResponse.build(jobSummaryResponse).newBuilder());
assertEquals(jobSummaryResponse(ImmutableSet.of()), thrift.getJobSummary(ROLE));
assertEquals(
jobSummaryResponse(ownedCronJobSummaryWithRunningTask), thrift.getJobSummary(ROLE));
}
@Benchmark
public Iterable<IScheduledTask> run() {
return storage.read(store -> store.getTaskStore().fetchTasks(Query.unscoped()));
}
@Test
public void testGetJobs() throws Exception {
TaskConfig ownedCronJobTask =
nonProductionTask()
.setJobName(JOB_KEY.getName())
.setOwner(ROLE_IDENTITY)
.setEnvironment(JOB_KEY.getEnvironment());
JobConfiguration ownedCronJob =
makeJob().setCronSchedule(CRON_SCHEDULE).setTaskConfig(ownedCronJobTask);
IScheduledTask ownedCronJobScheduledTask =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(ownedCronJobTask))
.setStatus(ScheduleStatus.ASSIGNED));
Identity otherOwner = new Identity("other", "other");
JobConfiguration unownedCronJob =
makeJob()
.setOwner(otherOwner)
.setCronSchedule(CRON_SCHEDULE)
.setKey(JOB_KEY.newBuilder().setRole("other"))
.setTaskConfig(ownedCronJobTask.deepCopy().setOwner(otherOwner));
TaskConfig ownedImmediateTaskInfo =
defaultTask(false)
.setJob(JOB_KEY.newBuilder().setName("immediate"))
.setJobName("immediate")
.setOwner(ROLE_IDENTITY);
Set<JobConfiguration> ownedCronJobOnly = ImmutableSet.of(ownedCronJob);
Set<JobConfiguration> unownedCronJobOnly = ImmutableSet.of(unownedCronJob);
Set<JobConfiguration> bothCronJobs = ImmutableSet.of(ownedCronJob, unownedCronJob);
IScheduledTask ownedImmediateTask =
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(ownedImmediateTaskInfo))
.setStatus(ScheduleStatus.ASSIGNED));
JobConfiguration ownedImmediateJob =
new JobConfiguration()
.setKey(JOB_KEY.newBuilder().setName("immediate"))
.setOwner(ROLE_IDENTITY)
.setInstanceCount(1)
.setTaskConfig(ownedImmediateTaskInfo);
Query.Builder query = Query.roleScoped(ROLE).active();
storageUtil.expectTaskFetch(query);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(ownedCronJobOnly));
storageUtil.expectTaskFetch(query);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(bothCronJobs));
storageUtil.expectTaskFetch(query, ownedImmediateTask);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(unownedCronJobOnly));
expect(storageUtil.jobStore.fetchJobs()).andReturn(ImmutableSet.of());
storageUtil.expectTaskFetch(query);
// Handle the case where a cron job has a running task (same JobKey present in both stores).
storageUtil.expectTaskFetch(query, ownedCronJobScheduledTask);
expect(storageUtil.jobStore.fetchJobs())
.andReturn(IJobConfiguration.setFromBuilders(ImmutableSet.of(ownedCronJob)));
control.replay();
assertJobsEqual(
ownedCronJob,
Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));
assertJobsEqual(
ownedCronJob,
Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));
Set<JobConfiguration> queryResult3 =
thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs();
assertJobsEqual(ownedImmediateJob, Iterables.getOnlyElement(queryResult3));
assertEquals(
ITaskConfig.build(ownedImmediateTaskInfo),
ITaskConfig.build(Iterables.getOnlyElement(queryResult3).getTaskConfig()));
assertTrue(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs().isEmpty());
assertJobsEqual(
ownedCronJob,
Iterables.getOnlyElement(thrift.getJobs(ROLE).getResult().getGetJobsResult().getConfigs()));
}
private TaskQuery setupPaginatedQuery(Iterable<IScheduledTask> tasks, int offset, int limit) {
TaskQuery query = new TaskQuery().setOffset(offset).setLimit(limit);
Builder builder = Query.arbitrary(query);
storageUtil.expectTaskFetch(builder, ImmutableSet.copyOf(tasks));
return query;
}
