src/main/java/org/apache/aurora/scheduler/TaskVars.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.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.StreamSupport;
import javax.inject.Inject;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Predicate;
import com.google.common.base.Strings;
import com.google.common.base.Supplier;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.eventbus.Subscribe;
import com.google.common.util.concurrent.AbstractIdleService;

import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.events.PubsubEvent.TasksDeleted;
import org.apache.aurora.scheduler.filter.SchedulingFilter.Veto;
import org.apache.aurora.scheduler.filter.SchedulingFilter.VetoGroup;
import org.apache.aurora.scheduler.filter.SchedulingFilter.VetoType;
import org.apache.aurora.scheduler.storage.AttributeStore;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.entities.IAttribute;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

/**
 * A container that tracks and exports stat counters for tasks.
 */
public class TaskVars extends AbstractIdleService implements EventSubscriber {
  private static final Logger LOG = LoggerFactory.getLogger(TaskVars.class);
  private static final ImmutableSet<ScheduleStatus> TRACKED_JOB_STATES =
      ImmutableSet.of(ScheduleStatus.LOST, ScheduleStatus.FAILED);

  @VisibleForTesting
  static final Map<VetoGroup, String> VETO_GROUPS_TO_COUNTERS = ImmutableMap.of(
      VetoGroup.STATIC, "scheduling_veto_static",
      VetoGroup.DYNAMIC, "scheduling_veto_dynamic",
      VetoGroup.MIXED, "scheduling_veto_mixed"
  );

  @VisibleForTesting
  static final Map<VetoType, String> VETO_TYPE_TO_COUNTERS = ImmutableMap.of(
      VetoType.CONSTRAINT_MISMATCH, "scheduling_veto_constraint_mismatch",
      VetoType.DEDICATED_CONSTRAINT_MISMATCH, "scheduling_veto_dedicated_constraint_mismatch",
      VetoType.INSUFFICIENT_RESOURCES, "scheduling_veto_insufficient_resources",
      VetoType.LIMIT_NOT_SATISFIED, "scheduling_veto_limit_not_satisfied",
      VetoType.MAINTENANCE, "scheduling_veto_maintenance"
  );

  private final LoadingCache<String, Counter> counters;
  private final LoadingCache<String, Counter> untrackedCounters;
  private final Storage storage;
  private volatile boolean exporting = false;

  @Inject
  TaskVars(Storage storage, final StatsProvider statProvider) {
    this.storage = requireNonNull(storage);
    requireNonNull(statProvider);
    counters = buildCache(statProvider);
    untrackedCounters = buildCache(statProvider.untracked());
  }

  private LoadingCache<String, Counter> buildCache(final StatsProvider provider) {
    return CacheBuilder.newBuilder().build(new CacheLoader<String, Counter>() {
      @Override
      public Counter load(String statName) {
        Counter counter = new Counter(provider);
        if (exporting) {
          counter.exportAs(statName);
        }
        return counter;
      }
    });
  }

  @VisibleForTesting
  static String getVarName(ScheduleStatus status) {
    return "task_store_" + status;
  }

  @VisibleForTesting
  static String rackStatName(String rack) {
    return "tasks_lost_rack_" + rack;
  }

  @VisibleForTesting
  static String dedicatedRoleStatName(String role) {
    return "tasks_lost_dedicated_" + role.replace("*", "_");
  }

  @VisibleForTesting
  static String jobStatName(IScheduledTask task, ScheduleStatus status) {
    return String.format(
        "tasks_%s_%s",
        status,
        JobKeys.canonicalString(task.getAssignedTask().getTask().getJob()));
  }

  private static final Predicate<IAttribute> IS_RACK = attr -> "rack".equals(attr.getName());

  private static final Function<IAttribute, String> ATTR_VALUE =
      attr -> Iterables.getOnlyElement(attr.getValues());

  private Counter getCounter(ScheduleStatus status) {
    return counters.getUnchecked(getVarName(status));
  }

  private void incrementCount(ScheduleStatus status) {
    getCounter(status).increment();
  }

  private void decrementCount(ScheduleStatus status) {
    getCounter(status).decrement();
  }

  private void updateRackCounters(IScheduledTask task, ScheduleStatus newState) {
    final String host = task.getAssignedTask().getSlaveHost();
    Optional<String> rack;
    if (Strings.isNullOrEmpty(task.getAssignedTask().getSlaveHost())) {
      rack = Optional.empty();
    } else {
      rack = storage.read(storeProvider ->
          StreamSupport.stream(
              AttributeStore.Util.attributesOrNone(storeProvider, host).spliterator(),
              false)
            .filter(IS_RACK)
            .findFirst()
            .map(ATTR_VALUE));
    }

    // Always dummy-read the lost-tasks-per-rack stat. This ensures that there is at least a zero
    // exported for all racks.
    rack.ifPresent(s -> counters.getUnchecked(rackStatName(s)));

    if (newState == ScheduleStatus.LOST) {
      if (rack.isPresent()) {
        counters.getUnchecked(rackStatName(rack.get())).increment();
      } else {
        LOG.warn("Failed to find rack attribute associated with host " + host);
      }
    }
  }

  private void updateDedicatedCounters(IScheduledTask task, ScheduleStatus newState) {
    final String host = task.getAssignedTask().getSlaveHost();
    ImmutableSet<String> dedicatedRoles;
    if (Strings.isNullOrEmpty(host)) {
      dedicatedRoles = ImmutableSet.of();
    } else {
      dedicatedRoles = storage.read(store ->
          StreamSupport.stream(
                AttributeStore.Util.attributesOrNone(store, host).spliterator(),
                false)
              .filter(attr -> "dedicated".equals(attr.getName()))
              .findFirst()
              .map(IAttribute::getValues)
              .orElse(ImmutableSet.of())
      );
    }

    // Always dummy-read the lost-tasks-per-role stat. This ensures that there is at least a zero
    // exported for all roles.
    dedicatedRoles.forEach(s -> counters.getUnchecked(dedicatedRoleStatName(s)));

    if (newState == ScheduleStatus.LOST) {
      dedicatedRoles.forEach(s -> counters.getUnchecked(dedicatedRoleStatName(s)).increment());
    }
  }

  private void updateJobCounters(IScheduledTask task, ScheduleStatus newState) {
    if (TRACKED_JOB_STATES.contains(newState)) {
      untrackedCounters.getUnchecked(jobStatName(task, newState)).increment();
    }
  }

  @Subscribe
  public void taskChangedState(TaskStateChange stateChange) {
    IScheduledTask task = stateChange.getTask();
    Optional<ScheduleStatus> previousState = stateChange.getOldState();

    if (stateChange.isTransition() && !previousState.equals(Optional.of(ScheduleStatus.INIT))) {
      decrementCount(previousState.get());
    }
    incrementCount(task.getStatus());

    updateRackCounters(task, task.getStatus());
    updateJobCounters(task, task.getStatus());
    updateDedicatedCounters(task, task.getStatus());
  }

  @Override
  protected void startUp() {
    // Dummy read the counter for each status counter. This is important to guarantee a stat with
    // value zero is present for each state, even if all states are not represented in the task
    // store.
    for (ScheduleStatus status : ScheduleStatus.values()) {
      getCounter(status);
    }

    exportCounters(counters.asMap());
    exportCounters(untrackedCounters.asMap());
  }

  @Override
  protected void shutDown() {
    // Ignored. VM shutdown is required to stop exporting task vars.
  }

  private void exportCounters(Map<String, Counter> counterMap) {
    // Initiate export of all counters.  This is not done initially to avoid exporting values that
    // do not represent the entire storage contents.
    exporting = true;
    for (Map.Entry<String, Counter> entry : counterMap.entrySet()) {
      entry.getValue().exportAs(entry.getKey());
    }
  }

  @Subscribe
  public void tasksDeleted(final TasksDeleted event) {
    for (IScheduledTask task : event.getTasks()) {
      decrementCount(task.getStatus());
    }
  }

  public void taskVetoed(Set<Veto> vetoes) {
    VetoGroup vetoGroup = Veto.identifyGroup(vetoes);
    if (vetoGroup != VetoGroup.EMPTY) {
      counters.getUnchecked(VETO_GROUPS_TO_COUNTERS.get(vetoGroup)).increment();
    }
    for (Veto veto : vetoes) {
      counters.getUnchecked(VETO_TYPE_TO_COUNTERS.get(veto.getVetoType())).increment();
    }
  }

  private static class Counter implements Supplier<Long> {
    private final AtomicLong value = new AtomicLong();
    private boolean exported = false;
    private final StatsProvider stats;

    Counter(StatsProvider stats) {
      this.stats = stats;
    }

    @Override
    public Long get() {
      return value.get();
    }

    private synchronized void exportAs(String name) {
      if (!exported) {
        stats.makeGauge(name, this);
        exported = true;
      }
    }

    private void increment() {
      value.incrementAndGet();
    }

    private void decrement() {
      value.decrementAndGet();
    }
  }
}