/**
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.aurora.scheduler.state;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import javax.annotation.Nullable;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Functions;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import org.apache.aurora.common.base.Command;
import org.apache.aurora.common.base.Consumers;
import org.apache.aurora.common.base.MorePreconditions;
import org.apache.aurora.common.stats.Stats;
import org.apache.aurora.common.util.StateMachine;
import org.apache.aurora.common.util.StateMachine.Rule;
import org.apache.aurora.common.util.StateMachine.Transition;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import static java.util.Objects.requireNonNull;
import static org.apache.aurora.scheduler.state.SideEffect.Action;
import static org.apache.aurora.scheduler.state.SideEffect.Action.DELETE;
import static org.apache.aurora.scheduler.state.SideEffect.Action.INCREMENT_FAILURES;
import static org.apache.aurora.scheduler.state.SideEffect.Action.KILL;
import static org.apache.aurora.scheduler.state.SideEffect.Action.RESCHEDULE;
import static org.apache.aurora.scheduler.state.SideEffect.Action.SAVE_STATE;
import static org.apache.aurora.scheduler.state.StateChangeResult.ILLEGAL;
import static org.apache.aurora.scheduler.state.StateChangeResult.ILLEGAL_WITH_SIDE_EFFECTS;
import static org.apache.aurora.scheduler.state.StateChangeResult.NOOP;
import static org.apache.aurora.scheduler.state.StateChangeResult.SUCCESS;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.ASSIGNED;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.DELETED;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.DRAINING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.FAILED;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.FINISHED;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.INIT;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.KILLED;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.KILLING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.LOST;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.PENDING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.PREEMPTING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.RESTARTING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.RUNNING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.STARTING;
import static org.apache.aurora.scheduler.state.TaskStateMachine.TaskState.THROTTLED;
/**
* State machine for a task.
* <p>
* This enforces the lifecycle of a task, and triggers the actions that should be taken in response
* to different state transitions. These responses are externally communicated by populating a
* provided work queue.
* <p>
* TODO(wfarner): Augment this class to force the one-time-use nature. This is probably best done
* by hiding the constructor and exposing only a static function to transition a task and get the
* resulting actions.
*/
class TaskStateMachine {
private static final Logger LOG = LoggerFactory.getLogger(TaskStateMachine.class);
private static final AtomicLong ILLEGAL_TRANSITIONS =
Stats.exportLong("scheduler_illegal_task_state_transitions");
private final StateMachine<TaskState> stateMachine;
private Optional<TaskState> previousState = Optional.absent();
private final Set<SideEffect> sideEffects = Sets.newHashSet();
private static final Function<ScheduleStatus, TaskState> STATUS_TO_TASK_STATE =
input -> TaskState.valueOf(input.name());
private static final Function<IScheduledTask, TaskState> SCHEDULED_TO_TASK_STATE =
Functions.compose(STATUS_TO_TASK_STATE, IScheduledTask::getStatus);
/**
* ScheduleStatus enum extension to account for cases where no direct state mapping exists.
* TODO:(maxim): Consider making this private.
*/
@VisibleForTesting
enum TaskState {
INIT(Optional.of(ScheduleStatus.INIT)),
THROTTLED(Optional.of(ScheduleStatus.THROTTLED)),
PENDING(Optional.of(ScheduleStatus.PENDING)),
ASSIGNED(Optional.of(ScheduleStatus.ASSIGNED)),
STARTING(Optional.of(ScheduleStatus.STARTING)),
RUNNING(Optional.of(ScheduleStatus.RUNNING)),
FINISHED(Optional.of(ScheduleStatus.FINISHED)),
PREEMPTING(Optional.of(ScheduleStatus.PREEMPTING)),
RESTARTING(Optional.of(ScheduleStatus.RESTARTING)),
DRAINING(Optional.of(ScheduleStatus.DRAINING)),
FAILED(Optional.of(ScheduleStatus.FAILED)),
KILLED(Optional.of(ScheduleStatus.KILLED)),
KILLING(Optional.of(ScheduleStatus.KILLING)),
LOST(Optional.of(ScheduleStatus.LOST)),
/**
* The task does not have an associated state as it has been deleted from the store.
*/
DELETED(Optional.<ScheduleStatus>absent());
private final Optional<ScheduleStatus> status;
TaskState(Optional<ScheduleStatus> status) {
this.status = status;
}
Optional<ScheduleStatus> getStatus() {
return status;
}
}
/**
* Creates a new task state machine representing a non-existent task. This allows for consistent
* state-reconciliation actions when the external system disagrees with the scheduler.
*
* @param name Name of the state machine, for logging.
*/
TaskStateMachine(String name) {
this(name, Optional.absent());
}
/**
* Creates a new task state machine representing an existent task. The state machine will be
* named with the tasks ID.
*.
* @param task Read-only task that this state machine manages.
*/
TaskStateMachine(IScheduledTask task) {
this(Tasks.id(task), Optional.of(task));
}
private TaskStateMachine(final String name, final Optional<IScheduledTask> task) {
MorePreconditions.checkNotBlank(name);
requireNonNull(task);
final TaskState initialState = task.transform(SCHEDULED_TO_TASK_STATE).or(DELETED);
if (task.isPresent()) {
Preconditions.checkState(
initialState != DELETED,
"A task that exists may not be in DELETED state.");
} else {
Preconditions.checkState(
initialState == DELETED,
"A task that does not exist must start in DELETED state.");
}
Consumer<Transition<TaskState>> manageTerminatedTasks = Consumers.combine(
ImmutableList.<Consumer<Transition<TaskState>>>builder()
// Kill a task that we believe to be terminated when an attempt is made to revive.
.add(
Consumers.filter(Transition.to(ASSIGNED, STARTING, RUNNING),
addFollowupClosure(KILL)))
// Remove a terminated task that is requested to be deleted.
.add(Consumers.filter(Transition.to(DELETED), addFollowupClosure(DELETE)))
.build());
final Consumer<Transition<TaskState>> manageRestartingTask =
transition -> {
switch (transition.getTo()) {
case ASSIGNED:
addFollowup(KILL);
break;
case STARTING:
addFollowup(KILL);
break;
case RUNNING:
addFollowup(KILL);
break;
case LOST:
addFollowup(KILL);
addFollowup(RESCHEDULE);
break;
case FINISHED:
addFollowup(RESCHEDULE);
break;
case FAILED:
addFollowup(RESCHEDULE);
break;
case KILLED:
addFollowup(RESCHEDULE);
break;
default:
// No-op.
}
};
// To be called on a task transitioning into the FINISHED state.
final Command rescheduleIfService = () -> {
if (task.get().getAssignedTask().getTask().isIsService()) {
addFollowup(RESCHEDULE);
}
};
// To be called on a task transitioning into the FAILED state.
final Command incrementFailuresMaybeReschedule = new Command() {
@Override
public void execute() {
addFollowup(INCREMENT_FAILURES);
// Max failures is ignored for service task.
boolean isService = task.get().getAssignedTask().getTask().isIsService();
// Max failures is ignored when set to -1.
int maxFailures = task.get().getAssignedTask().getTask().getMaxTaskFailures();
boolean belowMaxFailures =
maxFailures == -1 || task.get().getFailureCount() < (maxFailures - 1);
if (isService || belowMaxFailures) {
addFollowup(RESCHEDULE);
} else {
LOG.info("Task " + name + " reached failure limit, not rescheduling");
}
}
};
final Consumer<Transition<TaskState>> deleteIfKilling =
Consumers.filter(Transition.to(KILLING), addFollowupClosure(DELETE));
stateMachine = StateMachine.<TaskState>builder(name)
.logTransitions()
.initialState(initialState)
.addState(
Rule.from(INIT)
.to(PENDING, THROTTLED))
.addState(
Rule.from(PENDING)
.to(ASSIGNED, KILLING)
.withCallback(deleteIfKilling))
.addState(
Rule.from(THROTTLED)
.to(PENDING, KILLING)
.withCallback(deleteIfKilling))
.addState(
Rule.from(ASSIGNED)
.to(STARTING, RUNNING, FINISHED, FAILED, RESTARTING, DRAINING,
KILLED, KILLING, LOST, PREEMPTING)
.withCallback(
transition -> {
switch (transition.getTo()) {
case FINISHED:
rescheduleIfService.execute();
break;
case PREEMPTING:
addFollowup(KILL);
break;
case FAILED:
incrementFailuresMaybeReschedule.execute();
break;
case RESTARTING:
addFollowup(KILL);
break;
case DRAINING:
addFollowup(KILL);
break;
case KILLED:
addFollowup(RESCHEDULE);
break;
case LOST:
addFollowup(RESCHEDULE);
addFollowup(KILL);
break;
case KILLING:
addFollowup(KILL);
break;
default:
// No-op.
}
}
))
.addState(
Rule.from(STARTING)
.to(RUNNING, FINISHED, FAILED, RESTARTING, DRAINING, KILLING,
KILLED, LOST, PREEMPTING)
.withCallback(
transition -> {
switch (transition.getTo()) {
case FINISHED:
rescheduleIfService.execute();
break;
case RESTARTING:
addFollowup(KILL);
break;
case DRAINING:
addFollowup(KILL);
break;
case PREEMPTING:
addFollowup(KILL);
break;
case FAILED:
incrementFailuresMaybeReschedule.execute();
break;
case KILLED:
addFollowup(RESCHEDULE);
break;
case KILLING:
addFollowup(KILL);
break;
case LOST:
addFollowup(RESCHEDULE);
break;
default:
// No-op.
}
}
))
.addState(
Rule.from(RUNNING)
.to(FINISHED, RESTARTING, DRAINING, FAILED, KILLING, KILLED, LOST, PREEMPTING)
.withCallback(
transition -> {
switch (transition.getTo()) {
case FINISHED:
rescheduleIfService.execute();
break;
case PREEMPTING:
addFollowup(KILL);
break;
case RESTARTING:
addFollowup(KILL);
break;
case DRAINING:
addFollowup(KILL);
break;
case FAILED:
incrementFailuresMaybeReschedule.execute();
break;
case KILLED:
addFollowup(RESCHEDULE);
break;
case KILLING:
addFollowup(KILL);
break;
case LOST:
addFollowup(RESCHEDULE);
break;
default:
// No-op.
}
}
))
.addState(
Rule.from(FINISHED)
.to(DELETED)
.withCallback(manageTerminatedTasks))
.addState(
Rule.from(PREEMPTING)
.to(FINISHED, FAILED, KILLING, KILLED, LOST)
.withCallback(manageRestartingTask))
.addState(
Rule.from(RESTARTING)
.to(FINISHED, FAILED, KILLING, KILLED, LOST)
.withCallback(manageRestartingTask))
.addState(
Rule.from(DRAINING)
.to(FINISHED, FAILED, KILLING, KILLED, LOST)
.withCallback(manageRestartingTask))
.addState(
Rule.from(FAILED)
.to(DELETED)
.withCallback(manageTerminatedTasks))
.addState(
Rule.from(KILLED)
.to(DELETED)
.withCallback(manageTerminatedTasks))
// TODO(maxim): Re-evaluate if *DELETED states are valid transitions here.
.addState(
Rule.from(KILLING)
.to(FINISHED, FAILED, KILLED, LOST, DELETED)
.withCallback(manageTerminatedTasks))
.addState(
Rule.from(LOST)
.to(DELETED)
.withCallback(manageTerminatedTasks))
.addState(
Rule.from(DELETED)
.noTransitions()
.withCallback(manageTerminatedTasks))
// Since we want this action to be performed last in the transition sequence, the callback
// must be the last chained transition callback.
.onAnyTransition(
new Consumer<Transition<TaskState>>() {
@Override
public void accept(final Transition<TaskState> transition) {
if (transition.isValidStateChange()) {
TaskState from = transition.getFrom();
TaskState to = transition.getTo();
// TODO(wfarner): Clean up this hack. This is here to suppress unnecessary work
// (save followed by delete), but it shows a wart with this catch-all behavior.
// Strongly consider pushing the SAVE_STATE behavior to each transition handler.
boolean pendingDeleteHack =
!((from == PENDING || from == THROTTLED) && to == KILLING);
// Don't bother saving state of a task that is being removed.
if (to != DELETED && pendingDeleteHack) {
addFollowup(SAVE_STATE);
}
previousState = Optional.of(from);
} else {
LOG.error("Illegal state transition attempted: " + transition);
ILLEGAL_TRANSITIONS.incrementAndGet();
}
}
}
)
// TODO(wfarner): Consider alternatives to allow exceptions to surface. This would allow
// the state machine to surface illegal state transitions and propagate better information
// to the caller. As it stands, the caller must implement logic that really belongs in
// the state machine. For example, preventing RESTARTING->UPDATING transitions
// (or for that matter, almost any user-initiated state transition) is awkward.
.throwOnBadTransition(false)
.build();
}
private void addFollowup(Action action) {
addFollowup(new SideEffect(action, Optional.absent()));
}
private void addFollowup(SideEffect sideEffect) {
LOG.debug("Adding work command {} for {}", sideEffect, this);
sideEffects.add(sideEffect);
}
private Consumer<Transition<TaskState>> addFollowupClosure(final Action action) {
return item -> addFollowup(action);
}
/**
* Attempt to transition the state machine to the provided state.
* At the time this method returns, any work commands required to satisfy the state transition
* will be appended to the work queue.
*
* TODO(maxim): The current StateManager/TaskStateMachine interaction makes it hard to expose
* a dedicated task deletion method without leaking out the state machine implementation details.
* Consider refactoring here to allow for an unambiguous task deletion without resorting to
* Optional.absent().
*
* @param status Status to apply to the task or absent if a task deletion is required.
* @return {@code true} if the state change was allowed, {@code false} otherwise.
*/
public synchronized TransitionResult updateState(final Optional<ScheduleStatus> status) {
requireNonNull(status);
Preconditions.checkState(sideEffects.isEmpty());
/**
* Don't bother applying noop state changes. If we end up modifying task state without a
* state transition (e.g. storing resource consumption of a running task), we need to find
* a different way to suppress noop transitions.
*/
TaskState taskState = status.transform(STATUS_TO_TASK_STATE).or(DELETED);
if (stateMachine.getState() == taskState) {
return new TransitionResult(NOOP, ImmutableSet.of());
}
boolean success = stateMachine.transition(taskState);
ImmutableSet<SideEffect> transitionEffects = ImmutableSet.copyOf(sideEffects);
sideEffects.clear();
if (success) {
return new TransitionResult(SUCCESS, transitionEffects);
}
return new TransitionResult(
transitionEffects.isEmpty() ? ILLEGAL : ILLEGAL_WITH_SIDE_EFFECTS,
transitionEffects);
}
/**
* Gets the previous state of this state machine.
*
* @return The state machine's previous state, or {@code null} if the state machine has not
* transitioned since being created.
*/
@Nullable
ScheduleStatus getPreviousState() {
return previousState.transform(item -> item.getStatus().orNull()).orNull();
}
@Override
public String toString() {
return stateMachine.getName();
}
}