src/main/java/org/apache/aurora/scheduler/TaskStatusHandlerImpl.java

/**
 * 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;

import java.lang.annotation.Retention;
import java.lang.annotation.Target;
import java.util.ArrayDeque;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.atomic.AtomicReference;

import javax.inject.Inject;
import javax.inject.Qualifier;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.AbstractExecutionThreadService;
import com.google.common.util.concurrent.MoreExecutors;

import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.scheduler.base.Conversions;
import org.apache.aurora.scheduler.mesos.Driver;
import org.apache.aurora.scheduler.state.StateChangeResult;
import org.apache.aurora.scheduler.state.StateManager;
import org.apache.aurora.scheduler.stats.CachedCounters;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.mesos.v1.Protos.TaskStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.lang.annotation.ElementType.FIELD;
import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.ElementType.PARAMETER;
import static java.lang.annotation.RetentionPolicy.RUNTIME;
import static java.util.Objects.requireNonNull;

/**
 * A {@link TaskStatusHandler} implementation.
 */
@VisibleForTesting
public class TaskStatusHandlerImpl extends AbstractExecutionThreadService
    implements TaskStatusHandler {

  private static final Logger LOG = LoggerFactory.getLogger(TaskStatusHandlerImpl.class);

  @VisibleForTesting
  static final String MEMORY_LIMIT_DISPLAY = "Task used more memory than requested.";

  @VisibleForTesting
  static final String DISK_LIMIT_DISPLAY = "Task used more disk than requested.";

  private final Storage storage;
  private final StateManager stateManager;
  private final Driver driver;
  private final BlockingQueue<TaskStatus> pendingUpdates;
  private final int maxBatchSize;
  private final CachedCounters counters;

  private final AtomicReference<Thread> threadReference = new AtomicReference<>();

  /**
   * Binding annotation for the status update queue.
   */
  @VisibleForTesting
  @Qualifier
  @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
  public @interface StatusUpdateQueue { }

  /**
   * Binding annotation maximum size of a status update batch.
   */
  @VisibleForTesting
  @Qualifier
  @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
  public @interface MaxBatchSize { }

  @Inject
  TaskStatusHandlerImpl(
      Storage storage,
      StateManager stateManager,
      StatsProvider statsProvider,
      final Driver driver,
      @StatusUpdateQueue BlockingQueue<TaskStatus> pendingUpdates,
      @MaxBatchSize Integer maxBatchSize,
      CachedCounters counters) {

    this.storage = requireNonNull(storage);
    this.stateManager = requireNonNull(stateManager);
    this.driver = requireNonNull(driver);
    this.pendingUpdates = requireNonNull(pendingUpdates);
    this.maxBatchSize = requireNonNull(maxBatchSize);
    this.counters = requireNonNull(counters);
    requireNonNull(statsProvider);

    statsProvider.exportSize("status_updates_queue_size", this.pendingUpdates);

    addListener(
        new Listener() {
          @Override
          public void failed(State from, Throwable failure) {
            LOG.error("TaskStatusHandler failed: ", failure);
            driver.abort();
          }
        },
        MoreExecutors.newDirectExecutorService());
  }

  @Override
  public void statusUpdate(TaskStatus status) {
    pendingUpdates.add(status);
  }

  @Override
  protected void triggerShutdown() {
    Thread thread = threadReference.get();

    if (thread != null) {
      thread.interrupt();
    }
  }

  @Override
  protected void run() {
    threadReference.set(Thread.currentThread());

    while (isRunning()) {
      final Queue<TaskStatus> updates = new ArrayDeque<>();

      try {
        updates.add(pendingUpdates.take());
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        break;
      }

      // Process all other available updates, up to the limit on batch size.
      // TODO(bmahler): Expose histogram metrics of the batch sizes.
      pendingUpdates.drainTo(updates, maxBatchSize - updates.size());

      try {
        storage.write((NoResult.Quiet) storeProvider -> {
          for (TaskStatus status : updates) {
            ScheduleStatus translatedState = Conversions.convertProtoState(status.getState());

            StateChangeResult result = stateManager.changeState(
                storeProvider,
                status.getTaskId().getValue(),
                Optional.empty(),
                translatedState,
                formatMessage(status));

            if (status.hasReason()) {
              counters.get(statName(status, result)).incrementAndGet();
            }
          }
        });

        for (TaskStatus status : updates) {
          driver.acknowledgeStatusUpdate(status);
        }
      } catch (RuntimeException e) {
        LOG.error("Failed to process status update batch " + updates, e);
      }
    }
  }

  @VisibleForTesting
  static String statName(TaskStatus status, StateChangeResult result) {
    return "status_update_" + status.getReason() + "_" + result;
  }

  private static Optional<String> formatMessage(TaskStatus status) {
    Optional<String> message = Optional.empty();
    if (status.hasMessage()) {
      message = Optional.of(status.getMessage());
    }

    if (status.hasReason()) {
      switch (status.getReason()) {
        case REASON_CONTAINER_LIMITATION_MEMORY:
          // Add a failure explanation to the user
          if (!message.isPresent()) {
            message = Optional.of(MEMORY_LIMIT_DISPLAY);
          }
          break;

        case REASON_CONTAINER_LIMITATION_DISK:
          // Add a failure explanation to the user
          if (!message.isPresent()) {
            message = Optional.of(DISK_LIMIT_DISPLAY);
          }
          break;

        case REASON_EXECUTOR_UNREGISTERED:
          // Suppress "Unregistered executor" message as it bears no meaning to the user.
          message = Optional.empty();
          break;

        default:
          // Message is already populated above.
          break;
      }
    }

    return message;
  }
}