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

import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;

import javax.inject.Inject;
import javax.inject.Singleton;

import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.annotations.VisibleForTesting;
import com.google.inject.AbstractModule;
import com.google.inject.PrivateModule;
import com.google.inject.TypeLiteral;

import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.common.util.BackoffStrategy;
import org.apache.aurora.common.util.TruncatedBinaryBackoff;
import org.apache.aurora.scheduler.BatchWorker;
import org.apache.aurora.scheduler.BatchWorker.NoResult;
import org.apache.aurora.scheduler.SchedulerServicesModule;
import org.apache.aurora.scheduler.base.AsyncUtil;
import org.apache.aurora.scheduler.base.TaskGroupKey;
import org.apache.aurora.scheduler.config.CliOptions;
import org.apache.aurora.scheduler.config.types.TimeAmount;
import org.apache.aurora.scheduler.events.PubsubEventModule;
import org.apache.aurora.scheduler.preemptor.BiCache;
import org.apache.aurora.scheduler.preemptor.BiCache.BiCacheSettings;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Binding module for scheduling logic and higher-level state management.
 */
public class UpdaterModule extends AbstractModule {
  private static final Logger LOG = LoggerFactory.getLogger(UpdaterModule.class);

  @Parameters(separators = "=")
  public static class Options {
    @Parameter(names = "-enable_update_affinity",
        description = "Enable best-effort affinity of task updates.",
        arity = 1)
    public boolean enableAffinity = false;

    @Parameter(names = "-update_affinity_reservation_hold_time",
        description = "How long to wait for a reserved agent to reoffer freed up resources.")
    public TimeAmount affinityExpiration = new TimeAmount(3L, Time.MINUTES);

    @Parameter(names = "-sla_aware_action_max_batch_size",
        description = "The maximum number of sla aware update actions that can be processed"
            + " in a batch.")
    public int slaAwareActionMaxBatchSize = 300;

    @Parameter(names = "-sla_aware_kill_retry_min_delay",
        description = "Minimum amount of time to wait between attempting to perform an SLA-Aware"
            + " kill on a task.")
    public TimeAmount slaAwareKillRetryMinDelay = new TimeAmount(1L, Time.MINUTES);

    @Parameter(names = "-sla_aware_kill_retry_max_delay",
        description = "Maximum amount of time to wait between attempting to perform an SLA-Aware"
            + " kill on a task.")
    public TimeAmount slaAwareKillRetryMaxDelay = new TimeAmount(5L, Time.MINUTES);
  }

  private final ScheduledExecutorService executor;
  private final Optional<UpdateActionBatchWorker> batchWorker;
  private final Options options;

  public UpdaterModule(Options options) {
    this(
        AsyncUtil.singleThreadLoggingScheduledExecutor("updater-%d", LOG),
        Optional.empty(),
        options);
  }

  @VisibleForTesting
  UpdaterModule(ScheduledExecutorService executor,
                Optional<UpdateActionBatchWorker> batchWorker,
                Options options) {

    this.executor = Objects.requireNonNull(executor);
    this.batchWorker = batchWorker;
    this.options = options;
  }

  @Override
  protected void configure() {
    install(new PrivateModule() {
      @Override
      protected void configure() {
        if (options.enableAffinity) {
          bind(BiCacheSettings.class).toInstance(
              new BiCacheSettings(options.affinityExpiration, "update_affinity"));
          bind(new TypeLiteral<BiCache<IInstanceKey, TaskGroupKey>>() { }).in(Singleton.class);
          bind(UpdateAgentReserver.class).to(UpdateAgentReserver.UpdateAgentReserverImpl.class);
          bind(UpdateAgentReserver.UpdateAgentReserverImpl.class).in(Singleton.class);
        } else {
          bind(UpdateAgentReserver.class).to(UpdateAgentReserver.NullAgentReserver.class);
          bind(UpdateAgentReserver.NullAgentReserver.class).in(Singleton.class);
        }
        expose(UpdateAgentReserver.class);

        bind(BackoffStrategy.class).toInstance(new TruncatedBinaryBackoff(
            options.slaAwareKillRetryMinDelay,
            options.slaAwareKillRetryMaxDelay));
        bind(SlaKillController.class).in(Singleton.class);

        bind(ScheduledExecutorService.class).toInstance(executor);
        bind(UpdateFactory.class).to(UpdateFactory.UpdateFactoryImpl.class);
        bind(UpdateFactory.UpdateFactoryImpl.class).in(Singleton.class);
        bind(JobUpdateController.class).to(JobUpdateControllerImpl.class);
        bind(JobUpdateControllerImpl.class).in(Singleton.class);
        expose(JobUpdateController.class);
        bind(JobUpdateEventSubscriber.class);
        expose(JobUpdateEventSubscriber.class);
      }
    });

    if (batchWorker.isPresent()) {
      bind(UpdateActionBatchWorker.class).toInstance(batchWorker.get());
    } else {
      bind(UpdateActionBatchWorker.class).in(Singleton.class);
    }
    SchedulerServicesModule.addSchedulerActiveServiceBinding(binder())
        .to(UpdateActionBatchWorker.class);

    PubsubEventModule.bindSubscriber(binder(), JobUpdateEventSubscriber.class);
    SchedulerServicesModule.addSchedulerActiveServiceBinding(binder())
        .to(JobUpdateEventSubscriber.class);
  }

  public static class UpdateActionBatchWorker extends BatchWorker<NoResult> {

    @Inject
    UpdateActionBatchWorker(CliOptions options, Storage storage, StatsProvider statsProvider) {
      super(storage, statsProvider, options.updater.slaAwareActionMaxBatchSize);
    }

    @Override
    protected String serviceName() {
      return "UpdateActionBatchWorker";
    }
  }
}