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.List;
  import java.util.Optional;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.stream.Collectors;
  
  import javax.inject.Inject;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.base.Function;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.collect.Iterables;
  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.scheduler.base.Query;
  import org.apache.aurora.scheduler.base.Tasks;
  import org.apache.aurora.scheduler.mesos.Driver;
  import org.apache.aurora.scheduler.reconciliation.ReconciliationModule.BackgroundWorker;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.mesos.v1.Protos;
  import org.apache.mesos.v1.Protos.TaskStatus;
  
  import static java.util.Objects.requireNonNull;
  
  import static com.google.common.base.Preconditions.checkArgument;
  
  import static org.apache.aurora.common.quantity.Time.MINUTES;
  import static org.apache.aurora.common.quantity.Time.SECONDS;
  
  /**
   * A task reconciler that periodically triggers Mesos (implicit) and Aurora (explicit) task
   * reconciliation to synchronize global task states. More on task reconciliation:
   * http://mesos.apache.org/documentation/latest/reconciliation.
   */
  public class TaskReconciler extends AbstractIdleService {
  
    @VisibleForTesting
    static final String EXPLICIT_STAT_NAME = "reconciliation_explicit_runs";
  
    @VisibleForTesting
    static final String IMPLICIT_STAT_NAME = "reconciliation_implicit_runs";
  
    private final TaskReconcilerSettings settings;
    private final Storage storage;
    private final Driver driver;
    private final ScheduledExecutorService executor;
    private final AtomicLong explicitRuns;
    private final AtomicLong implicitRuns;
  
    static class TaskReconcilerSettings {
      private final Amount<Long, Time> explicitInterval;
      private final Amount<Long, Time> implicitInterval;
      private final long explicitDelayMinutes;
      private final long implicitDelayMinutes;
      private final long explicitBatchDelaySeconds;
      private final int explicitBatchSize;
  
      @VisibleForTesting
      TaskReconcilerSettings(
          Amount<Long, Time> initialDelay,
          Amount<Long, Time> explicitInterval,
          Amount<Long, Time> implicitInterval,
          Amount<Long, Time> scheduleSpread,
          Amount<Long, Time> explicitBatchInterval,
          int explicitBatchSize) {
  
        this.explicitInterval = requireNonNull(explicitInterval);
        this.implicitInterval = requireNonNull(implicitInterval);
        explicitDelayMinutes = requireNonNull(initialDelay).as(MINUTES);
        implicitDelayMinutes = initialDelay.as(MINUTES) + scheduleSpread.as(MINUTES);
        explicitBatchDelaySeconds = explicitBatchInterval.as(SECONDS);
        this.explicitBatchSize = explicitBatchSize;
  
        checkArgument(
            explicitDelayMinutes >= 0,
            "Invalid explicit reconciliation delay: %s", explicitDelayMinutes);
        checkArgument(
            implicitDelayMinutes >= 0L,
            "Invalid implicit reconciliation delay: %s", implicitDelayMinutes);
        checkArgument(
            explicitBatchDelaySeconds >= 0L,
            "Invalid explicit batch reconciliation delay: %s", explicitBatchDelaySeconds
        );
      }
    }
  
    @Inject
    TaskReconciler(
        TaskReconcilerSettings settings,
        Storage storage,
        Driver driver,
        @BackgroundWorker ScheduledExecutorService executor,
        StatsProvider stats) {
  
      this.settings = requireNonNull(settings);
      this.storage = requireNonNull(storage);
      this.driver = requireNonNull(driver);
      this.executor = requireNonNull(executor);
      this.explicitRuns = stats.makeCounter(EXPLICIT_STAT_NAME);
      this.implicitRuns = stats.makeCounter(IMPLICIT_STAT_NAME);
    }
  
    public void triggerExplicitReconciliation(Optional<Integer> batchSize) {
      doExplicitReconcile(batchSize.orElse(settings.explicitBatchSize));
    }
  
    public void triggerImplicitReconciliation() {
      doImplicitReconcile();
    }
  
    @Override
    protected void startUp() {
      scheduleExplicitReconciliation();
      scheduleImplicitReconciliation();
    }
  
    private void scheduleExplicitReconciliation() {
      executor.scheduleAtFixedRate(
          () -> doExplicitReconcile(settings.explicitBatchSize),
          settings.explicitDelayMinutes,
          settings.explicitInterval.as(MINUTES),
          MINUTES.getTimeUnit());
    }
  
    private void scheduleImplicitReconciliation() {
      executor.scheduleAtFixedRate(
          () -> doImplicitReconcile(),
          settings.implicitDelayMinutes,
          settings.implicitInterval.as(MINUTES),
          MINUTES.getTimeUnit());
    }
  
    private void doImplicitReconcile() {
      driver.reconcileTasks(ImmutableSet.of());
      implicitRuns.incrementAndGet();
    }
  
    private void doExplicitReconcile(int batchSize) {
      Iterable<List<IScheduledTask>> activeBatches = Iterables.partition(
          Storage.Util.fetchTasks(storage, Query.unscoped().byStatus(Tasks.SLAVE_ASSIGNED_STATES)),
          batchSize);
  
      long delay = 0;
      for (List<IScheduledTask> batch : activeBatches) {
        executor.schedule(() -> driver.reconcileTasks(
            batch.stream().map(TASK_TO_PROTO::apply).collect(Collectors.toList())),
            delay,
            SECONDS.getTimeUnit());
        delay += settings.explicitBatchDelaySeconds;
      }
      explicitRuns.incrementAndGet();
    }
  
    @Override
    protected void shutDown() {
      // Nothing to do - await VM shutdown.
    }
  
    @VisibleForTesting
    static final Function<IScheduledTask, TaskStatus> TASK_TO_PROTO = t -> TaskStatus.newBuilder()
        // TODO(maxim): State is required by protobuf but ignored by Mesos for reconciliation
        // purposes. This is the artifact of the native API. The new HTTP Mesos API will be
        // accepting task IDs instead. AURORA-1326 tracks solution on the scheduler side.
        // Setting TASK_RUNNING as a safe dummy value here.
        .setState(Protos.TaskState.TASK_RUNNING)
        .setAgentId(
            Protos.AgentID.newBuilder().setValue(t.getAssignedTask().getSlaveId()).build())
        .setTaskId(Protos.TaskID.newBuilder().setValue(t.getAssignedTask().getTaskId()).build())
        .build();
  }
  