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;
  
  import java.util.LinkedList;
  import java.util.List;
  import java.util.Optional;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.ScheduledExecutorService;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.AtomicLong;
  
  import javax.inject.Inject;
  
  import com.google.common.util.concurrent.AbstractExecutionThreadService;
  
  import org.apache.aurora.common.stats.SlidingStats;
  import org.apache.aurora.common.stats.StatsProvider;
  import org.apache.aurora.common.util.BackoffStrategy;
  import org.apache.aurora.scheduler.base.AsyncUtil;
  import org.apache.aurora.scheduler.storage.Storage;
  import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  import static com.google.common.base.Preconditions.checkState;
  
  /**
   * Generic helper that allows bundling multiple work items into a single {@link Storage}
   * transaction aiming to reduce the write lock contention.
   *
   * @param <T> Expected result type.
   */
  public class BatchWorker<T> extends AbstractExecutionThreadService {
    /**
     * Empty result placeholder.
     */
    public interface NoResult { }
  
    /**
     * Convenience wrapper for a non-repeatable no value work {@link Result}.
     */
    public static final NoResult NO_RESULT = new NoResult() { };
  
    private static final Logger LOG = LoggerFactory.getLogger(BatchWorker.class);
    private final Storage storage;
    private final int maxBatchSize;
    private final SlidingStats batchUnlocked;
    private final SlidingStats batchLocked;
    private final BlockingQueue<WorkItem<T>> workQueue = new LinkedBlockingQueue<>();
    private final ScheduledExecutorService scheduledExecutor;
    private final AtomicInteger lastBatchSize = new AtomicInteger(0);
    private final AtomicLong itemsProcessed;
    private final AtomicLong batchesProcessed;
  
    /**
     * Wraps result returned by the {@link RepeatableWork} item.
     *
     * @param <T> Expected result type.
     */
    public static class Result<T> {
      private final boolean isCompleted;
      private final T value;
  
      /**
       * Initializes a {@link Result} instance with {@code isCompleted} and {@code value}.
       * <p>
       * The {@code isCompleted} may be set to {@code False} for a {@link RepeatableWork} that has
       * not finished yet. Otherwise, it must be set to {@code True}.
       *
       * @param isCompleted Flag indicating if the {@link RepeatableWork} has completed.
       * @param value result value.
       */
      public Result(boolean isCompleted, T value) {
        this.isCompleted = isCompleted;
        this.value = value;
      }
    }
  
    /**
     * Encapsulates a potentially repeatable operation.
     */
    public interface RepeatableWork<T> {
      /**
       * Abstracts a unit of repeatable (i.e.: "repeat until completed") work.
       * <p>
       * The work unit may be repeated as instructed by the {@link Result}.
       *
       * @param storeProvider {@link MutableStoreProvider} instance.
       * @return {@link Result}
       */
      Result<T> apply(MutableStoreProvider storeProvider);
    }
  
    /**
     * Encapsulates a non-repeatable operation.
     */
    public interface Work<T> extends RepeatableWork<T> {
      @Override
      default Result<T> apply(MutableStoreProvider storeProvider) {
        T value = execute(storeProvider);
        return new Result<>(true, value);
      }
  
      /**
       * Abstracts a unit of non-repeatable (i.e.: "run exactly once") work.
       *
       * @param storeProvider {@link MutableStoreProvider} instance.
       * @return result value.
       */
      T execute(MutableStoreProvider storeProvider);
    }
  
    @Inject
    protected BatchWorker(
        Storage storage,
        StatsProvider statsProvider,
        int maxBatchSize) {
  
      this.storage = requireNonNull(storage);
      this.maxBatchSize = maxBatchSize;
  
      scheduledExecutor = AsyncUtil.singleThreadLoggingScheduledExecutor(serviceName() + "-%d", LOG);
      statsProvider.makeGauge(serviceName() + "_queue_size", () -> workQueue.size());
      statsProvider.makeGauge(
          serviceName() + "_last_processed_batch_size",
          () -> lastBatchSize.intValue());
      batchUnlocked = new SlidingStats(serviceName() + "_batch_unlocked", "nanos");
      batchLocked = new SlidingStats(serviceName() + "_batch_locked", "nanos");
      itemsProcessed = statsProvider.makeCounter(serviceName() + "_items_processed");
      batchesProcessed = statsProvider.makeCounter(serviceName() + "_batches_processed");
    }
  
    /**
     * Executes a non-repeatable {@link Work} and returns {@link CompletableFuture} to wait on.
     *
     * @param work A non-repeatable {@link Work} to execute.
     * @return {@link CompletableFuture} to wait on.
     */
    public CompletableFuture<T> execute(Work<T> work) {
      CompletableFuture<T> result = new CompletableFuture<>();
      workQueue.add(new WorkItem<>(
          work,
          result,
          Optional.empty(),
          Optional.empty()));
  
      return result;
    }
  
    /**
     * Executes a {@link RepeatableWork} until it completes and returns {@link CompletableFuture}
     * to wait on.
     *
     * @param backoffStrategy A {@link BackoffStrategy} instance to backoff subsequent runs.
     * @param work A {@link RepeatableWork} to execute.
     */
    public CompletableFuture<T> executeWithReplay(
        BackoffStrategy backoffStrategy,
        RepeatableWork<T> work) {
  
      CompletableFuture<T> result = new CompletableFuture<>();
      workQueue.add(new WorkItem<>(
          work,
          result,
          Optional.of(backoffStrategy),
          Optional.of(0L)));
  
      return result;
    }
  
    @Override
    protected void run() throws Exception {
      while (isRunning()) {
        List<WorkItem<T>> batch = new LinkedList<>();
  
        // Make the loop responsive to shutdown under light load by using
        // a short non-configurable timeout in poll().
        Optional<WorkItem<T>> head = Optional.ofNullable(workQueue.poll(3, TimeUnit.SECONDS));
        if (head.isPresent()) {
          workQueue.add(head.get());
          workQueue.drainTo(batch, maxBatchSize - batch.size());
          processBatch(batch);
        }
      }
    }
  
    private void processBatch(List<WorkItem<T>> batch) {
      if (!batch.isEmpty()) {
        long unlockedStart = System.nanoTime();
        storage.write((Storage.MutateWork.NoResult.Quiet) storeProvider -> {
          long lockedStart = System.nanoTime();
          for (WorkItem<T> item : batch) {
            Result<T> itemResult = item.work.apply(storeProvider);
            if (itemResult.isCompleted) {
              item.result.complete(itemResult.value);
            } else {
              // Work not finished yet - re-queue for a followup later.
              long backoffMsec = backoffFor(item);
              scheduledExecutor.schedule(
                  () -> workQueue.add(new WorkItem<>(
                      item.work,
                      item.result,
                      item.backoffStrategy,
                      Optional.of(backoffMsec))),
                  backoffMsec,
                  TimeUnit.MILLISECONDS);
            }
          }
          batchLocked.accumulate(System.nanoTime() - lockedStart);
        });
        batchUnlocked.accumulate(System.nanoTime() - unlockedStart);
        batchesProcessed.incrementAndGet();
        lastBatchSize.set(batch.size());
        itemsProcessed.addAndGet(batch.size());
      }
    }
  
    private long backoffFor(WorkItem<T> item) {
      checkState(item.backoffStrategy.isPresent());
      checkState(item.lastBackoffMsec.isPresent());
      return item.backoffStrategy.get().calculateBackoffMs(item.lastBackoffMsec.get());
    }
  
    private class WorkItem<V> {
      private final RepeatableWork<V> work;
      private final CompletableFuture<T> result;
      private final Optional<BackoffStrategy> backoffStrategy;
      private final Optional<Long> lastBackoffMsec;
  
      WorkItem(
          RepeatableWork<V> work,
          CompletableFuture<T> result,
          Optional<BackoffStrategy> backoffStrategy,
          Optional<Long> lastBackoffMsec) {
  
        this.work = work;
        this.result = result;
        this.backoffStrategy = backoffStrategy;
        this.lastBackoffMsec = lastBackoffMsec;
      }
    }
  }
  