content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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.bookkeeper.common.util;
  
  import static com.google.common.base.Preconditions.checkArgument;
  
  import com.google.common.util.concurrent.ForwardingExecutorService;
  import com.google.common.util.concurrent.ListenableFuture;
  import com.google.common.util.concurrent.SettableFuture;
  import com.google.common.util.concurrent.ThreadFactoryBuilder;
  import io.netty.util.concurrent.DefaultThreadFactory;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.ThreadFactory;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.stream.Collectors;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.common.util.affinity.CpuAffinity;
  import org.apache.bookkeeper.stats.Gauge;
  import org.apache.bookkeeper.stats.NullStatsLogger;
  import org.apache.bookkeeper.stats.OpStatsLogger;
  import org.apache.bookkeeper.stats.StatsLogger;
  import org.apache.bookkeeper.stats.ThreadRegistry;
  import org.apache.commons.lang3.StringUtils;
  import org.slf4j.MDC;
  
  /**
   * This class supports submitting tasks with an ordering key, so that tasks submitted
   * with the same key will always be executed in order, but tasks across
   * different keys can be unordered. This retains parallelism while retaining the
   * basic amount of ordering we want (e.g. , per ledger handle). Ordering is
   * achieved by hashing the key objects to threads by their {@link #hashCode()}
   * method.
   */
  @Slf4j
  public class OrderedExecutor implements ExecutorService {
      public static final int NO_TASK_LIMIT = -1;
      private static final int DEFAULT_MAX_ARRAY_QUEUE_SIZE = 10_000;
      protected static final long WARN_TIME_MICRO_SEC_DEFAULT = TimeUnit.SECONDS.toMicros(1);
  
      final String name;
      final ExecutorService[] threads;
      final long[] threadIds;
      final Random rand = new Random();
      final OpStatsLogger taskExecutionStats;
      final OpStatsLogger taskPendingStats;
      final boolean traceTaskExecution;
      final boolean preserveMdcForTaskExecution;
      final long warnTimeMicroSec;
      final int maxTasksInQueue;
      final boolean enableBusyWait;
      // we only want thread-scoped metrics on the server-side where it can be explicitly enabled
      final boolean enableThreadScopedMetrics;
  
      public static Builder newBuilder() {
          return new Builder();
      }
  
      /**
       * A builder class for an OrderedExecutor.
       */
      public static class Builder extends AbstractBuilder<OrderedExecutor> {
  
          @Override
          public OrderedExecutor build() {
              if (null == threadFactory) {
                  threadFactory = new DefaultThreadFactory("bookkeeper-ordered-safe-executor");
              }
              return new OrderedExecutor(name, numThreads, threadFactory, statsLogger,
                                             traceTaskExecution, preserveMdcForTaskExecution,
                                             warnTimeMicroSec, maxTasksInQueue, enableBusyWait,
                                             enableThreadScopedMetrics);
          }
      }
  
      /**
       * Abstract builder class to build {@link OrderedScheduler}.
       */
      public abstract static class AbstractBuilder<T extends OrderedExecutor> {
          protected String name = getClass().getSimpleName();
          protected int numThreads = Runtime.getRuntime().availableProcessors();
          protected ThreadFactory threadFactory = null;
          protected StatsLogger statsLogger = NullStatsLogger.INSTANCE;
          protected boolean traceTaskExecution = false;
          protected boolean preserveMdcForTaskExecution = false;
          protected long warnTimeMicroSec = WARN_TIME_MICRO_SEC_DEFAULT;
          protected int maxTasksInQueue = NO_TASK_LIMIT;
          protected boolean enableBusyWait = false;
          protected boolean enableThreadScopedMetrics = false;
  
          public AbstractBuilder<T> name(String name) {
              this.name = name;
              return this;
          }
  
          public AbstractBuilder<T> numThreads(int num) {
              this.numThreads = num;
              return this;
          }
  
          public AbstractBuilder<T> maxTasksInQueue(int num) {
              this.maxTasksInQueue = num;
              return this;
          }
  
          public AbstractBuilder<T> threadFactory(ThreadFactory threadFactory) {
              this.threadFactory = threadFactory;
              return this;
          }
  
          public AbstractBuilder<T> statsLogger(StatsLogger statsLogger) {
              this.statsLogger = statsLogger;
              return this;
          }
  
          public AbstractBuilder<T> traceTaskExecution(boolean enabled) {
              this.traceTaskExecution = enabled;
              return this;
          }
  
          public AbstractBuilder<T> preserveMdcForTaskExecution(boolean enabled) {
              this.preserveMdcForTaskExecution = enabled;
              return this;
          }
  
          public AbstractBuilder<T> traceTaskWarnTimeMicroSec(long warnTimeMicroSec) {
              this.warnTimeMicroSec = warnTimeMicroSec;
              return this;
          }
  
          public AbstractBuilder<T> enableBusyWait(boolean enableBusyWait) {
              this.enableBusyWait = enableBusyWait;
              return this;
          }
  
          public AbstractBuilder<T> enableThreadScopedMetrics(boolean enableThreadScopedMetrics) {
              this.enableThreadScopedMetrics = enableThreadScopedMetrics;
              return this;
          }
  
          @SuppressWarnings("unchecked")
          public T build() {
              if (null == threadFactory) {
                  threadFactory = new DefaultThreadFactory(name);
              }
              return (T) new OrderedExecutor(
                  name,
                  numThreads,
                  threadFactory,
                  statsLogger,
                  traceTaskExecution,
                  preserveMdcForTaskExecution,
                  warnTimeMicroSec,
                  maxTasksInQueue,
                  enableBusyWait,
                  enableThreadScopedMetrics);
          }
      }
  
      /**
       * Decorator class for a runnable that measure the execution time.
       */
      protected class TimedRunnable implements Runnable {
          final Runnable runnable;
          final long initNanos;
          final Class<?> runnableClass;
  
          TimedRunnable(Runnable runnable) {
              this.runnable = runnable;
              this.initNanos = MathUtils.nowInNano();
              this.runnableClass = runnable.getClass();
           }
  
          @Override
          public void run() {
              taskPendingStats.registerSuccessfulEvent(MathUtils.elapsedNanos(initNanos), TimeUnit.NANOSECONDS);
              long startNanos = MathUtils.nowInNano();
              try {
                  this.runnable.run();
              } finally {
                  long elapsedMicroSec = MathUtils.elapsedMicroSec(startNanos);
                  taskExecutionStats.registerSuccessfulEvent(elapsedMicroSec, TimeUnit.MICROSECONDS);
                  if (elapsedMicroSec >= warnTimeMicroSec) {
                      log.warn("Runnable {} took too long {} micros to execute.", runnableClass, elapsedMicroSec);
                  }
              }
          }
      }
  
      /**
       * Decorator class for a callable that measure the execution time.
       */
      protected class TimedCallable<T> implements Callable<T> {
          final Callable<T> callable;
          final long initNanos;
          final Class<?> callableClass;
  
          TimedCallable(Callable<T> callable) {
              this.callable = callable;
              this.initNanos = MathUtils.nowInNano();
              this.callableClass = callable.getClass();
          }
  
          @Override
          public T call() throws Exception {
              taskPendingStats.registerSuccessfulEvent(MathUtils.elapsedNanos(initNanos), TimeUnit.NANOSECONDS);
              long startNanos = MathUtils.nowInNano();
              try {
                  return this.callable.call();
              } finally {
                  long elapsedMicroSec = MathUtils.elapsedMicroSec(startNanos);
                  taskExecutionStats.registerSuccessfulEvent(elapsedMicroSec, TimeUnit.MICROSECONDS);
                  if (elapsedMicroSec >= warnTimeMicroSec) {
                      log.warn("Callable {} took too long {} micros to execute.", callableClass, elapsedMicroSec);
                  }
              }
          }
      }
  
      /**
       * Decorator class for a runnable that preserves MDC context.
       */
      static class ContextPreservingRunnable implements Runnable {
          private final Runnable runnable;
          private final Map<String, String> mdcContextMap;
  
          ContextPreservingRunnable(Runnable runnable) {
              this.runnable = runnable;
              this.mdcContextMap = MDC.getCopyOfContextMap();
          }
  
          @Override
          public void run() {
              MdcUtils.restoreContext(mdcContextMap);
              try {
                  runnable.run();
              } finally {
                  MDC.clear();
              }
          }
      }
  
      /**
       * Decorator class for a callable that preserves MDC context.
       */
      static class ContextPreservingCallable<T> implements Callable<T> {
          private final Callable<T> callable;
          private final Map<String, String> mdcContextMap;
  
          ContextPreservingCallable(Callable<T> callable) {
              this.callable = callable;
              this.mdcContextMap = MDC.getCopyOfContextMap();
          }
  
          @Override
          public T call() throws Exception {
              MdcUtils.restoreContext(mdcContextMap);
              try {
                  return callable.call();
              } finally {
                  MDC.clear();
              }
          }
      }
  
      protected ExecutorService createSingleThreadExecutor(ThreadFactory factory) {
          if (maxTasksInQueue > 0) {
              return new SingleThreadExecutor(factory, maxTasksInQueue, true);
          } else {
              return new SingleThreadExecutor(factory);
          }
      }
  
      protected ExecutorService getBoundedExecutor(ExecutorService executor) {
          checkArgument(executor instanceof ThreadPoolExecutor);
          return new BoundedExecutorService((ThreadPoolExecutor) executor, this.maxTasksInQueue);
      }
  
      protected ExecutorService addExecutorDecorators(ExecutorService executor) {
          return new ForwardingExecutorService() {
              @Override
              protected ExecutorService delegate() {
                  return executor;
              }
  
              @Override
              public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
                      throws InterruptedException {
                  return super.invokeAll(timedCallables(tasks));
              }
  
              @Override
              public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks,
                                                   long timeout, TimeUnit unit)
                      throws InterruptedException {
                  return super.invokeAll(timedCallables(tasks), timeout, unit);
              }
  
              @Override
              public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
                      throws InterruptedException, ExecutionException {
                  return super.invokeAny(timedCallables(tasks));
              }
  
              @Override
              public <T> T invokeAny(Collection<? extends Callable<T>> tasks,
                                     long timeout, TimeUnit unit)
                      throws InterruptedException, ExecutionException, TimeoutException {
                  return super.invokeAny(timedCallables(tasks), timeout, unit);
              }
  
              @Override
              public void execute(Runnable command) {
                  super.execute(timedRunnable(command));
              }
  
              @Override
              public <T> Future<T> submit(Callable<T> task) {
                  return super.submit(timedCallable(task));
              }
  
              @Override
              public Future<?> submit(Runnable task) {
                  return super.submit(timedRunnable(task));
              }
  
              @Override
              public <T> Future<T> submit(Runnable task, T result) {
                  return super.submit(timedRunnable(task), result);
              }
          };
      }
  
      /**
       * Constructs Safe executor.
       *
       * @param numThreads
       *            - number of threads
       * @param baseName
       *            - base name of executor threads
       * @param threadFactory
       *            - for constructing threads
       * @param statsLogger
       *            - for reporting executor stats
       * @param traceTaskExecution
       *            - should we stat task execution
       * @param preserveMdcForTaskExecution
       *            - should we preserve MDC for task execution
       * @param warnTimeMicroSec
       *            - log long task exec warning after this interval
       * @param maxTasksInQueue
       *            - maximum items allowed in a thread queue. -1 for no limit
       */
      protected OrderedExecutor(String baseName, int numThreads, ThreadFactory threadFactory,
                                  StatsLogger statsLogger, boolean traceTaskExecution,
                                  boolean preserveMdcForTaskExecution, long warnTimeMicroSec, int maxTasksInQueue,
                                  boolean enableBusyWait, boolean enableThreadScopedMetrics) {
          checkArgument(numThreads > 0);
          checkArgument(!StringUtils.isBlank(baseName));
  
          this.maxTasksInQueue = maxTasksInQueue;
          this.warnTimeMicroSec = warnTimeMicroSec;
          this.enableBusyWait = enableBusyWait;
          this.enableThreadScopedMetrics = enableThreadScopedMetrics;
          name = baseName;
          threads = new ExecutorService[numThreads];
          threadIds = new long[numThreads];
          for (int i = 0; i < numThreads; i++) {
              ExecutorService thread = createSingleThreadExecutor(
                      new ThreadFactoryBuilder().setNameFormat(name + "-" + getClass().getSimpleName() + "-" + i + "-%d")
                      .setThreadFactory(threadFactory).build());
              SingleThreadExecutor ste = null;
              if (thread instanceof SingleThreadExecutor) {
                  ste = (SingleThreadExecutor) thread;
              }
  
              if (traceTaskExecution || preserveMdcForTaskExecution) {
                  thread = addExecutorDecorators(thread);
              }
              threads[i] = thread;
  
              final int idx = i;
              try {
                  threads[idx].submit(() -> {
                      threadIds[idx] = Thread.currentThread().getId();
  
                      if (enableThreadScopedMetrics) {
                          ThreadRegistry.register(baseName, idx);
                      }
  
                      if (enableBusyWait) {
                          // Try to acquire 1 CPU core to the executor thread. If it fails we
                          // are just logging the error and continuing, falling back to
                          // non-isolated CPUs.
                          try {
                              CpuAffinity.acquireCore();
                          } catch (Throwable t) {
                              log.warn("Failed to acquire CPU core for thread {}: {}", Thread.currentThread().getName(),
                                      t.getMessage(), t);
                          }
                      }
                  }).get();
              } catch (InterruptedException e) {
                  Thread.currentThread().interrupt();
                  throw new RuntimeException("Couldn't start thread " + i, e);
              } catch (ExecutionException e) {
                  throw new RuntimeException("Couldn't start thread " + i, e);
              }
  
              if (ste != null) {
                  ste.registerMetrics(statsLogger);
              }
          }
  
          statsLogger.registerGauge(String.format("%s-threads", name), new Gauge<Number>() {
              @Override
              public Number getDefaultValue() {
                  return numThreads;
              }
  
              @Override
              public Number getSample() {
                  return numThreads;
              }
          });
  
          statsLogger.registerGauge(String.format("%s-max-queue-size", name), new Gauge<Number>() {
              @Override
              public Number getDefaultValue() {
                  return maxTasksInQueue;
              }
  
              @Override
              public Number getSample() {
                  return maxTasksInQueue;
              }
          });
  
          if (enableThreadScopedMetrics) {
              this.taskExecutionStats = statsLogger.scope(name).getThreadScopedOpStatsLogger("task_execution");
              this.taskPendingStats = statsLogger.scope(name).getThreadScopedOpStatsLogger("task_queued");
          } else {
              this.taskExecutionStats = statsLogger.scope(name).getOpStatsLogger("task_execution");
              this.taskPendingStats = statsLogger.scope(name).getOpStatsLogger("task_queued");
          }
          this.traceTaskExecution = traceTaskExecution;
          this.preserveMdcForTaskExecution = preserveMdcForTaskExecution;
      }
  
      /**
       * Flag describing executor's expectation in regards of MDC.
       * All tasks submitted through executor's submit/execute methods will automatically respect this.
       *
       * @return true if runnable/callable is expected to preserve MDC, false otherwise.
       */
      public boolean preserveMdc() {
          return preserveMdcForTaskExecution;
      }
  
      /**
       * Schedules a one time action to execute with an ordering guarantee on the key.
       * @param orderingKey
       * @param r
       */
      public void executeOrdered(Object orderingKey, Runnable r) {
          chooseThread(orderingKey).execute(r);
      }
  
      /**
       * Schedules a one time action to execute with an ordering guarantee on the key.
       * @param orderingKey
       * @param r
       */
      public void executeOrdered(long orderingKey, Runnable r) {
          chooseThread(orderingKey).execute(r);
      }
  
      /**
       * Schedules a one time action to execute with an ordering guarantee on the key.
       * @param orderingKey
       * @param r
       */
      public void executeOrdered(int orderingKey, Runnable r) {
          chooseThread(orderingKey).execute(r);
      }
  
      public <T> ListenableFuture<T> submitOrdered(long orderingKey, Callable<T> task) {
          SettableFuture<T> future = SettableFuture.create();
          executeOrdered(orderingKey, () -> {
              try {
                  T result = task.call();
                  future.set(result);
              } catch (Throwable t) {
                  future.setException(t);
              }
          });
  
          return future;
      }
  
  
      public long getThreadID(long orderingKey) {
          // skip hashcode generation in this special case
          if (threadIds.length == 1) {
              return threadIds[0];
          }
  
          return threadIds[chooseThreadIdx(orderingKey, threads.length)];
      }
  
      public ExecutorService chooseThread() {
          // skip random # generation in this special case
          if (threads.length == 1) {
              return threads[0];
          }
  
          return threads[rand.nextInt(threads.length)];
      }
  
      public ExecutorService chooseThread(Object orderingKey) {
          // skip hashcode generation in this special case
          if (threads.length == 1) {
              return threads[0];
          }
  
          if (null == orderingKey) {
              return threads[rand.nextInt(threads.length)];
          } else {
              return threads[chooseThreadIdx(orderingKey.hashCode(), threads.length)];
          }
      }
  
      /**
       * skip hashcode generation in this special case.
       *
       * @param orderingKey long ordering key
       * @return the thread for executing this order key
       */
      public ExecutorService chooseThread(long orderingKey) {
          if (threads.length == 1) {
              return threads[0];
          }
  
          return threads[chooseThreadIdx(orderingKey, threads.length)];
      }
  
      protected static int chooseThreadIdx(long orderingKey, int numThreads) {
          return MathUtils.signSafeMod(orderingKey >>> 1, numThreads);
      }
  
      protected Runnable timedRunnable(Runnable r) {
          final Runnable runMe = traceTaskExecution ? new TimedRunnable(r) : r;
          return preserveMdcForTaskExecution ? new ContextPreservingRunnable(runMe) : runMe;
      }
  
      protected <T> Callable<T> timedCallable(Callable<T> c) {
          final Callable<T> callMe = traceTaskExecution ? new TimedCallable<>(c) : c;
          return preserveMdcForTaskExecution ? new ContextPreservingCallable<>(callMe) : callMe;
      }
  
      protected <T> Collection<? extends Callable<T>> timedCallables(Collection<? extends Callable<T>> tasks) {
          if (traceTaskExecution || preserveMdcForTaskExecution) {
              return tasks.stream()
                      .map(this::timedCallable)
                      .collect(Collectors.toList());
          }
          return tasks;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> Future<T> submit(Callable<T> task) {
          return chooseThread().submit(task);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> Future<T> submit(Runnable task, T result) {
          return chooseThread().submit(task, result);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public Future<?> submit(Runnable task) {
          return chooseThread().submit(task);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
          throws InterruptedException {
          return chooseThread().invokeAll(tasks);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks,
                                           long timeout,
                                           TimeUnit unit)
          throws InterruptedException {
          return chooseThread().invokeAll(tasks, timeout, unit);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
          throws InterruptedException, ExecutionException {
          return chooseThread().invokeAny(tasks);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
          throws InterruptedException, ExecutionException, TimeoutException {
          return chooseThread().invokeAny(tasks, timeout, unit);
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public void execute(Runnable command) {
          chooseThread().execute(command);
      }
  
  
      /**
       * {@inheritDoc}
       */
      @Override
      public void shutdown() {
          for (int i = 0; i < threads.length; i++) {
              threads[i].shutdown();
          }
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public List<Runnable> shutdownNow() {
          List<Runnable> runnables = new ArrayList<Runnable>();
          for (ExecutorService executor : threads) {
              runnables.addAll(executor.shutdownNow());
          }
          return runnables;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public boolean isShutdown() {
          for (ExecutorService executor : threads) {
              if (!executor.isShutdown()) {
                  return false;
              }
          }
          return true;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
          boolean ret = true;
          for (int i = 0; i < threads.length; i++) {
              ret = ret && threads[i].awaitTermination(timeout, unit);
          }
          return ret;
      }
  
      /**
       * {@inheritDoc}
       */
      @Override
      public boolean isTerminated() {
          for (ExecutorService executor : threads) {
              if (!executor.isTerminated()) {
                  return false;
              }
          }
          return true;
      }
  
      /**
       * Force threads shutdown (cancel active requests) after specified delay,
       * to be used after shutdown() rejects new requests.
       */
      public void forceShutdown(long timeout, TimeUnit unit) {
          for (int i = 0; i < threads.length; i++) {
              try {
                  if (!threads[i].awaitTermination(timeout, unit)) {
                      threads[i].shutdownNow();
                  }
              } catch (InterruptedException exception) {
                  threads[i].shutdownNow();
                  Thread.currentThread().interrupt();
              }
          }
      }
  
  }
  