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.mesos;
  
  import java.util.Collection;
  import java.util.Optional;
  import java.util.concurrent.CountDownLatch;
  
  import com.google.common.collect.Collections2;
  import com.google.common.eventbus.Subscribe;
  import com.google.common.util.concurrent.AbstractIdleService;
  import com.google.common.util.concurrent.Futures;
  import com.google.common.util.concurrent.SettableFuture;
  import com.google.inject.Inject;
  
  import org.apache.aurora.common.base.Command;
  import org.apache.aurora.scheduler.events.PubsubEvent.DriverRegistered;
  import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.mesos.v1.Protos.Credential;
  import org.apache.mesos.v1.Protos.Filters;
  import org.apache.mesos.v1.Protos.FrameworkID;
  import org.apache.mesos.v1.Protos.FrameworkInfo;
  import org.apache.mesos.v1.Protos.Offer.Operation;
  import org.apache.mesos.v1.Protos.OfferID;
  import org.apache.mesos.v1.Protos.TaskID;
  import org.apache.mesos.v1.Protos.TaskStatus;
  import org.apache.mesos.v1.scheduler.Mesos;
  import org.apache.mesos.v1.scheduler.Protos.Call;
  import org.apache.mesos.v1.scheduler.Scheduler;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  import static com.google.common.base.Preconditions.checkState;
  
  /**
   * A driver implementation that uses the V1 API drivers from libmesos.
   */
  class VersionedSchedulerDriverService extends AbstractIdleService
      implements Driver, EventSubscriber {
    private static final Logger LOG = LoggerFactory.getLogger(VersionedSchedulerDriverService.class);
  
    private final Storage storage;
    private final DriverSettings driverSettings;
    private final Scheduler scheduler;
    private final VersionedDriverFactory factory;
    private final FrameworkInfoFactory infoFactory;
    private final SettableFuture<Mesos> mesosFuture = SettableFuture.create();
    private final CountDownLatch terminationLatch = new CountDownLatch(1);
    private final CountDownLatch registrationLatch = new CountDownLatch(1);
  
    @Inject
    VersionedSchedulerDriverService(
        Storage storage,
        DriverSettings settings,
        Scheduler scheduler,
        VersionedDriverFactory factory,
        FrameworkInfoFactory infoFactory) {
      this.storage = requireNonNull(storage);
      this.driverSettings = requireNonNull(settings);
      this.scheduler = requireNonNull(scheduler);
      this.factory = requireNonNull(factory);
      this.infoFactory = requireNonNull(infoFactory);
    }
  
    private FrameworkID getFrameworkId() {
      String id = storage.read(storeProvider ->
          storeProvider.getSchedulerStore().fetchFrameworkId().get());
      return FrameworkID.newBuilder().setValue(id).build();
    }
  
    @Override
    protected void startUp() throws Exception {
      Optional<String> frameworkId = storage.read(
          storeProvider -> storeProvider.getSchedulerStore().fetchFrameworkId());
  
      LOG.info("Connecting to mesos master: " + driverSettings.getMasterUri());
      if (!driverSettings.getCredentials().isPresent()) {
        LOG.warn("Connecting to master without authentication!");
      }
  
      FrameworkInfo.Builder frameworkBuilder = infoFactory.getFrameworkInfo().toBuilder();
  
      if (frameworkId.isPresent()) {
        LOG.info("Found persisted framework ID: " + frameworkId);
        frameworkBuilder.setId(FrameworkID.newBuilder().setValue(frameworkId.get()));
      } else {
        LOG.warn("Did not find a persisted framework ID, connecting as a new framework.");
      }
  
      Credential credential = driverSettings.getCredentials().orElse(null);
      Mesos mesos = factory.create(
          scheduler,
          frameworkBuilder.build(),
          driverSettings.getMasterUri(),
          Optional.ofNullable(credential));
  
      mesosFuture.set(mesos);
    }
  
    @Override
    protected void shutDown() throws Exception {
      terminationLatch.countDown();
    }
  
    @Override
    public void acceptOffers(OfferID offerId, Collection<Operation> operations, Filters filter) {
      whenRegistered(() -> {
        Collection<Operation.Type> opTypes = Collections2.transform(operations, Operation::getType);
        LOG.info("Accepting offer {} with ops {}", offerId.getValue(), opTypes);
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder()
                .setFrameworkId(getFrameworkId())
                .setType(Call.Type.ACCEPT)
                .setAccept(
                    Call.Accept.newBuilder()
                        .addOfferIds(offerId)
                        .addAllOperations(operations)
                        .setFilters(filter))
                .build());
      });
  
    }
  
    @Override
    public void acceptInverseOffer(OfferID offerID, Filters filter) {
      whenRegistered(() -> {
        LOG.info("Accepting Inverse Offer {}", offerID.getValue());
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder()
                .setFrameworkId(getFrameworkId())
                .setType(Call.Type.ACCEPT_INVERSE_OFFERS)
                .setAcceptInverseOffers(
                    Call.AcceptInverseOffers.newBuilder()
                      .addInverseOfferIds(offerID)
                      .setFilters(filter)
                )
            .build());
      });
    }
  
    @Override
    public void declineOffer(OfferID offerId, Filters filter) {
      whenRegistered(() -> {
        LOG.info("Declining offer {}", offerId.getValue());
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder().setType(Call.Type.DECLINE)
                .setFrameworkId(getFrameworkId())
                .setDecline(
                    Call.Decline.newBuilder()
                        .setFilters(filter)
                        .addOfferIds(offerId))
                .build()
        );
      });
    }
  
    @Override
    public void killTask(String taskId) {
      whenRegistered(() -> {
        LOG.info("Killing task {}", taskId);
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder().setType(Call.Type.KILL)
                .setFrameworkId(getFrameworkId())
                .setKill(
                    Call.Kill.newBuilder()
                        .setTaskId(TaskID.newBuilder().setValue(taskId)))
                .build()
        );
  
      });
    }
  
    @Override
    public void acknowledgeStatusUpdate(TaskStatus status) {
      // The Mesos API says frameworks are only supposed to acknowledge status updates
      // with a UUID. The V0Driver accepts them just fine but the V1Driver logs every time
      // a status update is given without a uuid. To silence logs, we drop them here.
  
      whenRegistered(() -> {
        if (!status.hasUuid()) {
          return;
        }
  
        LOG.info("Acking status update for {} with uuid: {}",
            status.getTaskId().getValue(),
            status.getUuid());
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder().setType(Call.Type.ACKNOWLEDGE)
                .setFrameworkId(getFrameworkId())
                .setAcknowledge(
                    Call.Acknowledge.newBuilder()
                        .setAgentId(status.getAgentId())
                        .setTaskId(status.getTaskId())
                        .setUuid(status.getUuid()))
                .build()
        );
      });
  
    }
  
    @Override
    public void reconcileTasks(Collection<TaskStatus> statuses) {
      whenRegistered(() -> {
        Collection<Call.Reconcile.Task> tasks = Collections2.transform(statuses, taskStatus ->
            Call.Reconcile.Task.newBuilder()
                .setTaskId(taskStatus.getTaskId())
                .build());
  
        Futures.getUnchecked(mesosFuture).send(
            Call.newBuilder()
                .setType(Call.Type.RECONCILE)
                .setFrameworkId(getFrameworkId())
                .setReconcile(
                    Call.Reconcile.newBuilder()
                        .addAllTasks(tasks))
                .build()
        );
      });
    }
  
    @Override
    public void blockUntilStopped() {
      ensureRunning();
      try {
        terminationLatch.await();
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
    }
  
    @Override
    public void abort() {
      terminationLatch.countDown();
      stopAsync();
    }
  
    @Subscribe
    public void registered(DriverRegistered event) {
      registrationLatch.countDown();
    }
  
    private void whenRegistered(Command c) {
      ensureRunning();
      // We need to block until registered because thats when we are guaranteed to have our
      // framework id. Without it, we cannot construct any Call objects.
      try {
        registrationLatch.await();
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
      c.execute();
    }
  
    private void ensureRunning() {
      checkState(isRunning(), "Driver is not running.");
    }
  }
  