org.apache.aurora.scheduler.mesos.SchedulerDriverService.java Source code

Introduction

Here is the source code for org.apache.aurora.scheduler.mesos.SchedulerDriverService.java

Source

/**
 * 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.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicLong;

import javax.inject.Inject;

import com.google.common.base.Optional;
import com.google.common.collect.Collections2;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.AbstractIdleService;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.SettableFuture;

import org.apache.aurora.common.stats.Stats;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.mesos.Protos.Filters;
import org.apache.mesos.Protos.Offer.Operation;
import org.apache.mesos.Protos.OfferID;
import org.apache.mesos.Protos.Status;
import org.apache.mesos.Protos.TaskID;
import org.apache.mesos.Protos.TaskStatus;
import org.apache.mesos.Scheduler;
import org.apache.mesos.SchedulerDriver;
import org.apache.mesos.v1.Protos;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

import static com.google.common.base.Preconditions.checkState;

import static org.apache.mesos.Protos.Status.DRIVER_RUNNING;

/**
 * Manages the lifecycle of the scheduler driver, and provides a more constrained API to use it.
 */
class SchedulerDriverService extends AbstractIdleService implements Driver {
    private static final Logger LOG = LoggerFactory.getLogger(SchedulerDriverService.class);

    private final AtomicLong killFailures = Stats.exportLong("scheduler_driver_kill_failures");
    private final DriverFactory driverFactory;
    private final FrameworkInfoFactory infoFactory;

    private final Scheduler scheduler;
    private final Storage storage;
    private final DriverSettings driverSettings;
    private final SettableFuture<SchedulerDriver> driverFuture = SettableFuture.create();

    @Inject
    SchedulerDriverService(Scheduler scheduler, Storage storage, DriverSettings driverSettings,
            DriverFactory driverFactory, FrameworkInfoFactory infoFactory) {

        this.scheduler = requireNonNull(scheduler);
        this.storage = requireNonNull(storage);
        this.driverSettings = requireNonNull(driverSettings);
        this.driverFactory = requireNonNull(driverFactory);
        this.infoFactory = requireNonNull(infoFactory);
    }

    @Override
    protected void startUp() {
        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!");
        }

        Protos.FrameworkInfo.Builder frameworkBuilder = infoFactory.getFrameworkInfo().toBuilder();

        if (frameworkId.isPresent()) {
            LOG.info("Found persisted framework ID: " + frameworkId);
            frameworkBuilder.setId(Protos.FrameworkID.newBuilder().setValue(frameworkId.get()));
        } else {
            LOG.warn("Did not find a persisted framework ID, connecting as a new framework.");
        }

        SchedulerDriver schedulerDriver = driverFactory.create(scheduler, driverSettings.getCredentials(),
                frameworkBuilder.build(), driverSettings.getMasterUri());
        Status status = schedulerDriver.start();
        LOG.info("Driver started with code " + status);

        driverFuture.set(schedulerDriver);
    }

    @Override
    public void blockUntilStopped() {
        Futures.getUnchecked(driverFuture).join();
    }

    @Override
    protected void shutDown() throws ExecutionException, InterruptedException {
        // WARNING: stop() and stop(false) are dangerous, avoid at all costs. See the docs for
        // SchedulerDriver for more details.
        driverFuture.get().stop(true /* failover */);
    }

    @Override
    public void abort() {
        Futures.getUnchecked(driverFuture).abort();
    }

    @Override
    public void acceptOffers(Protos.OfferID offerId, Collection<Protos.Offer.Operation> operations,
            Protos.Filters filter) {
        ensureRunning();

        OfferID convertedOfferId = ProtosConversion.convert(offerId);
        Collection<Operation> convertedOperations = Collections2.transform(operations, ProtosConversion::convert);
        Filters convertedFilter = ProtosConversion.convert(filter);

        Futures.getUnchecked(driverFuture).acceptOffers(ImmutableList.of(convertedOfferId), convertedOperations,
                convertedFilter);
    }

    @Override
    public void acceptInverseOffer(Protos.OfferID offerID, Protos.Filters filter) {
        throw new UnsupportedOperationException("SchedulerDriver does not support inverse offers");
    }

    @Override
    public void declineOffer(Protos.OfferID offerId, Protos.Filters filter) {
        ensureRunning();

        OfferID convertedOfferId = ProtosConversion.convert(offerId);
        Filters convertedFilter = ProtosConversion.convert(filter);

        Futures.getUnchecked(driverFuture).declineOffer(convertedOfferId, convertedFilter);
    }

    @Override
    public void killTask(String taskId) {
        ensureRunning();
        Status status = Futures.getUnchecked(driverFuture).killTask(TaskID.newBuilder().setValue(taskId).build());

        if (status != DRIVER_RUNNING) {
            LOG.error("Attempt to kill task {} failed with code {}", taskId, status);
            killFailures.incrementAndGet();
        }
    }

    @Override
    public void acknowledgeStatusUpdate(Protos.TaskStatus status) {
        ensureRunning();

        TaskStatus convertedStatus = ProtosConversion.convert(status);
        Futures.getUnchecked(driverFuture).acknowledgeStatusUpdate(convertedStatus);
    }

    @Override
    public void reconcileTasks(Collection<Protos.TaskStatus> statuses) {
        ensureRunning();

        Collection<TaskStatus> convertedStatuses = Collections2.transform(statuses, ProtosConversion::convert);
        Futures.getUnchecked(driverFuture).reconcileTasks(convertedStatuses);
    }

    private void ensureRunning() {
        checkState(isRunning(), "Driver is not running.");
    }
}