src/main/java/org/apache/aurora/scheduler/maintenance/MaintenanceController.java - aurora - Git at Google

 /**
  * 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.maintenance;

 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import javax.inject.Inject;
 import javax.inject.Qualifier;

 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Joiner;
 import com.google.common.base.Predicates;
 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.collect.Sets;
 import com.google.common.eventbus.Subscribe;
 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.gen.HostMaintenanceRequest;
 import org.apache.aurora.gen.HostStatus;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.PercentageSlaPolicy;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.SlaPolicy;
 import org.apache.aurora.scheduler.BatchWorker;
 import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.base.InstanceKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.config.types.TimeAmount;
 import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
 import org.apache.aurora.scheduler.sla.SlaManager;
 import org.apache.aurora.scheduler.state.StateManager;
 import org.apache.aurora.scheduler.storage.AttributeStore;
 import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
 import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
 import org.apache.aurora.scheduler.storage.entities.IHostMaintenanceRequest;
 import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ISlaPolicy;
 import org.apache.mesos.v1.Protos;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.ElementType.METHOD;
 import static java.lang.annotation.ElementType.PARAMETER;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 import static java.util.Objects.requireNonNull;

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

 import static org.apache.aurora.gen.MaintenanceMode.DRAINED;
 import static org.apache.aurora.gen.MaintenanceMode.DRAINING;

 /**
  * Logic that puts hosts into maintenance mode, and triggers draining of hosts upon request.
  * All state-changing functions return their results.  Additionally, all state-changing functions
  * will ignore requests to change state of unknown hosts and subsequently omit these hosts from
  * return values.
  */
 public interface MaintenanceController {

   /**
    * Places hosts in maintenance mode.
    * Hosts in maintenance mode are less-preferred for scheduling.
    * No change will be made for hosts that are not recognized, and unrecognized hosts will not be
    * included in the result.
    *
    * @param hosts Hosts to put into maintenance mode.
    * @return The adjusted state of the hosts.
    */
   Set<IHostStatus> startMaintenance(Set<String> hosts);

   /**
    * Initiate a drain of all active tasks on {@code hosts}.
    *
    * @param hosts Hosts to drain.
    * @return The adjusted state of the hosts.  Hosts without any active tasks will be immediately
    *         moved to DRAINED.
    */
   Set<IHostStatus> drain(Set<String> hosts);

   /**
    * Initiate an SLA-aware drain of all active tasks on {@code hosts}.
    *
    * @param hosts Hosts to drain.
    * @param defaultSlaPolicy SlaPolicy to use if a task does not have an SlaPolicy.
    * @param timeoutSecs Interval after which tasks will be forcefully drained without checking SLA.
    * @return The adjusted state of the hosts. Hosts without any active tasks will be immediately
    *         moved to DRAINED.
    */
   Set<IHostStatus> slaDrain(Set<String> hosts, SlaPolicy defaultSlaPolicy, long timeoutSecs);

   /**
    * Drain tasks defined by the inverse offer.
    * This method doesn't set any host attributes.
    *
    * @param inverseOffer the inverse offer to use.
    */
   void drainForInverseOffer(Protos.InverseOffer inverseOffer);

   /**
    * Fetches the current maintenance mode of {$code host}.
    *
    * @param host Host to fetch state for.
    * @return Maintenance mode of host, {@link MaintenanceMode#NONE} if the host is not known.
    */
   MaintenanceMode getMode(String host);

   /**
    * Fetches the current state of {@code hosts}.
    *
    * @param hosts Hosts to fetch state for.
    * @return The state of the hosts.
    */
   Set<IHostStatus> getStatus(Set<String> hosts);

   /**
    * Moves {@code hosts} out of maintenance mode, returning them to mode NONE.
    *
    * @param hosts Hosts to move out of maintenance mode.
    * @return The adjusted state of the hosts.
    */
   Set<IHostStatus> endMaintenance(Set<String> hosts);

   /**
    * Records the maintenance requests for hosts and drains any active tasks from the host
    * asynchronously.
    *
    * Tasks are drained iff it will satisfy the required SLA for the task. Task's SLA is either the
    * {@link SlaPolicy} configured as part of the TaskConfig or the default {@link SlaPolicy}
    * specified as part of the maintenance request. If neither then the task is drained immediately.
    *
    * In order to avoid tasks from blocking maintenance perpetually each maintenance request has a
    * timeout after which all tasks forcefully drained.
    */
   class MaintenanceControllerImpl
       extends AbstractScheduledService implements MaintenanceController, EventSubscriber {
     private static final Logger LOG = LoggerFactory.getLogger(MaintenanceControllerImpl.class);

     @VisibleForTesting
     @Qualifier
     @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
     public @interface PollingInterval { }

     @VisibleForTesting
     static final String DRAINING_MESSAGE = "Draining machine for maintenance.";

     private static final String COUNTDOWN_MS_PARAM = "forceMaintenanceCountdownMs";

     private static final String MAINTENANCE_COUNTDOWN_STAT_NAME = "maintenance_countdown_ms";
     private static final String MISSING_MAINTENANCE_REQUEST = "missing_maintenance_request";
     private static final SlaPolicy ZERO_PERCENT_SLA = SlaPolicy.percentageSlaPolicy(
         new PercentageSlaPolicy()
             .setPercentage(0)
             .setDurationSecs(0));

     private final Storage storage;
     private final Amount<Long, Time> pollingInterval;
     private final TaskEventBatchWorker batchWorker;
     private final SlaManager slaManager;
     private final StateManager stateManager;

     private final AtomicLong missingMaintenanceCounter;
     private final LoadingCache<String, AtomicLong> maintenanceCountDownByTask;

     @Inject
     public MaintenanceControllerImpl(
         Storage storage,
         @PollingInterval Amount<Long, Time> pollingInterval,
         TaskEventBatchWorker batchWorker,
         SlaManager slaManager,
         StateManager stateManager,
         StatsProvider statsProvider) {

       this.storage = requireNonNull(storage);
       this.pollingInterval = checkNotNull(pollingInterval);
       this.batchWorker = requireNonNull(batchWorker);
       this.slaManager = requireNonNull(slaManager);
       this.stateManager = requireNonNull(stateManager);
       this.missingMaintenanceCounter = statsProvider.makeCounter(MISSING_MAINTENANCE_REQUEST);
       this.maintenanceCountDownByTask = CacheBuilder.newBuilder().build(
           new CacheLoader<String, AtomicLong>() {
             @Override
             public AtomicLong load(String key) {
               return statsProvider.makeCounter(key);
             }
           }
       );
     }

     private Set<String> drainTasksOnHost(String host, StoreProvider store) {
       Query.Builder query = Query.slaveScoped(host).active();

       List<IScheduledTask> candidates = new ArrayList<>(store.getTaskStore().fetchTasks(query));

       if (candidates.isEmpty()) {
         LOG.info("No tasks to drain on host: {}", host);
         return Collections.emptySet();
       }

       // shuffle the candidates to avoid head-of-line blocking
       Collections.shuffle(candidates);
       candidates.forEach(task -> {
         try {
           drainTask(task, store);
         } catch (ExecutionException e) {
           LOG.error("Exception when trying to drain task: {}", Tasks.id(task), e);
         }
       });

       return candidates.stream().map(Tasks::id).collect(Collectors.toSet());
     }

     private Set<IHostStatus> watchDrainingTasks(MutableStoreProvider store, Set<String> hosts) {
       LOG.info("Hosts to drain: " + hosts);
       Set<String> emptyHosts = Sets.newHashSet();
       for (String host : hosts) {
         Set<String> drainedTasks = drainTasksOnHost(host, store);
         // If there are no tasks on the host, immediately transition to DRAINED.
         if (drainedTasks.isEmpty()) {
           emptyHosts.add(host);
         }
       }

       return ImmutableSet.<IHostStatus>builder()
           .addAll(setMaintenanceMode(store, emptyHosts, DRAINED))
           .addAll(setMaintenanceMode(store, Sets.difference(hosts, emptyHosts), DRAINING))
           .build();
     }

     /**
      * Notifies the MaintenanceController that a task has changed state.
      *
      * @param change Event
      */
     @Subscribe
     public void taskChangedState(final TaskStateChange change) {
       if (Tasks.isTerminated(change.getNewState())) {
         final String host = change.getTask().getAssignedTask().getSlaveHost();
         batchWorker.execute(store -> {
           // If the task _was_ associated with a draining host, and it was the last task on the
           // host.
           Optional<IHostAttributes> attributes =
               store.getAttributeStore().getHostAttributes(host);
           if (attributes.isPresent() && attributes.get().getMode() == DRAINING) {
             Query.Builder builder = Query.slaveScoped(host).active();
             Iterable<IScheduledTask> activeTasks = store.getTaskStore().fetchTasks(builder);
             if (Iterables.isEmpty(activeTasks)) {
               LOG.info("Moving host {} into DRAINED", host);
               setMaintenanceMode(store, ImmutableSet.of(host), DRAINED);
               store.getHostMaintenanceStore().removeHostMaintenanceRequest(host);
             } else {
               LOG.info("Host {} is DRAINING with active tasks: {}", host, Tasks.ids(activeTasks));
             }
           }
           return BatchWorker.NO_RESULT;
         });
       }
     }

     @Override
     public Set<IHostStatus> startMaintenance(Set<String> hosts) {
       return storage.write(
           storeProvider -> setMaintenanceMode(storeProvider, hosts, MaintenanceMode.SCHEDULED));
     }

     private void recordMaintenanceRequests(
         MutableStoreProvider store,
         Set<String> hosts,
         SlaPolicy defaultSlaPolicy,
         long timeoutSecs) {

       hosts.forEach(
           host -> store.getHostMaintenanceStore().saveHostMaintenanceRequest(
               IHostMaintenanceRequest.build(
                   new HostMaintenanceRequest()
                       .setHost(host)
                       .setDefaultSlaPolicy(defaultSlaPolicy)
                       .setTimeoutSecs(timeoutSecs)
                       .setCreatedTimestampMs(System.currentTimeMillis()))));
     }

     @Override
     public Set<IHostStatus> drain(Set<String> hosts) {
       return storage.write(store -> {
         // Create a dummy maintenance request zero percent sla and timeout to force drain.
         recordMaintenanceRequests(store, hosts, ZERO_PERCENT_SLA, 0);
         return watchDrainingTasks(store, hosts);
       });
     }

     @Override
     public Set<IHostStatus> slaDrain(
         Set<String> hosts,
         SlaPolicy defaultSlaPolicy,
         long timeoutSecs) {

       // We can have only one maintenance request per host at any time.
       // So we will simply overwrite any existing request. If the current one is actively handled,
       // during the write, the new one will just be a no-op, since the host is already being
       // drained. If host is in DRAINED it will be moved back into DRAINING and then back into
       // DRAINED without having to perform any work.
       return storage.write(store -> {
         recordMaintenanceRequests(store, hosts, defaultSlaPolicy, timeoutSecs);
         return setMaintenanceMode(store, hosts, DRAINING);
       });
     }

     private Optional<String> getHostname(Protos.InverseOffer offer) {
       if (offer.getUrl().getAddress().hasHostname()) {
         return Optional.of(offer.getUrl().getAddress().getHostname());
       } else {
         return Optional.empty();
       }
     }

     @Override
     public void drainForInverseOffer(Protos.InverseOffer offer) {
       // TaskStore does not allow for querying by agent id.
       Optional<String> hostname = getHostname(offer);

       if (hostname.isPresent()) {
         String host = hostname.get();
         storage.write(storeProvider -> {
           // Create a dummy maintenance request zero percent sla and timeout to force drain.
           recordMaintenanceRequests(storeProvider, ImmutableSet.of(host), ZERO_PERCENT_SLA, 0);
           return drainTasksOnHost(host, storeProvider);
         });
       } else {
         LOG.error("Unable to drain tasks on agent {} because "
             + "no hostname attached to inverse offer {}.", offer.getAgentId(), offer.getId());
       }
     }

     private static final Function<IHostAttributes, String> HOST_NAME =
         IHostAttributes::getHost;

     private static final Function<IHostAttributes, IHostStatus> ATTRS_TO_STATUS =
         attributes -> IHostStatus.build(
             new HostStatus().setHost(attributes.getHost()).setMode(attributes.getMode()));

     private static final Function<IHostStatus, MaintenanceMode> GET_MODE = IHostStatus::getMode;

     @Override
     public MaintenanceMode getMode(final String host) {
       return storage.read(storeProvider -> storeProvider.getAttributeStore().getHostAttributes(host)
           .map(ATTRS_TO_STATUS)
           .map(GET_MODE)
           .orElse(MaintenanceMode.NONE));
     }

     @Override
     public Set<IHostStatus> getStatus(final Set<String> hosts) {
       return storage.read(storeProvider -> {
         // Warning - this is filtering _all_ host attributes.  If using this to frequently query
         // for a small set of hosts, a getHostAttributes variant should be added.
         return storeProvider.getAttributeStore().getHostAttributes().stream()
             .filter(Predicates.compose(Predicates.in(hosts), HOST_NAME))
             .map(ATTRS_TO_STATUS)
             .collect(Collectors.toSet());
       });
     }

     @Override
     public Set<IHostStatus> endMaintenance(final Set<String> hosts) {
       return storage.write(
           storeProvider -> {
             hosts.forEach(
                 h -> storeProvider.getHostMaintenanceStore().removeHostMaintenanceRequest(h));
             return setMaintenanceMode(storeProvider, hosts, MaintenanceMode.NONE);
           });
     }

     private Set<IHostStatus> setMaintenanceMode(
         MutableStoreProvider storeProvider,
         Set<String> hosts,
         MaintenanceMode mode) {

       AttributeStore.Mutable store = storeProvider.getAttributeStore();
       ImmutableSet.Builder<IHostStatus> statuses = ImmutableSet.builder();
       for (String host : hosts) {
         LOG.info("Setting maintenance mode to {} for host {}", mode, host);
         Optional<IHostAttributes> toSave = AttributeStore.Util.mergeMode(store, host, mode);
         if (toSave.isPresent()) {
           store.saveHostAttributes(toSave.get());
           LOG.info("Updated host attributes: " + toSave.get());
           statuses.add(IHostStatus.build(new HostStatus().setHost(host).setMode(mode)));
         }
       }
       return statuses.build();
     }

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

     @Timed
     @Override
     protected void runOneIteration() {
       LOG.info("Looking for hosts in DRAINING state");
       storage.read(store -> {
         store.getAttributeStore()
             .getHostAttributes()
             .stream()
             .filter(h -> h.getMode() == DRAINING)
             .forEach(h -> {
               if (drainTasksOnHost(h.getHost(), store).isEmpty()) {
                 storage.write(mutable -> setMaintenanceMode(
                     mutable,
                     ImmutableSet.of(h.getHost()),
                     DRAINED));
               }
             });
         return null;
       });
     }

     @Override
     protected Scheduler scheduler() {
       return Scheduler.newFixedDelaySchedule(
           pollingInterval.getValue(),
           pollingInterval.getValue(),
           pollingInterval.getUnit().getTimeUnit());
     }

     private void drainTask(IScheduledTask task, StoreProvider store) throws ExecutionException {
       String host = task.getAssignedTask().getSlaveHost();
       Optional<IHostMaintenanceRequest> hostMaintenanceRequest =
           store.getHostMaintenanceStore().getHostMaintenanceRequest(host);
       if (!hostMaintenanceRequest.isPresent()) {
         LOG.error("No maintenance request found for host: {}. Assuming SLA not satisfied.", host);
         missingMaintenanceCounter.incrementAndGet();
         return;
       }

       boolean force = false;
       long startMs = hostMaintenanceRequest.get().getCreatedTimestampMs();
       long timeoutMs = TimeAmount.of(
             hostMaintenanceRequest.get().getTimeoutSecs(),
             Time.SECONDS)
           .as(Time.MILLISECONDS);
       long endMs = startMs + timeoutMs;
       long remainingMs = endMs - System.currentTimeMillis();
       maintenanceCountDownByTask.get(
           Joiner.on("_")
               .join(
                   MAINTENANCE_COUNTDOWN_STAT_NAME,
                   InstanceKeys.toString(Tasks.getJob(task), Tasks.getInstanceId(task))
               )
             )
           .getAndSet(remainingMs);

       if (remainingMs < 0) {
         LOG.warn("Maintenance request timed out for host: {} after {} secs. Forcing drain of {}.",
             host, hostMaintenanceRequest.get().getTimeoutSecs(), Tasks.id(task));
         force = true;
       }

       final ISlaPolicy slaPolicy = task.getAssignedTask().getTask().isSetSlaPolicy()
           ? task.getAssignedTask().getTask().getSlaPolicy()
           : hostMaintenanceRequest.get().getDefaultSlaPolicy();

       slaManager.checkSlaThenAct(
           task,
           slaPolicy,
           storeProvider -> stateManager.changeState(
               storeProvider,
               Tasks.id(task),
               Optional.empty(),
               ScheduleStatus.DRAINING,
               Optional.of(DRAINING_MESSAGE)),
           ImmutableMap.of(COUNTDOWN_MS_PARAM, Long.toString(remainingMs)),
           force);
     }
   }
 }
	/**
	* 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.maintenance;

	import java.lang.annotation.Retention;
	import java.lang.annotation.Target;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.List;
	import java.util.Optional;
	import java.util.Set;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.atomic.AtomicLong;
	import java.util.stream.Collectors;
	import javax.inject.Inject;
	import javax.inject.Qualifier;

	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.base.Function;
	import com.google.common.base.Joiner;
	import com.google.common.base.Predicates;
	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.collect.Sets;
	import com.google.common.eventbus.Subscribe;
	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.gen.HostMaintenanceRequest;
	import org.apache.aurora.gen.HostStatus;
	import org.apache.aurora.gen.MaintenanceMode;
	import org.apache.aurora.gen.PercentageSlaPolicy;
	import org.apache.aurora.gen.ScheduleStatus;
	import org.apache.aurora.gen.SlaPolicy;
	import org.apache.aurora.scheduler.BatchWorker;
	import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
	import org.apache.aurora.scheduler.base.InstanceKeys;
	import org.apache.aurora.scheduler.base.Query;
	import org.apache.aurora.scheduler.base.Tasks;
	import org.apache.aurora.scheduler.config.types.TimeAmount;
	import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
	import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
	import org.apache.aurora.scheduler.sla.SlaManager;
	import org.apache.aurora.scheduler.state.StateManager;
	import org.apache.aurora.scheduler.storage.AttributeStore;
	import org.apache.aurora.scheduler.storage.Storage;
	import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
	import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
	import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
	import org.apache.aurora.scheduler.storage.entities.IHostMaintenanceRequest;
	import org.apache.aurora.scheduler.storage.entities.IHostStatus;
	import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
	import org.apache.aurora.scheduler.storage.entities.ISlaPolicy;
	import org.apache.mesos.v1.Protos;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import static java.lang.annotation.ElementType.FIELD;
	import static java.lang.annotation.ElementType.METHOD;
	import static java.lang.annotation.ElementType.PARAMETER;
	import static java.lang.annotation.RetentionPolicy.RUNTIME;
	import static java.util.Objects.requireNonNull;

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

	import static org.apache.aurora.gen.MaintenanceMode.DRAINED;
	import static org.apache.aurora.gen.MaintenanceMode.DRAINING;

	/**
	* Logic that puts hosts into maintenance mode, and triggers draining of hosts upon request.
	* All state-changing functions return their results. Additionally, all state-changing functions
	* will ignore requests to change state of unknown hosts and subsequently omit these hosts from
	* return values.
	*/
	public interface MaintenanceController {

	/**
	* Places hosts in maintenance mode.
	* Hosts in maintenance mode are less-preferred for scheduling.
	* No change will be made for hosts that are not recognized, and unrecognized hosts will not be
	* included in the result.
	*
	* @param hosts Hosts to put into maintenance mode.
	* @return The adjusted state of the hosts.
	*/
	Set<IHostStatus> startMaintenance(Set<String> hosts);

	/**
	* Initiate a drain of all active tasks on {@code hosts}.
	*
	* @param hosts Hosts to drain.
	* @return The adjusted state of the hosts. Hosts without any active tasks will be immediately
	* moved to DRAINED.
	*/
	Set<IHostStatus> drain(Set<String> hosts);

	/**
	* Initiate an SLA-aware drain of all active tasks on {@code hosts}.
	*
	* @param hosts Hosts to drain.
	* @param defaultSlaPolicy SlaPolicy to use if a task does not have an SlaPolicy.
	* @param timeoutSecs Interval after which tasks will be forcefully drained without checking SLA.
	* @return The adjusted state of the hosts. Hosts without any active tasks will be immediately
	* moved to DRAINED.
	*/
	Set<IHostStatus> slaDrain(Set<String> hosts, SlaPolicy defaultSlaPolicy, long timeoutSecs);

	/**
	* Drain tasks defined by the inverse offer.
	* This method doesn't set any host attributes.
	*
	* @param inverseOffer the inverse offer to use.
	*/
	void drainForInverseOffer(Protos.InverseOffer inverseOffer);

	/**
	* Fetches the current maintenance mode of {$code host}.
	*
	* @param host Host to fetch state for.
	* @return Maintenance mode of host, {@link MaintenanceMode#NONE} if the host is not known.
	*/
	MaintenanceMode getMode(String host);

	/**
	* Fetches the current state of {@code hosts}.
	*
	* @param hosts Hosts to fetch state for.
	* @return The state of the hosts.
	*/
	Set<IHostStatus> getStatus(Set<String> hosts);

	/**
	* Moves {@code hosts} out of maintenance mode, returning them to mode NONE.
	*
	* @param hosts Hosts to move out of maintenance mode.
	* @return The adjusted state of the hosts.
	*/
	Set<IHostStatus> endMaintenance(Set<String> hosts);

	/**
	* Records the maintenance requests for hosts and drains any active tasks from the host
	* asynchronously.
	*
	* Tasks are drained iff it will satisfy the required SLA for the task. Task's SLA is either the
	* {@link SlaPolicy} configured as part of the TaskConfig or the default {@link SlaPolicy}
	* specified as part of the maintenance request. If neither then the task is drained immediately.
	*
	* In order to avoid tasks from blocking maintenance perpetually each maintenance request has a
	* timeout after which all tasks forcefully drained.
	*/
	class MaintenanceControllerImpl
	extends AbstractScheduledService implements MaintenanceController, EventSubscriber {
	private static final Logger LOG = LoggerFactory.getLogger(MaintenanceControllerImpl.class);

	@VisibleForTesting
	@Qualifier
	@Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
	public @interface PollingInterval { }

	@VisibleForTesting
	static final String DRAINING_MESSAGE = "Draining machine for maintenance.";

	private static final String COUNTDOWN_MS_PARAM = "forceMaintenanceCountdownMs";

	private static final String MAINTENANCE_COUNTDOWN_STAT_NAME = "maintenance_countdown_ms";
	private static final String MISSING_MAINTENANCE_REQUEST = "missing_maintenance_request";
	private static final SlaPolicy ZERO_PERCENT_SLA = SlaPolicy.percentageSlaPolicy(
	new PercentageSlaPolicy()
	.setPercentage(0)
	.setDurationSecs(0));

	private final Storage storage;
	private final Amount<Long, Time> pollingInterval;
	private final TaskEventBatchWorker batchWorker;
	private final SlaManager slaManager;
	private final StateManager stateManager;

	private final AtomicLong missingMaintenanceCounter;
	private final LoadingCache<String, AtomicLong> maintenanceCountDownByTask;

	@Inject
	public MaintenanceControllerImpl(
	Storage storage,
	@PollingInterval Amount<Long, Time> pollingInterval,
	TaskEventBatchWorker batchWorker,
	SlaManager slaManager,
	StateManager stateManager,
	StatsProvider statsProvider) {

	this.storage = requireNonNull(storage);
	this.pollingInterval = checkNotNull(pollingInterval);
	this.batchWorker = requireNonNull(batchWorker);
	this.slaManager = requireNonNull(slaManager);
	this.stateManager = requireNonNull(stateManager);
	this.missingMaintenanceCounter = statsProvider.makeCounter(MISSING_MAINTENANCE_REQUEST);
	this.maintenanceCountDownByTask = CacheBuilder.newBuilder().build(
	new CacheLoader<String, AtomicLong>() {
	@Override
	public AtomicLong load(String key) {
	return statsProvider.makeCounter(key);
	}
	}
	);
	}

	private Set<String> drainTasksOnHost(String host, StoreProvider store) {
	Query.Builder query = Query.slaveScoped(host).active();

	List<IScheduledTask> candidates = new ArrayList<>(store.getTaskStore().fetchTasks(query));

	if (candidates.isEmpty()) {
	LOG.info("No tasks to drain on host: {}", host);
	return Collections.emptySet();
	}

	// shuffle the candidates to avoid head-of-line blocking
	Collections.shuffle(candidates);
	candidates.forEach(task -> {
	try {
	drainTask(task, store);
	} catch (ExecutionException e) {
	LOG.error("Exception when trying to drain task: {}", Tasks.id(task), e);
	}
	});

	return candidates.stream().map(Tasks::id).collect(Collectors.toSet());
	}

	private Set<IHostStatus> watchDrainingTasks(MutableStoreProvider store, Set<String> hosts) {
	LOG.info("Hosts to drain: " + hosts);
	Set<String> emptyHosts = Sets.newHashSet();
	for (String host : hosts) {
	Set<String> drainedTasks = drainTasksOnHost(host, store);
	// If there are no tasks on the host, immediately transition to DRAINED.
	if (drainedTasks.isEmpty()) {
	emptyHosts.add(host);
	}
	}

	return ImmutableSet.<IHostStatus>builder()
	.addAll(setMaintenanceMode(store, emptyHosts, DRAINED))
	.addAll(setMaintenanceMode(store, Sets.difference(hosts, emptyHosts), DRAINING))
	.build();
	}

	/**
	* Notifies the MaintenanceController that a task has changed state.
	*
	* @param change Event
	*/
	@Subscribe
	public void taskChangedState(final TaskStateChange change) {
	if (Tasks.isTerminated(change.getNewState())) {
	final String host = change.getTask().getAssignedTask().getSlaveHost();
	batchWorker.execute(store -> {
	// If the task _was_ associated with a draining host, and it was the last task on the
	// host.
	Optional<IHostAttributes> attributes =
	store.getAttributeStore().getHostAttributes(host);
	if (attributes.isPresent() && attributes.get().getMode() == DRAINING) {
	Query.Builder builder = Query.slaveScoped(host).active();
	Iterable<IScheduledTask> activeTasks = store.getTaskStore().fetchTasks(builder);
	if (Iterables.isEmpty(activeTasks)) {
	LOG.info("Moving host {} into DRAINED", host);
	setMaintenanceMode(store, ImmutableSet.of(host), DRAINED);
	store.getHostMaintenanceStore().removeHostMaintenanceRequest(host);
	} else {
	LOG.info("Host {} is DRAINING with active tasks: {}", host, Tasks.ids(activeTasks));
	}
	}
	return BatchWorker.NO_RESULT;
	});
	}
	}

	@Override
	public Set<IHostStatus> startMaintenance(Set<String> hosts) {
	return storage.write(
	storeProvider -> setMaintenanceMode(storeProvider, hosts, MaintenanceMode.SCHEDULED));
	}

	private void recordMaintenanceRequests(
	MutableStoreProvider store,
	Set<String> hosts,
	SlaPolicy defaultSlaPolicy,
	long timeoutSecs) {

	hosts.forEach(
	host -> store.getHostMaintenanceStore().saveHostMaintenanceRequest(
	IHostMaintenanceRequest.build(
	new HostMaintenanceRequest()
	.setHost(host)
	.setDefaultSlaPolicy(defaultSlaPolicy)
	.setTimeoutSecs(timeoutSecs)
	.setCreatedTimestampMs(System.currentTimeMillis()))));
	}

	@Override
	public Set<IHostStatus> drain(Set<String> hosts) {
	return storage.write(store -> {
	// Create a dummy maintenance request zero percent sla and timeout to force drain.
	recordMaintenanceRequests(store, hosts, ZERO_PERCENT_SLA, 0);
	return watchDrainingTasks(store, hosts);
	});
	}

	@Override
	public Set<IHostStatus> slaDrain(
	Set<String> hosts,
	SlaPolicy defaultSlaPolicy,
	long timeoutSecs) {

	// We can have only one maintenance request per host at any time.
	// So we will simply overwrite any existing request. If the current one is actively handled,
	// during the write, the new one will just be a no-op, since the host is already being
	// drained. If host is in DRAINED it will be moved back into DRAINING and then back into
	// DRAINED without having to perform any work.
	return storage.write(store -> {
	recordMaintenanceRequests(store, hosts, defaultSlaPolicy, timeoutSecs);
	return setMaintenanceMode(store, hosts, DRAINING);
	});
	}

	private Optional<String> getHostname(Protos.InverseOffer offer) {
	if (offer.getUrl().getAddress().hasHostname()) {
	return Optional.of(offer.getUrl().getAddress().getHostname());
	} else {
	return Optional.empty();
	}
	}

	@Override
	public void drainForInverseOffer(Protos.InverseOffer offer) {
	// TaskStore does not allow for querying by agent id.
	Optional<String> hostname = getHostname(offer);

	if (hostname.isPresent()) {
	String host = hostname.get();
	storage.write(storeProvider -> {
	// Create a dummy maintenance request zero percent sla and timeout to force drain.
	recordMaintenanceRequests(storeProvider, ImmutableSet.of(host), ZERO_PERCENT_SLA, 0);
	return drainTasksOnHost(host, storeProvider);
	});
	} else {
	LOG.error("Unable to drain tasks on agent {} because "
	+ "no hostname attached to inverse offer {}.", offer.getAgentId(), offer.getId());
	}
	}

	private static final Function<IHostAttributes, String> HOST_NAME =
	IHostAttributes::getHost;

	private static final Function<IHostAttributes, IHostStatus> ATTRS_TO_STATUS =
	attributes -> IHostStatus.build(
	new HostStatus().setHost(attributes.getHost()).setMode(attributes.getMode()));

	private static final Function<IHostStatus, MaintenanceMode> GET_MODE = IHostStatus::getMode;

	@Override
	public MaintenanceMode getMode(final String host) {
	return storage.read(storeProvider -> storeProvider.getAttributeStore().getHostAttributes(host)
	.map(ATTRS_TO_STATUS)
	.map(GET_MODE)
	.orElse(MaintenanceMode.NONE));
	}

	@Override
	public Set<IHostStatus> getStatus(final Set<String> hosts) {
	return storage.read(storeProvider -> {
	// Warning - this is filtering _all_ host attributes. If using this to frequently query
	// for a small set of hosts, a getHostAttributes variant should be added.
	return storeProvider.getAttributeStore().getHostAttributes().stream()
	.filter(Predicates.compose(Predicates.in(hosts), HOST_NAME))
	.map(ATTRS_TO_STATUS)
	.collect(Collectors.toSet());
	});
	}

	@Override
	public Set<IHostStatus> endMaintenance(final Set<String> hosts) {
	return storage.write(
	storeProvider -> {
	hosts.forEach(
	h -> storeProvider.getHostMaintenanceStore().removeHostMaintenanceRequest(h));
	return setMaintenanceMode(storeProvider, hosts, MaintenanceMode.NONE);
	});
	}

	private Set<IHostStatus> setMaintenanceMode(
	MutableStoreProvider storeProvider,
	Set<String> hosts,
	MaintenanceMode mode) {

	AttributeStore.Mutable store = storeProvider.getAttributeStore();
	ImmutableSet.Builder<IHostStatus> statuses = ImmutableSet.builder();
	for (String host : hosts) {
	LOG.info("Setting maintenance mode to {} for host {}", mode, host);
	Optional<IHostAttributes> toSave = AttributeStore.Util.mergeMode(store, host, mode);
	if (toSave.isPresent()) {
	store.saveHostAttributes(toSave.get());
	LOG.info("Updated host attributes: " + toSave.get());
	statuses.add(IHostStatus.build(new HostStatus().setHost(host).setMode(mode)));
	}
	}
	return statuses.build();
	}

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

	@Timed
	@Override
	protected void runOneIteration() {
	LOG.info("Looking for hosts in DRAINING state");
	storage.read(store -> {
	store.getAttributeStore()
	.getHostAttributes()
	.stream()
	.filter(h -> h.getMode() == DRAINING)
	.forEach(h -> {
	if (drainTasksOnHost(h.getHost(), store).isEmpty()) {
	storage.write(mutable -> setMaintenanceMode(
	mutable,
	ImmutableSet.of(h.getHost()),
	DRAINED));
	}
	});
	return null;
	});
	}

	@Override
	protected Scheduler scheduler() {
	return Scheduler.newFixedDelaySchedule(
	pollingInterval.getValue(),
	pollingInterval.getValue(),
	pollingInterval.getUnit().getTimeUnit());
	}

	private void drainTask(IScheduledTask task, StoreProvider store) throws ExecutionException {
	String host = task.getAssignedTask().getSlaveHost();
	Optional<IHostMaintenanceRequest> hostMaintenanceRequest =
	store.getHostMaintenanceStore().getHostMaintenanceRequest(host);
	if (!hostMaintenanceRequest.isPresent()) {
	LOG.error("No maintenance request found for host: {}. Assuming SLA not satisfied.", host);
	missingMaintenanceCounter.incrementAndGet();
	return;
	}

	boolean force = false;
	long startMs = hostMaintenanceRequest.get().getCreatedTimestampMs();
	long timeoutMs = TimeAmount.of(
	hostMaintenanceRequest.get().getTimeoutSecs(),
	Time.SECONDS)
	.as(Time.MILLISECONDS);
	long endMs = startMs + timeoutMs;
	long remainingMs = endMs - System.currentTimeMillis();
	maintenanceCountDownByTask.get(
	Joiner.on("_")
	.join(
	MAINTENANCE_COUNTDOWN_STAT_NAME,
	InstanceKeys.toString(Tasks.getJob(task), Tasks.getInstanceId(task))
	)
	)
	.getAndSet(remainingMs);

	if (remainingMs < 0) {
	LOG.warn("Maintenance request timed out for host: {} after {} secs. Forcing drain of {}.",
	host, hostMaintenanceRequest.get().getTimeoutSecs(), Tasks.id(task));
	force = true;
	}

	final ISlaPolicy slaPolicy = task.getAssignedTask().getTask().isSetSlaPolicy()
	? task.getAssignedTask().getTask().getSlaPolicy()
	: hostMaintenanceRequest.get().getDefaultSlaPolicy();

	slaManager.checkSlaThenAct(
	task,
	slaPolicy,
	storeProvider -> stateManager.changeState(
	storeProvider,
	Tasks.id(task),
	Optional.empty(),
	ScheduleStatus.DRAINING,
	Optional.of(DRAINING_MESSAGE)),
	ImmutableMap.of(COUNTDOWN_MS_PARAM, Long.toString(remainingMs)),
	force);
	}
	}
	}