base/task/thread_pool/task_source.h

// Copyright 2019 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef BASE_TASK_THREAD_POOL_TASK_SOURCE_H_
#define BASE_TASK_THREAD_POOL_TASK_SOURCE_H_

#include <stddef.h>

#include "base/base_export.h"
#include "base/containers/intrusive_heap.h"
#include "base/dcheck_is_on.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/raw_ptr_exclusion.h"
#include "base/memory/ref_counted.h"
#include "base/sequence_token.h"
#include "base/task/common/checked_lock.h"
#include "base/task/task_traits.h"
#include "base/task/thread_pool/task.h"
#include "base/task/thread_pool/task_source_sort_key.h"
#include "base/threading/sequence_local_storage_map.h"
#include "base/time/time.h"

namespace base {
namespace internal {

class TaskTracker;

enum class TaskSourceExecutionMode {
  kParallel,
  kSequenced,
  kSingleThread,
  kJob,
  kMax = kJob,
};

struct BASE_EXPORT ExecutionEnvironment {
  SequenceToken token;
  raw_ptr<SequenceLocalStorageMap> sequence_local_storage;
};

// A TaskSource is a virtual class that provides a series of Tasks that must be
// executed immediately or in the future.
//
// When a task source has delayed tasks but no immediate tasks, the scheduler
// must call OnBecomeReady() after HasReadyTasks(now) == true, which is
// guaranteed once now >= GetDelayedSortKey().
//
// A task source is registered when it's ready to be added to the immediate
// queue. A task source is ready to be queued when either:
// 1- It has new tasks that can run concurrently as a result of external
//    operations, e.g. posting a new immediate task to an empty Sequence or
//    increasing max concurrency of a JobTaskSource;
// 2- A worker finished running a task from it and both DidProcessTask() and
//    WillReEnqueue() returned true; or
// 3- A worker is about to run a task from it and WillRunTask() returned
//    kAllowedNotSaturated.
// 4- A delayed task became ready and OnBecomeReady() returns true.
//
// A worker may perform the following sequence of operations on a
// RegisteredTaskSource after obtaining it from the queue:
// 1- Check whether a task can run with WillRunTask() (and register/enqueue the
//    task source again if not saturated).
// 2- (optional) Iff (1) determined that a task can run, access the next task
//    with TakeTask().
// 3- (optional) Execute the task.
// 4- Inform the task source that a task was processed with DidProcessTask(),
//    and re-enqueue the task source iff requested. The task source is ready to
//    run immediately iff WillReEnqueue() returns true.
// When a task source is registered multiple times, many overlapping chains of
// operations may run concurrently, as permitted by WillRunTask(). This allows
// tasks from the same task source to run in parallel.
// However, the following invariants are kept:
// - The number of workers concurrently running tasks never goes over the
//   intended concurrency.
// - If the task source has more tasks that can run concurrently, it must be
//   queued.
//
// Note: there is a known refcounted-ownership cycle in the ThreadPool
// architecture: TaskSource -> TaskRunner -> TaskSource -> ... This is okay so
// long as the other owners of TaskSource (PriorityQueue and WorkerThread in
// alternation and ThreadGroupImpl::WorkerThreadDelegateImpl::GetWork()
// temporarily) keep running it (and taking Tasks from it as a result). A
// dangling reference cycle would only occur should they release their reference
// to it while it's not empty. In other words, it is only correct for them to
// release it when DidProcessTask() returns false.
//
// This class is thread-safe.
class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
 public:
  // Indicates whether WillRunTask() allows TakeTask() to be called on a
  // RegisteredTaskSource.
  enum class RunStatus {
    // TakeTask() cannot be called.
    kDisallowed,
    // TakeTask() may called, and the TaskSource has not reached its maximum
    // concurrency (i.e. the TaskSource still needs to be queued).
    kAllowedNotSaturated,
    // TakeTask() may called, and the TaskSource has reached its maximum
    // concurrency (i.e. the TaskSource no longer needs to be queued).
    kAllowedSaturated,
  };

  // A Transaction can perform multiple operations atomically on a
  // TaskSource. While a Transaction is alive, it is guaranteed that nothing
  // else will access the TaskSource; the TaskSource's lock is held for the
  // lifetime of the Transaction. No Transaction must be held when ~TaskSource()
  // is called.
  class BASE_EXPORT Transaction {
   public:
    Transaction(Transaction&& other);
    Transaction(const Transaction&) = delete;
    Transaction& operator=(const Transaction&) = delete;
    ~Transaction();

    operator bool() const { return !!task_source_; }

    // Sets TaskSource priority to |priority|.
    void UpdatePriority(TaskPriority priority);

    // Returns the traits of all Tasks in the TaskSource.
    TaskTraits traits() const { return task_source_->traits_; }

    TaskSource* task_source() const { return task_source_; }

    void Release();

   protected:
    explicit Transaction(TaskSource* task_source);

   private:
    friend class TaskSource;

    // This field is not a raw_ptr<> because it was filtered by the rewriter
    // for: #union
    RAW_PTR_EXCLUSION TaskSource* task_source_;
  };

  // |traits| is metadata that applies to all Tasks in the TaskSource.
  // |task_runner| is a reference to the TaskRunner feeding this TaskSource.
  // |task_runner| can be nullptr only for tasks with no TaskRunner, in which
  // case |execution_mode| must be kParallel. Otherwise, |execution_mode| is the
  // execution mode of |task_runner|.
  TaskSource(const TaskTraits& traits,
             TaskRunner* task_runner,
             TaskSourceExecutionMode execution_mode);
  TaskSource(const TaskSource&) = delete;
  TaskSource& operator=(const TaskSource&) = delete;

  // Begins a Transaction. This method cannot be called on a thread which has an
  // active TaskSource::Transaction.
  [[nodiscard]] Transaction BeginTransaction();

  virtual ExecutionEnvironment GetExecutionEnvironment() = 0;

  // Thread-safe but the returned value may immediately be obsolete. As such
  // this should only be used as a best-effort guess of how many more workers
  // are needed. This may be called on an empty task source.
  virtual size_t GetRemainingConcurrency() const = 0;

  // Returns a TaskSourceSortKey representing the priority of the TaskSource.
  virtual TaskSourceSortKey GetSortKey() const = 0;
  // Returns a Timeticks object representing the next delayed runtime of the
  // TaskSource.
  virtual TimeTicks GetDelayedSortKey() const = 0;
  // Returns true if there are tasks ready to be executed. Thread-safe but the
  // returned value may immediately be obsolete.
  virtual bool HasReadyTasks(TimeTicks now) const = 0;
  // Returns true if the TaskSource should be moved to the immediate queue
  // due to ready delayed tasks. Note: Returns false if the TaskSource contains
  // ready delayed tasks, but expects to already be in the immediate queue.
  virtual bool OnBecomeReady() = 0;

  // Support for IntrusiveHeap in ThreadGroup::PriorityQueue.
  void SetImmediateHeapHandle(const HeapHandle& handle);
  void ClearImmediateHeapHandle();
  HeapHandle GetImmediateHeapHandle() const {
    return immediate_pq_heap_handle_;
  }

  HeapHandle immediate_heap_handle() const { return immediate_pq_heap_handle_; }

  // Support for IntrusiveHeap in ThreadGroup::DelayedPriorityQueue.
  void SetDelayedHeapHandle(const HeapHandle& handle);
  void ClearDelayedHeapHandle();
  HeapHandle GetDelayedHeapHandle() const { return delayed_pq_heap_handle_; }

  HeapHandle delayed_heap_handle() const { return delayed_pq_heap_handle_; }

  // Returns the shutdown behavior of all Tasks in the TaskSource. Can be
  // accessed without a Transaction because it is never mutated.
  TaskShutdownBehavior shutdown_behavior() const {