task.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 "pw_allocator/internal/control_block.h"
#include "pw_allocator/shared_ptr.h"
#include "pw_assert/assert.h"
#include "pw_async2/internal/lock.h"
#include "pw_async2/poll.h"
#include "pw_async2/waker.h"
#include "pw_containers/intrusive_forward_list.h"
#include "pw_containers/intrusive_list.h"
#include "pw_containers/intrusive_queue.h"
#include "pw_log/tokenized_args.h"
#include "pw_sync/lock_annotations.h"
 
namespace pw::async2 {
 
 
#define PW_ASYNC_TASK_NAME(name) PW_LOG_TOKEN_EXPR("pw_async2", name)
 
class Dispatcher;
 
 
 
class Context {
 public:
  Context(const Context&) = delete;
  Context(Context&&) = delete;
 
  Context& operator=(const Context&) = delete;
  Context& operator=(Context&&) = delete;
 
  void ReEnqueue();  // Implemented inline in task.h after Task is defined.
 
  PendingType Unschedule();
 
 private:
  friend class Task;
 
  constexpr Context() = default;
};
 
 
 
enum class RunTaskResult {
  kActive,
 
  kDeregistered,
 
  kCompleted,
};
 
class Task : public IntrusiveQueue<Task>::Item, private Context {
 public:
  explicit constexpr Task(log::Token name = kDefaultName) : name_(name) {}
 
  Task(const Task&) = delete;
  Task(Task&&) = delete;
  Task& operator=(const Task&) = delete;
  Task& operator=(Task&&) = delete;
 
  virtual ~Task();
 
  Poll<> Pend(Context& cx) { return DoPend(cx); }
 
  bool IsRegistered() const;
 
  void Deregister() PW_LOCKS_EXCLUDED(internal::lock());
 
  void Join() PW_LOCKS_EXCLUDED(internal::lock());
 
 protected:
  using Context = ::pw::async2::Context;
 
 private:
  friend class Dispatcher;
  friend class Waker;
  friend Context;
 
  static constexpr log::Token kDefaultName =
      PW_LOG_TOKEN("pw_async2", "(anonymous)");
 
  bool TryDeregister() PW_LOCKS_EXCLUDED(internal::lock());
 
  virtual Poll<> DoPend(Context&) = 0;
 
  // Sets this task to use the provided dispatcher.
  void PostTo(Dispatcher& dispatcher)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    PW_DASSERT(state_ == State::kUnposted);
    PW_DASSERT(dispatcher_ == nullptr);
    state_ = State::kWoken;
    dispatcher_ = &dispatcher;
  }
 
  // Removes the task from the dispatcher. Returns the ControlBlock* if the
  // dispatcher has a shared reference to this task.
  allocator::internal::ControlBlock* Unpost()
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Unposts and releases the dispatcher's shared reference to the task, if any.
  // Unconditionally releases the lock since it cannot be held while a task
  // destructor is called.
  void UnpostAndReleaseRef() PW_UNLOCK_FUNCTION(internal::lock());
 
  // Unposts and releases the dispatcher's shared reference to the task, if any.
  // If the dispatcher has a shared reference to the task, the lock is
  // released to destroy the task, then reacquired. This should ONLY be called
  // from the Dispatcher's destructor, since no tasks should be posted or run
  // while the Dispatcher is being destroyed.
  void UnpostAndReleaseRefFromDispatcherDestructor()
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  void MarkRunning() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    state_ = State::kRunning;
  }
 
  // Called by the dispatcher to run this task. The task has already been marked
  // as running.
  RunTaskResult RunInDispatcher() PW_LOCKS_EXCLUDED(internal::lock());
 
  // Called by a waker to wake this task.
  void Wake() PW_UNLOCK_FUNCTION(internal::lock());
 
  // Unlinks all `Waker` objects associated with this `Task.`
  void RemoveAllWakersLocked() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Adds a `Waker` to the linked list of `Waker`s tracked by this `Task`.
  void AddWakerLocked(Waker& waker)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    wakers_.push_front(waker);
  }
 
  // Removes a `Waker` from the linked list of `Waker`s tracked by this `Task`.
  //
  // Precondition: the provided waker *must* be in the list of `Waker`s tracked
  // by this `Task`.
  void RemoveWakerLocked(Waker& waker)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    wakers_.remove(waker);
  }
 
  void ReleaseSharedRef(allocator::internal::ControlBlock* control_block)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    // Create a SharedPtr to decrement the ref count and possibly delete this.
    SharedPtr<Task> temp(this, control_block);
  }
 
  // Sets the control block for an allocated BEFORE the task is posted. Since
  // the task hasn't been posted yet, it's not necessary to hold the lock.
  void SetControlBlockBeforePosted(
      allocator::internal::ControlBlock& control_block)
      PW_NO_LOCK_SAFETY_ANALYSIS {
    control_block_ = &control_block;
  }
 
  enum class State : unsigned char {
    kUnposted,
    kSleeping,
    kRunning,
    kDeregisteredButRunning,
    kWokenWhileRunning,
    kWoken,
  };
 
  // The current state of the task.
  State state_ PW_GUARDED_BY(internal::lock()) = State::kUnposted;
 
  // Indicates whether this task needs a waker if it is Pending.
  enum : bool {
    kWakerNeeded,
    kNoWakerNeeded
  } waker_requirement_ = kWakerNeeded;
 
  // A pointer to the dispatcher this task is associated with.
  //
  // This will be non-null when `state_` is anything other than `kUnposted`.
  //
  // This value must be cleared by the dispatcher upon destruction in order to
  // prevent null access.
  Dispatcher* dispatcher_ PW_GUARDED_BY(internal::lock()) = nullptr;
 
  // The memory block that contains this task, if it was dynamically allocated.
  allocator::internal::ControlBlock* control_block_
      PW_GUARDED_BY(internal::lock()) = nullptr;
 
  // Linked list of `Waker` s that may awaken this `Task`.
  IntrusiveForwardList<Waker> wakers_ PW_GUARDED_BY(internal::lock());
 
  // Optional user-facing name for the task. If set, it will be included in
  // debug logs.
  log::Token name_;
};
 
 
 
inline PendingType Context::Unschedule() {
  static_cast<Task&>(*this).waker_requirement_ = Task::kNoWakerNeeded;
  return Pending();
}
 
inline void Context::ReEnqueue() {
  Waker(static_cast<Task&>(*this), {}).Wake();
}
 
 
}  // namespace pw::async2