src/main/java/org/apache/aurora/scheduler/pruning/JobUpdateHistoryPruner.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.pruning;

import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import javax.inject.Inject;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Ordering;
import com.google.common.util.concurrent.AbstractScheduledService;

import org.apache.aurora.common.inject.TimedInterceptor.Timed;
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.common.util.Clock;
import org.apache.aurora.gen.JobUpdateQuery;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateSummary;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

import static org.apache.aurora.scheduler.storage.JobUpdateStore.TERMINAL_STATES;

/**
 * Prunes per-job update history on a periodic basis.
 */
class JobUpdateHistoryPruner extends AbstractScheduledService {
  private static final Logger LOG = LoggerFactory.getLogger(JobUpdateHistoryPruner.class);
  @VisibleForTesting
  static final String JOB_UPDATES_PRUNED = "job_updates_pruned";

  private final Clock clock;
  private final Storage storage;
  private final HistoryPrunerSettings settings;
  private final AtomicLong prunedUpdatesCount;

  static class HistoryPrunerSettings {
    private final Amount<Long, Time> pruneInterval;
    private final Amount<Long, Time> maxHistorySize;
    private final int maxUpdatesPerJob;

    HistoryPrunerSettings(
        Amount<Long, Time> pruneInterval,
        Amount<Long, Time> maxHistorySize,
        int maxUpdatesPerJob) {

      this.pruneInterval = requireNonNull(pruneInterval);
      this.maxHistorySize = requireNonNull(maxHistorySize);
      this.maxUpdatesPerJob = maxUpdatesPerJob;
    }
  }

  @Inject
  JobUpdateHistoryPruner(
      Clock clock,
      Storage storage,
      HistoryPrunerSettings settings,
      StatsProvider statsProvider) {

    this.clock = requireNonNull(clock);
    this.storage = requireNonNull(storage);
    this.settings = requireNonNull(settings);
    this.prunedUpdatesCount = statsProvider.makeCounter(JOB_UPDATES_PRUNED);
  }

  @Override
  protected Scheduler scheduler() {
    return Scheduler.newFixedDelaySchedule(
        settings.pruneInterval.as(Time.MILLISECONDS),
        settings.pruneInterval.as(Time.MILLISECONDS),
        TimeUnit.MILLISECONDS);
  }

  @VisibleForTesting
  void runForTest() {
    runOneIteration();
  }

  @Timed("job_update_store_prune_history")
  @Override
  protected void runOneIteration() {
    storage.write((NoResult.Quiet) storeProvider -> {
      List<IJobUpdateSummary> completedUpdates = storeProvider.getJobUpdateStore()
          .fetchJobUpdates(IJobUpdateQuery.build(
              new JobUpdateQuery().setUpdateStatuses(TERMINAL_STATES)))
          .stream()
          .map(u -> u.getUpdate().getSummary())
          .collect(Collectors.toList());

      long cutoff = clock.nowMillis() - settings.maxHistorySize.as(Time.MILLISECONDS);
      Predicate<IJobUpdateSummary> expiredFilter =
          s -> s.getState().getCreatedTimestampMs() < cutoff;

      // Set up a predicate to detect updates that have no tasks left in the store.
      Set<IJobKey> currentJobs = storeProvider.getTaskStore().getJobKeys();
      Predicate<IJobUpdateSummary> orphanFilter = u -> !currentJobs.contains(u.getKey().getJob());

      ImmutableSet.Builder<IJobUpdateKey> pruneBuilder = ImmutableSet.builder();

      // Gather updates based on time threshold or ones that have been orphaned.
      pruneBuilder.addAll(completedUpdates
          .stream()
          .filter(expiredFilter.or(orphanFilter))
          .map(IJobUpdateSummary::getKey)
          .collect(Collectors.toList()));

      Multimap<IJobKey, IJobUpdateSummary> updatesByJob = Multimaps.index(
          // Avoid counting to-be-removed expired updates.
          completedUpdates.stream().filter(expiredFilter.negate()).iterator(),
          s -> s.getKey().getJob());

      updatesByJob.asMap().values().forEach(updates -> {
        if (updates.size() > settings.maxUpdatesPerJob) {
          Ordering<IJobUpdateSummary> creationOrder = Ordering.natural()
              .onResultOf(s -> s.getState().getCreatedTimestampMs());
          pruneBuilder.addAll(creationOrder
              .leastOf(updates, updates.size() - settings.maxUpdatesPerJob)
              .stream()
              .map(IJobUpdateSummary::getKey)
              .iterator());
        }
      });

      Set<IJobUpdateKey> pruned = pruneBuilder.build();
      if (!pruned.isEmpty()) {
        storeProvider.getJobUpdateStore().removeJobUpdates(pruned);
      }

      prunedUpdatesCount.addAndGet(pruned.size());
      LOG.info(pruned.isEmpty()
          ? "No job update history to prune."
          : "Pruned job update history: " + Joiner.on(",").join(pruned));
    });
  }
}