content: /*
   * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
   * (the "License"). You may not use this work except in compliance with the License, which is
   * available at www.apache.org/licenses/LICENSE-2.0
   *
   * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
   * either express or implied, as more fully set forth in the License.
   *
   * See the NOTICE file distributed with this work for information regarding copyright ownership.
   */
  
  package alluxio.worker.grpc;
  
  import alluxio.Constants;
  import alluxio.conf.Configuration;
  import alluxio.conf.PropertyKey;
  import alluxio.metrics.MetricKey;
  import alluxio.metrics.MetricsSystem;
  import alluxio.security.User;
  import alluxio.security.authentication.AuthenticatedClientUser;
  import alluxio.util.ThreadFactoryUtils;
  import alluxio.util.executor.UniqueBlockingQueue;
  import alluxio.worker.block.DefaultBlockWorker;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Collection;
  import java.util.List;
  import java.util.concurrent.AbstractExecutorService;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Future;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.SynchronousQueue;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * Executors for gRPC block server.
   */
  @ThreadSafe
  public final class GrpcExecutors {
    private static final Logger LOG = LoggerFactory.getLogger(GrpcExecutors.class);
    private static final long THREAD_STOP_MS = Constants.SECOND_MS * 10;
    private static final int THREADS_MIN = 4;
  
    private static final ThreadPoolExecutor CACHE_MANAGER_THREAD_POOL_EXECUTOR =
        new ThreadPoolExecutor(THREADS_MIN,
            Configuration.getInt(PropertyKey.WORKER_NETWORK_ASYNC_CACHE_MANAGER_THREADS_MAX),
            THREAD_STOP_MS, TimeUnit.MILLISECONDS, new UniqueBlockingQueue<>(
            Configuration.getInt(PropertyKey.WORKER_NETWORK_ASYNC_CACHE_MANAGER_QUEUE_MAX)),
            ThreadFactoryUtils.build("CacheManagerExecutor-%d", true));
    // Async caching is an optimization internal to Alluxio, which can be aborted any time
    public static final ExecutorService CACHE_MANAGER_EXECUTOR =
        new ImpersonateThreadPoolExecutor(CACHE_MANAGER_THREAD_POOL_EXECUTOR, false);
  
    // Used by BlockWorkerClientServiceHandler.readBlock() by DataReader threads,
    // where each DataReader reads a block content for reply.
    // The thread pool queue is always empty.
    private static final ThreadPoolExecutor BLOCK_READER_THREAD_POOL_EXECUTOR =
        new ThreadPoolExecutor(THREADS_MIN, Configuration.getInt(
            PropertyKey.WORKER_NETWORK_BLOCK_READER_THREADS_MAX), THREAD_STOP_MS,
            TimeUnit.MILLISECONDS, new SynchronousQueue<>(),
            ThreadFactoryUtils.build("BlockDataReaderExecutor-%d", true));
    public static final ExecutorService BLOCK_READER_EXECUTOR =
        new ImpersonateThreadPoolExecutor(BLOCK_READER_THREAD_POOL_EXECUTOR, true);
  
    // Used for replying data to the client in BlockReadHandler.
    // The thread pool has a small queue of a constant size.
    private static final ThreadPoolExecutor BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR =
        new ThreadPoolExecutor(THREADS_MIN,
            Configuration.getInt(PropertyKey.WORKER_NETWORK_BLOCK_READER_THREADS_MAX),
            THREAD_STOP_MS, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(32),
            ThreadFactoryUtils.build("BlockDataReaderSerializedExecutor-%d", true),
            new ThreadPoolExecutor.CallerRunsPolicy());
    public static final ExecutorService BLOCK_READER_SERIALIZED_RUNNER_EXECUTOR =
        new ImpersonateThreadPoolExecutor(BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR, true);
  
    // Used for writing blocks. The queue is always empty.
    private static final ThreadPoolExecutor BLOCK_WRITE_THREAD_POOL_EXECUTOR =
        new ThreadPoolExecutor(THREADS_MIN, Configuration.getInt(
            PropertyKey.WORKER_NETWORK_BLOCK_WRITER_THREADS_MAX), THREAD_STOP_MS,
            TimeUnit.MILLISECONDS, new SynchronousQueue<>(),
            ThreadFactoryUtils.build("BlockDataWriterExecutor-%d", true));
    public static final ExecutorService BLOCK_WRITER_EXECUTOR =
            new ImpersonateThreadPoolExecutor(BLOCK_WRITE_THREAD_POOL_EXECUTOR, true);
  
    static {
      MetricsSystem.registerCachedGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_CACHE_MANAGER_THREAD_ACTIVE_COUNT.getName()),
          CACHE_MANAGER_THREAD_POOL_EXECUTOR::getActiveCount, 5, TimeUnit.SECONDS);
      MetricsSystem.registerCachedGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_CACHE_MANAGER_THREAD_CURRENT_COUNT.getName()),
          CACHE_MANAGER_THREAD_POOL_EXECUTOR::getPoolSize, 5, TimeUnit.SECONDS);
      MetricsSystem.registerCachedGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_CACHE_MANAGER_THREAD_QUEUE_WAITING_TASK_COUNT.getName()),
          CACHE_MANAGER_THREAD_POOL_EXECUTOR.getQueue()::size, 5, TimeUnit.SECONDS);
      // This value is not updated after the process starts,
      // so it can be cached for a long time to reduce the number of queries
      MetricsSystem.registerCachedGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_CACHE_MANAGER_THREAD_MAX_COUNT.getName()),
          CACHE_MANAGER_THREAD_POOL_EXECUTOR::getMaximumPoolSize, 30, TimeUnit.MINUTES);
      MetricsSystem.registerCachedGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_CACHE_MANAGER_COMPLETED_TASK_COUNT.getName()),
          CACHE_MANAGER_THREAD_POOL_EXECUTOR::getCompletedTaskCount, 5, TimeUnit.SECONDS);
  
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_READER_THREAD_ACTIVE_COUNT.getName()),
          BLOCK_READER_THREAD_POOL_EXECUTOR::getActiveCount);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_READER_THREAD_CURRENT_COUNT.getName()),
          BLOCK_READER_THREAD_POOL_EXECUTOR::getPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_READER_THREAD_MAX_COUNT.getName()),
          BLOCK_READER_THREAD_POOL_EXECUTOR::getMaximumPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_READER_COMPLETED_TASK_COUNT.getName()),
          BLOCK_READER_THREAD_POOL_EXECUTOR::getCompletedTaskCount);
  
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_SERIALIZED_THREAD_ACTIVE_COUNT.getName()),
          BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR::getActiveCount);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_SERIALIZED_THREAD_CURRENT_COUNT.getName()),
          BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR::getPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_SERIALIZED_THREAD_MAX_COUNT.getName()),
          BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR::getMaximumPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_SERIALIZED_COMPLETED_TASK_COUNT.getName()),
          BLOCK_SERIALIZED_THREAD_POOL_EXECUTOR::getCompletedTaskCount);
  
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_WRITER_THREAD_ACTIVE_COUNT.getName()),
          BLOCK_WRITE_THREAD_POOL_EXECUTOR::getActiveCount);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_WRITER_THREAD_CURRENT_COUNT.getName()),
          BLOCK_WRITE_THREAD_POOL_EXECUTOR::getPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_WRITER_THREAD_MAX_COUNT.getName()),
          BLOCK_WRITE_THREAD_POOL_EXECUTOR::getMaximumPoolSize);
      MetricsSystem.registerGaugeIfAbsent(MetricsSystem.getMetricName(
          MetricKey.WORKER_BLOCK_WRITER_COMPLETED_TASK_COUNT.getName()),
          BLOCK_WRITE_THREAD_POOL_EXECUTOR::getCompletedTaskCount);
    }
  
    /**
     * Private constructor.
     */
    private GrpcExecutors() {}
  
    /**
     * This executor passes impersonation information to the real worker thread.
     * The proxy user is tracked by {@link AuthenticatedClientUser#sUserThreadLocal}.
     * This executor delegates operations to the underlying executor while setting the
     * ThreadLocal context for execution.
     * */
    private static class ImpersonateThreadPoolExecutor extends AbstractExecutorService {
      private final ExecutorService mDelegate;
      private final boolean mTracked;
  
      public ImpersonateThreadPoolExecutor(ExecutorService service, boolean tracked) {
        mDelegate = service;
        mTracked = tracked;
      }
  
      @Override
      public void execute(final Runnable command) {
        // If there's no impersonation, proxyUser is just null
        User proxyUser = AuthenticatedClientUser.getOrNull();
        mDelegate.execute(() -> {
          if (mTracked) {
            DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.inc();
          }
          try {
  //          SleepUtils.sleepMs(1000);
            AuthenticatedClientUser.set(proxyUser);
            command.run();
          } finally {
            if (mTracked) {
              DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.dec();
            }
            AuthenticatedClientUser.remove();
          }
        });
      }
  
      @Override
      public <T> Future<T> submit(Callable<T> task) {
        // If there's no impersonation, proxyUser is just null
        User proxyUser = AuthenticatedClientUser.getOrNull();
        return mDelegate.submit(() -> {
          if (mTracked) {
            DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.inc();
          }
          try {
  //          SleepUtils.sleepMs(1000);
            AuthenticatedClientUser.set(proxyUser);
            return task.call();
          } finally {
            if (mTracked) {
              DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.dec();
            }
            AuthenticatedClientUser.remove();
          }
        });
      }
  
      @Override
      public <T> Future<T> submit(Runnable task, T result) {
        // If there's no impersonation, proxyUser is just null
        User proxyUser = AuthenticatedClientUser.getOrNull();
        return mDelegate.submit(() -> {
          if (mTracked) {
            DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.inc();
          }
          try {
  //          SleepUtils.sleepMs(1000);
            AuthenticatedClientUser.set(proxyUser);
            task.run();
          } finally {
            if (mTracked) {
              DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.dec();
            }
            AuthenticatedClientUser.remove();
          }
        }, result);
      }
  
      @Override
      public Future<?> submit(Runnable task) {
        // If there's no impersonation, proxyUser is just null
        User proxyUser = AuthenticatedClientUser.getOrNull();
        return mDelegate.submit(() -> {
          if (mTracked) {
            DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.inc();
          }
          try {
  //          SleepUtils.sleepMs(1000);
            AuthenticatedClientUser.set(proxyUser);
            task.run();
          } finally {
            if (mTracked) {
              DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.dec();
            }
            AuthenticatedClientUser.remove();
          }
        });
      }
  
      @Override
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
          throws InterruptedException {
        // Not used. Also the active counter is hard, so we do not support it.
        throw new UnsupportedOperationException("invokeAll(Collection) is not supported");
      }
  
      @Override
      public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout,
          TimeUnit unit) throws InterruptedException {
        // Not used. Also the active counter is hard, so we do not support it.
        throw new UnsupportedOperationException(
            "invokeAll(Collection,long,TimeUnit) is not supported");
      }
  
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks) {
        // Not used. Also the active counter is hard, so we do not support it.
        throw new UnsupportedOperationException("invokeAny(Callable) is not supported");
      }
  
      @Override
      public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) {
        // Not used. Also the active counter is hard, so we do not support it.
        throw new UnsupportedOperationException(
            "invokeAny(Callable,long,TimeUnit) is not supported");
      }
  
      @Override
      public void shutdown() {
        long operationCount = DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.getCount();
        if (operationCount > 0) {
          LOG.warn("{} operations have not completed at shutdown()", operationCount);
        }
        mDelegate.shutdown();
      }
  
      @Override
      public List<Runnable> shutdownNow() {
        long operationCount = DefaultBlockWorker.Metrics.WORKER_ACTIVE_OPERATIONS.getCount();
        if (operationCount > 0) {
          LOG.warn("{} operations have not completed at shutdownNow()", operationCount);
        }
        return mDelegate.shutdownNow();
      }
  
      @Override
      public boolean isShutdown() {
        return mDelegate.isShutdown();
      }
  
      @Override
      public boolean isTerminated() {
        return mDelegate.isTerminated();
      }
  
      @Override
      public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return mDelegate.awaitTermination(timeout, unit);
      }
    }
  }
  