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 "pw_async2/context.h"
#include "pw_async2/internal/lock.h"
#include "pw_async2/task.h"
#include "pw_async2/waker.h"
#include "pw_containers/intrusive_list.h"
#include "pw_sync/lock_annotations.h"
 
namespace pw::async2 {
namespace internal {
 
template <typename T>
using PendOutputOf = typename decltype(std::declval<T>().Pend(
    std::declval<Context&>()))::value_type;
 
}  // namespace internal
 
 
class Dispatcher {
 public:
  Dispatcher(Dispatcher&) = delete;
  Dispatcher& operator=(Dispatcher&) = delete;
 
  Dispatcher(Dispatcher&&) = delete;
  Dispatcher& operator=(Dispatcher&&) = delete;
 
  virtual ~Dispatcher() { Deregister(); }
 
  void Post(Task& task) PW_LOCKS_EXCLUDED(internal::lock());
 
  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());
    set_wants_wake();
    return PopTaskToRunLocked();
  }
 
  enum RunTaskResult {
    kActive = Task::kActive,
 
    kDeregistered = Task::kDeregistered,
 
    kCompleted = Task::kCompleted,
  };
 
  RunTaskResult RunTask(Task& task) PW_LOCKS_EXCLUDED(internal::lock());
 
 private:
  friend class Task;
  friend class Waker;
 
  // Allow DispatcherForTestFacade to wrap another dispatcher (call Do*).
  template <typename>
  friend class DispatcherForTestFacade;
 
  virtual void DoWake() PW_LOCKS_EXCLUDED(internal::lock()) = 0;
 
  void Wake() {
    if (wanted_wake()) {
      DoWake();
    }
  }
 
  Task* PopTaskToRunLocked() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Removes references to this `Dispatcher` from all linked `Task`s and
  // `Waker`s.
  void Deregister() PW_LOCKS_EXCLUDED(internal::lock());
 
  static void UnpostTaskList(IntrusiveList<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 set_wants_wake() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    wants_wake_.store(true, std::memory_order_relaxed);
  }
 
  [[nodiscard]] bool wanted_wake() PW_NO_LOCK_SAFETY_ANALYSIS {
    return wants_wake_.exchange(false, std::memory_order_relaxed);
  }
 
  IntrusiveList<Task> woken_ PW_GUARDED_BY(internal::lock());
  IntrusiveList<Task> sleeping_ PW_GUARDED_BY(internal::lock());
 
  // Latches wake requests to avoid duplicate DoWake calls.
  std::atomic<bool> wants_wake_ PW_GUARDED_BY(internal::lock()) = false;
};
 
 
}  // namespace pw::async2