dispatcher.h

// Copyright 2025 The Pigweed Authors
//
// 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
//
//     https://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.
#pragma once
 
#include <atomic>
#include <mutex>
#include <type_traits>
#include <utility>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/shared_ptr.h"
#include "pw_async2/context.h"
#include "pw_async2/func_task.h"
#include "pw_async2/future_task.h"
#include "pw_async2/internal/lock.h"
#include "pw_async2/task.h"
#include "pw_async2/waker.h"
#include "pw_containers/intrusive_forward_list.h"
#include "pw_sync/lock_annotations.h"
 
// Coroutines are supported if the build target depends on //pw_async2:coro.
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
#include <functional>  // std::invoke
 
#include "pw_async2/coro.h"
#include "pw_async2/coro_task.h"
#include "pw_async2/fallible_coro_task.h"
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
namespace pw::async2 {
 
 
class Dispatcher {
 public:
  Dispatcher(Dispatcher&) = delete;
  Dispatcher& operator=(Dispatcher&) = delete;
 
  Dispatcher(Dispatcher&&) = delete;
  Dispatcher& operator=(Dispatcher&&) = delete;
 
  virtual ~Dispatcher() PW_LOCKS_EXCLUDED(internal::lock()) { Destroy(); }
 
  void Post(Task& task) PW_LOCKS_EXCLUDED(internal::lock()) {
    {
      std::lock_guard lock(internal::lock());
      PostLocked(task);
    }
    Wake();
  }
 
  template <typename T>
  void PostShared(const SharedPtr<T>& task)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    PW_ASSERT(PostAllocatedTask(task));
  }
 
  template <typename TaskType,
            typename... Args,
            typename = std::enable_if_t<std::is_base_of_v<Task, TaskType>>>
  [[nodiscard]] SharedPtr<TaskType> Post(Allocator& allocator, Args&&... args) {
    auto task = allocator.MakeShared<TaskType>(std::forward<Args>(args)...);
    if (!PostAllocatedTask(task)) {
      return nullptr;
    }
    return task;
  }
 
  template <
      typename Func = void,
      int&... kExplicitGuard,
      typename Arg,
      typename ActualFunc =
          std::conditional_t<std::is_void_v<Func>, std::decay_t<Arg>, Func>,
      typename = std::enable_if_t<!std::is_base_of_v<Task, ActualFunc>>>
  [[nodiscard]] SharedPtr<FuncTask<ActualFunc>> Post(Allocator& allocator,
                                                     Arg&& func) {
    return Post<FuncTask<ActualFunc>>(allocator, std::forward<Arg>(func));
  }
 
  template <typename Func>
  [[nodiscard]] SharedPtr<RunOnceTask<Func>> RunOnce(Allocator& allocator,
                                                     Func&& func) {
    return Post<RunOnceTask<Func>>(allocator, std::forward<Func>(func));
  }
 
  template <typename Fut>
  [[nodiscard]] SharedPtr<FutureTask<Fut>> PostFuture(Allocator& allocator,
                                                      Fut&& future) {
    return Post<FutureTask<Fut>>(allocator, std::forward<Fut>(future));
  }
 
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
  template <typename T>
  [[nodiscard]] SharedPtr<CoroTask<T>> Post(Allocator& allocator,
                                            Coro<T>&& coro) {
    if (!coro.ok()) {
      return nullptr;
    }
    return Post<CoroTask<T>>(allocator, std::move(coro));
  }
 
  template <typename T,
            typename E = void,
            int&... kExplicitGuard,
            typename Arg,
            typename ErrorHandler =
                std::conditional_t<std::is_void_v<E>, std::decay_t<Arg>, E>>
  [[nodiscard]] SharedPtr<FallibleCoroTask<T, ErrorHandler>> Post(
      Allocator& allocator, Coro<T>&& coro, Arg&& error_handler) {
    if (!coro.ok()) {
      return nullptr;
    }
    return Post<FallibleCoroTask<T, ErrorHandler>>(
        allocator, std::move(coro), std::forward<Arg>(error_handler));
  }
 
  template <auto kCoroFunc, typename... Args>
  [[nodiscard]] auto Post(CoroContext coro_cx, Args&&... args) {
    if constexpr (std::is_member_function_pointer_v<decltype(kCoroFunc)>) {
      return PostSharedMemberCoro<kCoroFunc>(coro_cx,
                                             std::forward<Args>(args)...);
    } else {
      return Post(coro_cx.allocator(),
                  std::invoke(kCoroFunc, coro_cx, std::forward<Args>(args)...));
    }
  }
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
  void LogRegisteredTasks() PW_LOCKS_EXCLUDED(internal::lock());
 
 protected:
  constexpr Dispatcher() = default;
 
  bool PopAndRunAllReadyTasks() PW_LOCKS_EXCLUDED(internal::lock());
 
  Task* PopTaskToRun() PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    return PopTaskToRunLocked();
  }
 
  Task* PopTaskToRun(bool& has_posted_tasks)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    Task* task = PopTaskToRunLocked();
    has_posted_tasks = task != nullptr || !sleeping_.empty();
    return task;
  }
 
  Task* PopSingleTaskForThisWake() PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    SetWantsWake();
    return PopTaskToRunLocked();
  }
 
  RunTaskResult RunTask(Task& task) PW_LOCKS_EXCLUDED(internal::lock()) {
    return task.RunInDispatcher();
  }
 
 private:
  friend class Task;
  friend class Waker;
 
  // Allow DispatcherForTestFacade to wrap another dispatcher (call Do*).
  template <typename>
  friend class DispatcherForTestFacade;
 
  // Posts a task, but does not wake the dispatcher, which is done after the
  // lock is released.
  void PostLocked(Task& task) PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  virtual void DoWake() PW_LOCKS_EXCLUDED(internal::lock()) = 0;
 
  // Prefer to call Wake() without the lock held, but it can be called with it
  // when necessary (see WakeTask()).
  void Wake() {
    if (wants_wake_.exchange(false, std::memory_order_relaxed)) {
      DoWake();
    }
  }
 
  Task* PopTaskToRunLocked() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Removes references to this `Dispatcher` from all linked `Task`s and
  // `Waker`s. Use a separate function so thread safety analysis applies.
  void Destroy() PW_LOCKS_EXCLUDED(internal::lock());
 
  static void UnpostTaskList(IntrusiveForwardList<Task>& list)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  void RemoveWokenTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    woken_.remove(task);
  }
  void RemoveSleepingTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    sleeping_.remove(task);
  }
  void AddSleepingTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    sleeping_.push_front(task);
  }
 
  // For use by `Waker`.
  void WakeTask(Task& task) PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  void LogTaskWakers(const Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Indicates that this Dispatcher should be woken when Wake() is called. This
  // prevents unnecessary wakes when, for example, multiple wakers wake the same
  // task or multiple tasks are posted before the dipsatcher runs.
  //
  // Must be called while the lock is held to prevent missed wakes.
  void SetWantsWake() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    wants_wake_.store(true, std::memory_order_relaxed);
  }
 
  // Checks if `task` contains a value and posts it to the dispatcher as an
  // allocated task. Returns true if the task posted successfully.
  template <typename T>
  bool PostAllocatedTask(const SharedPtr<T>& task)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    return PostAllocatedTask(
        task.get(),
        task.GetControlBlock(
            allocator::internal::ControlBlockHandle::GetInstance_DO_NOT_USE()));
  }
 
  bool PostAllocatedTask(Task* task,
                         allocator::internal::ControlBlock* control_block)
      PW_LOCKS_EXCLUDED(internal::lock());
 
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
  template <auto kCoroMemberFunc, typename Receiver, typename... Args>
  [[nodiscard]] auto PostSharedMemberCoro(CoroContext coro_cx,
                                          Receiver&& receiver,
                                          Args&&... args) {
    return Post(coro_cx.allocator(),
                std::invoke(kCoroMemberFunc,
                            std::forward<Receiver>(receiver),
                            coro_cx,
                            std::forward<Args>(args)...));
  }
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
  // TODO: b/491844340 - Evaluate IntrusiveForwardList performance and consider
  //     alternatives.
  IntrusiveForwardList<Task> woken_ PW_GUARDED_BY(internal::lock());
  IntrusiveForwardList<Task> sleeping_ PW_GUARDED_BY(internal::lock());
 
  // Latches wake requests to avoid duplicate DoWake calls.
  std::atomic<bool> wants_wake_ = false;
};
 
 
}  // namespace pw::async2