text: /*
   *     Copyright 2022-Present Couchbase, Inc.
   *
   *   Use of this software is governed by the Business Source License included
   *   in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
   *   in that file, in accordance with the Business Source License, use of this
   *   software will be governed by the Apache License, Version 2.0, included in
   *   the file licenses/APL2.txt.
   */
  #pragma once
  
  #include "globaltask.h"
  
  #include <folly/concurrency/UnboundedQueue.h>
  #include <folly/executors/CPUThreadPoolExecutor.h>
  
  /**
   * Custom folly CPU thread pool executor which puts work in a separate queue
   * that we control so that we can remove tasks as and when a bucket goes away
   * without having to actually run them (seen to noticeably slow down bucket
   * deinitialization which can case rebalance failures).
   */
  class CancellableCPUExecutor {
  public:
      CancellableCPUExecutor(size_t numThreads,
                             std::shared_ptr<folly::ThreadFactory> threadFactory)
          : cpuPool(numThreads, std::move(threadFactory)) {
      }
  
      size_t numThreads() const {
          return cpuPool.numThreads();
      }
  
      void setNumThreads(size_t numThreads) {
          cpuPool.setNumThreads(numThreads);
      }
  
      folly::ThreadPoolExecutor::PoolStats getPoolStats() const {
          return cpuPool.getPoolStats();
      }
  
      size_t getPendingTaskCount() const {
          return cpuPool.getPendingTaskCount();
      }
  
      size_t getTaskQueueSize() const {
          return cpuPool.getTaskQueueSize();
      }
  
      void join() {
          cpuPool.join();
      }
  
      /**
       * Add a task to run
       *
       * @param task which includes the taskable and taskId that we need in this
       *        class. Task can be null if the work is not strictly associated
       *        with a GlobalTask (i.e. a reset of a task pointer).
       * @param func the function to run (which runs the task)
       */
      void add(GlobalTask* task, folly::Func func);
  
      /**
       * Remove tasks for the given taskable from our work queue. Should always
       * be run on the SchedulerPool (EventBase) to ensure that tasks are not
       * scheduled on this pool while we run.
       *
       * @param taskable to remove tasks for
       * @return vector of GlobalTask* for this taskable
       */
      std::vector<GlobalTask*> removeTasksForTaskable(const Taskable& taskable);
  
  private:
      /**
       * Element of our queue
       */
      struct QueueElem {
          QueueElem(GlobalTask* task, folly::Func func)
              : task(task), func(std::move(func)) {
          }
  
          bool isInternalExecutorTask() const {
              return !task;
          }
  
          // Taking a GlobalTask* here because it contains both the Taskable
          // reference and the taskId which we need to cancel tasks. We take a
          // GlobalTask* over an ExTask to avoid shared_ptr promotions for the
          // sake of efficiency. We require this to be nullable as some tasks
          // (the reset of task ptrs on the cpuPool) are not associated with a
          // GlobalTask.
          GlobalTask* task;
          folly::Func func;
      };
  
      /**
       * Queue of tasks to run. Uses the same queue type as
       * folly::CPUThreadPoolExecutor. The task here contians the actual work we
       * want to run.
       */
      folly::UMPMCQueue<QueueElem, false, 6> tasks;
  
      /**
       * Underlying CPUThreadPoolExecutor which runs "notification" tasks to run a
       * task in our tasks queue.
       */
      folly::CPUThreadPoolExecutor cpuPool;
  };
  