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.benchmark;
  
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.Executor;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.TimeUnit;
  
  import javax.inject.Singleton;
  
  import com.google.common.eventbus.EventBus;
  import com.google.common.eventbus.Subscribe;
  import com.google.common.util.concurrent.AbstractExecutionThreadService;
  import com.google.common.util.concurrent.Uninterruptibles;
  import com.google.inject.AbstractModule;
  import com.google.inject.Guice;
  import com.google.inject.Injector;
  import com.google.inject.TypeLiteral;
  
  import org.apache.aurora.benchmark.fakes.FakeDriver;
  import org.apache.aurora.benchmark.fakes.FakeOfferManager;
  import org.apache.aurora.benchmark.fakes.FakeRescheduleCalculator;
  import org.apache.aurora.benchmark.fakes.FakeSchedulerDriver;
  import org.apache.aurora.benchmark.fakes.FakeStatsProvider;
  import org.apache.aurora.common.application.ShutdownStage;
  import org.apache.aurora.common.base.Command;
  import org.apache.aurora.common.base.Commands;
  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.common.util.Clock;
  import org.apache.aurora.common.util.testing.FakeClock;
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.ServerInfo;
  import org.apache.aurora.scheduler.TaskIdGenerator;
  import org.apache.aurora.scheduler.TaskStatusHandler;
  import org.apache.aurora.scheduler.TaskStatusHandlerImpl;
  import org.apache.aurora.scheduler.TierModule;
  import org.apache.aurora.scheduler.base.AsyncUtil;
  import org.apache.aurora.scheduler.base.TaskTestUtil;
  import org.apache.aurora.scheduler.config.CliOptions;
  import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
  import org.apache.aurora.scheduler.events.EventSink;
  import org.apache.aurora.scheduler.events.PubsubEvent;
  import org.apache.aurora.scheduler.filter.SchedulingFilter;
  import org.apache.aurora.scheduler.filter.SchedulingFilterImpl;
  import org.apache.aurora.scheduler.mesos.Driver;
  import org.apache.aurora.scheduler.mesos.DriverFactory;
  import org.apache.aurora.scheduler.mesos.DriverSettings;
  import org.apache.aurora.scheduler.mesos.FrameworkInfoFactory;
  import org.apache.aurora.scheduler.mesos.MesosCallbackHandler;
  import org.apache.aurora.scheduler.mesos.MesosCallbackHandler.MesosCallbackHandlerImpl;
  import org.apache.aurora.scheduler.mesos.MesosSchedulerImpl;
  import org.apache.aurora.scheduler.mesos.ProtosConversion;
  import org.apache.aurora.scheduler.mesos.SchedulerDriverModule;
  import org.apache.aurora.scheduler.mesos.TestExecutorSettings;
  import org.apache.aurora.scheduler.offers.OfferManager;
  import org.apache.aurora.scheduler.scheduling.RescheduleCalculator;
  import org.apache.aurora.scheduler.state.ClusterStateImpl;
  import org.apache.aurora.scheduler.state.StateModule;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.apache.aurora.scheduler.storage.entities.IServerInfo;
  import org.apache.aurora.scheduler.storage.mem.MemStorageModule;
  import org.apache.mesos.Scheduler;
  import org.apache.mesos.v1.Protos;
  import org.openjdk.jmh.annotations.Benchmark;
  import org.openjdk.jmh.annotations.BenchmarkMode;
  import org.openjdk.jmh.annotations.Fork;
  import org.openjdk.jmh.annotations.Level;
  import org.openjdk.jmh.annotations.Measurement;
  import org.openjdk.jmh.annotations.Mode;
  import org.openjdk.jmh.annotations.OutputTimeUnit;
  import org.openjdk.jmh.annotations.Param;
  import org.openjdk.jmh.annotations.Scope;
  import org.openjdk.jmh.annotations.Setup;
  import org.openjdk.jmh.annotations.State;
  import org.openjdk.jmh.annotations.TearDown;
  import org.openjdk.jmh.annotations.Threads;
  import org.openjdk.jmh.annotations.Warmup;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Performance benchmarks for status update processing throughput. Note that we
   * need to send many updates and wait for all transitions to occur within one run
   * of the benchmark. This is because we don't want to assume that status updates
   * are processed synchronously.
   */
  @BenchmarkMode(Mode.Throughput)
  @OutputTimeUnit(TimeUnit.SECONDS)
  @Warmup(iterations = 1, time = 10, timeUnit = TimeUnit.SECONDS)
  @Measurement(iterations = 5, time = 10, timeUnit = TimeUnit.SECONDS)
  @Fork(1)
  @Threads(1)
  @State(Scope.Thread)
  public class StatusUpdateBenchmark {
    /**
     * Simulates a slow storage backend by introducing latency on top of an
     * underlying storage implementation.
     *
     * TODO(bmahler): Consider specifying read and write latency separately.
     */
    private static final class SlowStorageWrapper implements Storage {
      private final Storage underlyingStorage;
      private Optional<Amount<Long, Time>> latency = Optional.empty();
  
      private SlowStorageWrapper(Storage underlyingStorage) {
        this.underlyingStorage = requireNonNull(underlyingStorage);
      }
  
      private void setLatency(Amount<Long, Time> latency) {
        this.latency = Optional.of(latency);
      }
  
      private void maybeSleep() {
        if (latency.isPresent()) {
          Uninterruptibles.sleepUninterruptibly(
              latency.get().getValue(),
              latency.get().getUnit().getTimeUnit());
        }
      }
  
      @Override
      public <T, E extends Exception> T read(Work<T, E> work) throws StorageException, E {
        maybeSleep();
        return underlyingStorage.read(work);
      }
  
      @Override
      public <T, E extends Exception> T write(MutateWork<T, E> work) throws StorageException, E {
        maybeSleep();
        return underlyingStorage.write(work);
      }
  
      @Override
      public void prepare() throws StorageException {
        underlyingStorage.prepare();
      }
    }
  
    // Benchmark with 1000 tasks to easily observe the kilo-qps of status
    // update processing. Consider varying this number if needed.
    private static final int NUM_TASKS = 1000;
  
    // Vary the storage latency to observe the effect on throughput.
    @Param({"5", "25", "100"})
    private long latencyMilliseconds;
  
    private Scheduler scheduler;
    private AbstractExecutionThreadService statusHandler;
    private SlowStorageWrapper storage;
    private EventBus eventBus;
    private Set<IScheduledTask> tasks;
    private CountDownLatch countDownLatch;
  
    /**
     * Run once per trial to set up the benchmark.
     */
    @Setup(Level.Trial)
    public void setUpBenchmark() {
      eventBus = new EventBus();
      storage = new SlowStorageWrapper(MemStorageModule.newEmptyStorage());
  
      Injector injector = Guice.createInjector(
          new StateModule(new CliOptions()),
          new TierModule(TaskTestUtil.TIER_CONFIG),
          new AbstractModule() {
            @Override
            protected void configure() {
              bind(Driver.class).toInstance(new FakeDriver());
              bind(Scheduler.class).to(MesosSchedulerImpl.class);
              bind(MesosSchedulerImpl.class).in(Singleton.class);
              bind(MesosCallbackHandler.class).to(MesosCallbackHandlerImpl.class);
              bind(MesosCallbackHandlerImpl.class).in(Singleton.class);
              bind(Executor.class)
                  .annotatedWith(SchedulerDriverModule.SchedulerExecutor.class)
                  .toInstance(AsyncUtil.singleThreadLoggingScheduledExecutor(
                      "SchedulerImpl-%d",
                      LoggerFactory.getLogger(StatusUpdateBenchmark.class)));
              bind(DriverFactory.class)
                  .toInstance((s, credentials, frameworkInfo, master) -> new FakeSchedulerDriver());
              bind(OfferManager.class).toInstance(new FakeOfferManager());
              bind(TaskIdGenerator.class).to(TaskIdGenerator.TaskIdGeneratorImpl.class);
              bind(SchedulingFilter.class).to(SchedulingFilterImpl.class);
              bind(Command.class).annotatedWith(ShutdownStage.class).toInstance(Commands.NOOP);
              bind(Thread.UncaughtExceptionHandler.class).toInstance(
                  (t, e) -> {
                    // no-op
                  });
              bind(Storage.class).toInstance(storage);
              bind(DriverSettings.class).toInstance(
                  new DriverSettings(
                      "fakemaster",
                      Optional.empty()));
              bind(FrameworkInfoFactory.class).toInstance(() -> Protos.FrameworkInfo.newBuilder()
                      .setUser("framework user")
                      .setName("test framework")
                      .build());
              bind(RescheduleCalculator.class).toInstance(new FakeRescheduleCalculator());
              bind(Clock.class).toInstance(new FakeClock());
              bind(ExecutorSettings.class).toInstance(TestExecutorSettings.THERMOS_EXECUTOR);
              bind(StatsProvider.class).toInstance(new FakeStatsProvider());
              bind(EventSink.class).toInstance(eventBus::post);
              bind(new TypeLiteral<BlockingQueue<Protos.TaskStatus>>() { })
                  .annotatedWith(TaskStatusHandlerImpl.StatusUpdateQueue.class)
                  .toInstance(new LinkedBlockingQueue<>());
              bind(new TypeLiteral<Integer>() { })
                  .annotatedWith(TaskStatusHandlerImpl.MaxBatchSize.class)
                  .toInstance(1000);
              bind(TaskStatusHandler.class).to(TaskStatusHandlerImpl.class);
              bind(TaskStatusHandlerImpl.class).in(Singleton.class);
              bind(IServerInfo.class).toInstance(IServerInfo.build(new ServerInfo("jmh", "")));
            }
          }
      );
  
      eventBus.register(injector.getInstance(ClusterStateImpl.class));
      scheduler = injector.getInstance(Scheduler.class);
      eventBus.register(this);
  
      statusHandler = injector.getInstance(TaskStatusHandlerImpl.class);
      statusHandler.startAsync();
    }
  
    @TearDown(Level.Trial)
    public void tearDown() {
      statusHandler.stopAsync();
    }
  
    /**
     * Runs before each iteration of the benchmark in order to vary the storage
     * latency across iterations, based on the latency parameter.
     */
    @Setup(Level.Iteration)
    public void setIterationLatency() {
      storage.setLatency(Amount.of(latencyMilliseconds, Time.MILLISECONDS));
    }
  
    /**
     * Runs before each invocation of the benchmark in order to store the tasks
     * that we will transition in the benchmark.
     */
    @Setup(Level.Invocation)
    public void createTasks() {
      tasks = new Tasks.Builder()
          .setScheduleStatus(ScheduleStatus.STARTING)
          .build(NUM_TASKS);
  
      storage.write(
          (NoResult.Quiet) storeProvider -> storeProvider.getUnsafeTaskStore().saveTasks(tasks));
  
      countDownLatch = new CountDownLatch(tasks.size());
    }
  
    @Subscribe
    public void taskChangedState(PubsubEvent.TaskStateChange stateChange) {
      countDownLatch.countDown();
    }
  
    @Benchmark
    public boolean runBenchmark() throws InterruptedException {
      for (String taskId : org.apache.aurora.scheduler.base.Tasks.ids(tasks)) {
        Protos.TaskStatus status = Protos.TaskStatus.newBuilder()
            .setState(Protos.TaskState.TASK_RUNNING)
            .setSource(Protos.TaskStatus.Source.SOURCE_EXECUTOR)
            .setMessage("message")
            .setTimestamp(1D)
            .setTaskId(Protos.TaskID.newBuilder().setValue(taskId).build())
            .build();
  
        scheduler.statusUpdate(new FakeSchedulerDriver(), ProtosConversion.convert(status));
      }
  
      // Wait for all task transitions to complete.
      countDownLatch.await();
  
      // Return an unguessable value.
      return System.currentTimeMillis() % 5 == 0;
    }
  }
  