content: /**
   * 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.reconciliation;
  
  import java.util.EnumSet;
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicLong;
  
  import javax.inject.Inject;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.eventbus.Subscribe;
  import com.google.common.util.concurrent.AbstractIdleService;
  
  import org.apache.aurora.common.quantity.Amount;
  import org.apache.aurora.common.quantity.Time;
  import org.apache.aurora.common.stats.StatsProvider;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor;
  import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
  import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
  import org.apache.aurora.scheduler.state.StateChangeResult;
  import org.apache.aurora.scheduler.state.StateManager;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Observes task transitions and identifies tasks that are 'stuck' in a transient state.  Stuck
   * tasks will be transitioned to the LOST state.
   */
  class TaskTimeout extends AbstractIdleService implements EventSubscriber {
    private static final Logger LOG = LoggerFactory.getLogger(TaskTimeout.class);
  
    @VisibleForTesting
    static final Amount<Long, Time> NOT_STARTED_RETRY = Amount.of(5L, Time.SECONDS);
  
    @VisibleForTesting
    static final String TIMED_OUT_TASKS_COUNTER = "timed_out_tasks";
  
    @VisibleForTesting
    static final Optional<String> TIMEOUT_MESSAGE = Optional.of("Task timed out");
  
    @VisibleForTesting
    static final Set<ScheduleStatus> TRANSIENT_STATES = EnumSet.of(
        ScheduleStatus.ASSIGNED,
        ScheduleStatus.PREEMPTING,
        ScheduleStatus.RESTARTING,
        ScheduleStatus.KILLING,
        ScheduleStatus.DRAINING);
  
    private final ScheduledExecutorService executor;
    private final Storage storage;
    private final StateManager stateManager;
    private final Amount<Long, Time> timeout;
    private final AtomicLong timedOutTasks;
  
    @Inject
    TaskTimeout(
        @AsyncExecutor ScheduledExecutorService executor,
        Storage storage,
        StateManager stateManager,
        Amount<Long, Time> timeout,
        StatsProvider statsProvider) {
  
      this.executor = requireNonNull(executor);
      this.storage = requireNonNull(storage);
      this.stateManager = requireNonNull(stateManager);
      this.timeout = requireNonNull(timeout);
      this.timedOutTasks = statsProvider.makeCounter(TIMED_OUT_TASKS_COUNTER);
    }
  
    private static boolean isTransient(ScheduleStatus status) {
      return TRANSIENT_STATES.contains(status);
    }
  
    @Override
    protected void startUp() {
      // No work to do here for startup, however we leverage the state tracking in
      // AbstractIdleService.
    }
  
    @Override
    protected void shutDown() {
      // Nothing to do for shutting down.
    }
  
    private class TimedOutTaskHandler implements Runnable {
      private final String taskId;
      private final ScheduleStatus newState;
  
      TimedOutTaskHandler(String taskId, ScheduleStatus newState) {
        this.taskId = taskId;
        this.newState = newState;
      }
  
      @Override
      public void run() {
        if (isRunning()) {
          Optional<IScheduledTask> task = storage.read(
              storeProvider -> storeProvider.getTaskStore().fetchTask(taskId));
          // Double-Checked Locking: acquire storage write lock only if necessary
          if (task.isPresent() && task.get().getStatus() == newState) {
            // This query acts as a CAS by including the state that we expect the task to be in
            // if the timeout is still valid.  Ideally, the future would have already been
            // canceled, but in the event of a state transition race, including transientState
            // prevents an unintended task timeout.
            // Note: This requires LOST transitions trigger Driver.killTask.
            StateChangeResult result = storage.write(storeProvider -> stateManager.changeState(
                storeProvider,
                taskId,
                Optional.of(newState),
                ScheduleStatus.LOST,
                TIMEOUT_MESSAGE));
  
            if (result == StateChangeResult.SUCCESS) {
              LOG.info("Timeout reached for task " + taskId + ":" + taskId);
              timedOutTasks.incrementAndGet();
            }
          }
        } else {
          // Our service is not yet started.  We don't want to lose track of the task, so
          // we will try again later.
          LOG.debug("Retrying timeout of task {} in {}", taskId, NOT_STARTED_RETRY);
          // TODO(wfarner): This execution should not wait for a transaction, but a second executor
          // would be weird.
          executor.schedule(this, NOT_STARTED_RETRY.as(Time.MILLISECONDS), TimeUnit.MILLISECONDS);
        }
      }
    }
  
    @Subscribe
    public void recordStateChange(TaskStateChange change) {
      if (isTransient(change.getNewState())) {
        executor.schedule(
            new TimedOutTaskHandler(change.getTaskId(), change.getNewState()),
            timeout.as(Time.MILLISECONDS),
            TimeUnit.MILLISECONDS);
      }
    }
  }
  