@Override
public void start(final IJobUpdate update, final AuditData auditData)
throws UpdateStateException {
requireNonNull(update);
requireNonNull(auditData);
storage.write(
(NoResult<UpdateStateException>)
storeProvider -> {
IJobUpdateSummary summary = update.getSummary();
IJobUpdateInstructions instructions = update.getInstructions();
IJobKey job = summary.getKey().getJob();
// Validate the update configuration by making sure we can create an updater for it.
updateFactory.newUpdate(update.getInstructions(), true);
if (instructions.getInitialState().isEmpty() && !instructions.isSetDesiredState()) {
throw new IllegalArgumentException("Update instruction is a no-op.");
}
List<IJobUpdateSummary> activeJobUpdates =
storeProvider.getJobUpdateStore().fetchJobUpdateSummaries(queryActiveByJob(job));
if (!activeJobUpdates.isEmpty()) {
throw new UpdateStateException(
"An active update already exists for this job, "
+ "please terminate it before starting another. "
+ "Active updates are those in states "
+ Updates.ACTIVE_JOB_UPDATE_STATES);
}
LOG.info("Starting update for job " + job);
ILock lock;
try {
lock =
lockManager.acquireLock(
ILockKey.build(LockKey.job(job.newBuilder())), auditData.getUser());
} catch (LockException e) {
throw new UpdateStateException(e.getMessage(), e);
}
storeProvider
.getJobUpdateStore()
.saveJobUpdate(update, Optional.of(requireNonNull(lock.getToken())));
JobUpdateStatus status = ROLLING_FORWARD;
if (isCoordinatedUpdate(instructions)) {
status = ROLL_FORWARD_AWAITING_PULSE;
pulseHandler.initializePulseState(update, status);
}
recordAndChangeJobUpdateStatus(
storeProvider, summary.getKey(), addAuditData(newEvent(status), auditData));
});
}
@Override
public ILock acquireLock(final ILockKey lockKey, final String user) throws LockException {
return storage.write(
storeProvider -> {
LockStore.Mutable lockStore = storeProvider.getLockStore();
Optional<ILock> existingLock = lockStore.fetchLock(lockKey);
if (existingLock.isPresent()) {
throw new LockException(
String.format(
"Operation for: %s is already in progress. Started at: %s. Current owner: %s.",
formatLockKey(lockKey),
new Date(existingLock.get().getTimestampMs()).toString(),
existingLock.get().getUser()));
}
ILock lock =
ILock.build(
new Lock()
.setKey(lockKey.newBuilder())
.setToken(tokenGenerator.createNew().toString())
.setTimestampMs(clock.nowMillis())
.setUser(user));
lockStore.saveLock(lock);
return lock;
});
}
@Override
public void releaseLock(final ILock lock) {
storage.write(
(NoResult.Quiet) storeProvider -> storeProvider.getLockStore().removeLock(lock.getKey()));
}
private void changeJobUpdateStatus(
MutableStoreProvider storeProvider,
IJobUpdateKey key,
JobUpdateEvent proposedEvent,
boolean recordChange)
throws UpdateStateException {
JobUpdateStatus status;
boolean record;
JobUpdateStore.Mutable updateStore = storeProvider.getJobUpdateStore();
Optional<String> updateLock = updateStore.getLockToken(key);
if (updateLock.isPresent()) {
status = proposedEvent.getStatus();
record = recordChange;
} else {
LOG.severe("Update " + key + " does not have a lock");
status = ERROR;
record = true;
}
LOG.info(String.format("Update %s is now in state %s", key, status));
if (record) {
updateStore.saveJobUpdateEvent(
key,
IJobUpdateEvent.build(proposedEvent.setTimestampMs(clock.nowMillis()).setStatus(status)));
}
if (TERMINAL_STATES.contains(status)) {
if (updateLock.isPresent()) {
lockManager.releaseLock(
ILock.build(
new Lock()
.setKey(LockKey.job(key.getJob().newBuilder()))
.setToken(updateLock.get())));
}
pulseHandler.remove(key);
} else {
pulseHandler.updatePulseStatus(key, status);
}
MonitorAction action = JobUpdateStateMachine.getActionForStatus(status);
IJobKey job = key.getJob();
if (action == STOP_WATCHING) {
updates.remove(job);
} else if (action == ROLL_FORWARD || action == ROLL_BACK) {
if (action == ROLL_BACK) {
updates.remove(job);
} else {
checkState(!updates.containsKey(job), "Updater already exists for " + job);
}
IJobUpdate jobUpdate = updateStore.fetchJobUpdate(key).get();
UpdateFactory.Update update;
try {
update = updateFactory.newUpdate(jobUpdate.getInstructions(), action == ROLL_FORWARD);
} catch (RuntimeException e) {
LOG.log(Level.WARNING, "Uncaught exception: " + e, e);
changeJobUpdateStatus(
storeProvider,
key,
newEvent(ERROR).setMessage("Internal scheduler error: " + e.getMessage()),
true);
return;
}
updates.put(job, update);
evaluateUpdater(storeProvider, update, jobUpdate.getSummary(), ImmutableMap.of());
}
}
final class Fixtures {
static final String ROLE = "bar_role";
static final String USER = "******";
static final Identity ROLE_IDENTITY = new Identity(ROLE, USER);
static final String JOB_NAME = "job_foo";
static final IJobKey JOB_KEY = JobKeys.from(ROLE, "devel", JOB_NAME);
static final ILockKey LOCK_KEY = ILockKey.build(LockKey.job(JOB_KEY.newBuilder()));
static final ILock LOCK = ILock.build(new Lock().setKey(LOCK_KEY.newBuilder()).setToken("token"));
static final JobConfiguration CRON_JOB = makeJob().setCronSchedule("* * * * *");
static final String TASK_ID = "task_id";
static final String UPDATE_ID = "82d6d790-3212-11e3-aa6e-0800200c9a74";
static final IJobUpdateKey UPDATE_KEY =
IJobUpdateKey.build(new JobUpdateKey(JOB_KEY.newBuilder(), UPDATE_ID));
static final UUID UU_ID = UUID.fromString(UPDATE_ID);
private static final Function<String, ResponseDetail> MESSAGE_TO_DETAIL =
new Function<String, ResponseDetail>() {
@Override
public ResponseDetail apply(String message) {
return new ResponseDetail().setMessage(message);
}
};
static final String CRON_SCHEDULE = "0 * * * *";
static final IResourceAggregate QUOTA =
IResourceAggregate.build(new ResourceAggregate(10.0, 1024, 2048));
static final QuotaCheckResult ENOUGH_QUOTA = new QuotaCheckResult(SUFFICIENT_QUOTA);
static final QuotaCheckResult NOT_ENOUGH_QUOTA = new QuotaCheckResult(INSUFFICIENT_QUOTA);
private Fixtures() {
// Utility class.
}
static JobConfiguration makeJob() {
return makeJob(nonProductionTask(), 1);
}
static JobConfiguration makeJob(TaskConfig task, int shardCount) {
return new JobConfiguration()
.setOwner(ROLE_IDENTITY)
.setInstanceCount(shardCount)
.setTaskConfig(task)
.setKey(JOB_KEY.newBuilder());
}
static TaskConfig defaultTask(boolean production) {
return new TaskConfig()
.setJob(JOB_KEY.newBuilder())
.setOwner(new Identity(ROLE, USER))
.setEnvironment("devel")
.setJobName(JOB_NAME)
.setContactEmail("*****@*****.**")
.setExecutorConfig(new ExecutorConfig("aurora", "data"))
.setNumCpus(1)
.setRamMb(1024)
.setDiskMb(1024)
.setProduction(production)
.setRequestedPorts(ImmutableSet.of())
.setTaskLinks(ImmutableMap.of())
.setMaxTaskFailures(1)
.setContainer(Container.mesos(new MesosContainer()));
}
static TaskConfig nonProductionTask() {
return defaultTask(false);
}
static Response jobSummaryResponse(Set<JobSummary> jobSummaries) {
return okResponse(Result.jobSummaryResult(new JobSummaryResult().setSummaries(jobSummaries)));
}
static Response response(ResponseCode code, Optional<Result> result, String... messages) {
Response response = Responses.empty().setResponseCode(code).setResult(result.orNull());
if (messages.length > 0) {
response.setDetails(
FluentIterable.from(Arrays.asList(messages)).transform(MESSAGE_TO_DETAIL).toList());
}
return response;
}
static Response okResponse(Result result) {
return response(OK, Optional.of(IResult.build(result).newBuilder()));
}
static JobConfiguration makeProdJob() {
return makeJob(productionTask(), 1);
}
static TaskConfig productionTask() {
return defaultTask(true);
}
static JobConfiguration makeJob(TaskConfig task) {
return makeJob(task, 1);
}
static Iterable<IScheduledTask> makeDefaultScheduledTasks(int n) {
return makeDefaultScheduledTasks(n, defaultTask(true));
}
static Iterable<IScheduledTask> makeDefaultScheduledTasks(int n, TaskConfig config) {
List<IScheduledTask> tasks = Lists.newArrayList();
for (int i = 0; i < n; i++) {
tasks.add(
IScheduledTask.build(
new ScheduledTask()
.setAssignedTask(new AssignedTask().setTask(config).setInstanceId(i))));
}
return tasks;
}
static Response assertOkResponse(Response response) {
return assertResponse(OK, response);
}
static Response assertResponse(ResponseCode expected, Response response) {
assertEquals(expected, response.getResponseCode());
return response;
}
}
@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);
}