dispatcher_for_test.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 <utility>
 
#include "pw_async2/runnable_dispatcher.h"
#include "pw_async2/waker.h"
 
namespace pw::async2 {
namespace internal {
 
// This should be updated to work for futures instead of generic pendables.
template <typename Pendable>
class PendableAsTaskWithOutput : public Task {
 public:
  using value_type = PendOutputOf<Pendable>;
  PendableAsTaskWithOutput(Pendable& pendable)
      : pendable_(pendable), output_(Pending()) {}
 
  Poll<value_type> TakePoll() { return std::move(output_); }
 
 private:
  Poll<> DoPend(Context& cx) final {
    output_ = pendable_.Pend(cx);
    return output_.Readiness();
  }
  Pendable& pendable_;
  Poll<value_type> output_;
};
 
}  // namespace internal
 
 
template <typename Native>
class DispatcherForTestFacade final : public RunnableDispatcher {
 public:
  DispatcherForTestFacade() = default;
 
  DispatcherForTestFacade(const DispatcherForTestFacade&) = delete;
  DispatcherForTestFacade& operator=(const DispatcherForTestFacade&) = delete;
 
  DispatcherForTestFacade(DispatcherForTestFacade&&) = delete;
  DispatcherForTestFacade& operator=(DispatcherForTestFacade&&) = delete;
 
  ~DispatcherForTestFacade() override;
 
  void AllowBlocking() { blocking_is_allowed_ = true; }
 
  template <typename Pendable>
  Poll<internal::PendOutputOf<Pendable>> RunInTaskUntilStalled(
      Pendable& pendable) PW_LOCKS_EXCLUDED(internal::lock()) {
    internal::PendableAsTaskWithOutput<Pendable> task(pendable);
    native().Post(task);
    native().RunUntilStalled();
 
    // Ensure that the task is no longer registered, as it will be destroyed
    // once we return.
    //
    // This operation will not block because we are on the dispatcher thread
    // and the dispatcher is not currently running (we just ran it).
    task.Deregister();
 
    return task.TakePoll();
  }
 
  void RunToCompletionUntilReleased();
 
  void Release();
 
  uint32_t tasks_polled() const {
    return tasks_polled_.load(std::memory_order_relaxed);
  }
 
  uint32_t tasks_completed() const {
    return tasks_completed_.load(std::memory_order_relaxed);
  }
 
  uint32_t wake_count() const {
    return wake_count_.load(std::memory_order_relaxed);
  }
 
 private:
  // Task to keep the Dispatcher busy in RunToCompletionUntilReleased().
  class IdleTask : public Task {
   public:
    IdleTask() : Task(PW_ASYNC_TASK_NAME("IdleTask")) {}
 
    ~IdleTask() override;
 
    void Complete() {
      should_complete_.store(true, std::memory_order_relaxed);
      std::move(waker_).Wake();
    }
 
    void Reset() {
      PW_DASSERT(!IsRegistered());
      should_complete_.store(false, std::memory_order_relaxed);
    }
 
   private:
    Poll<> DoPend(Context& cx) override;
 
    Waker waker_;
    std::atomic<bool> should_complete_ = false;
  };
 
  // These functions are implemented in dispatcher_for_test.cc for the
  // NativeDispatcherForTest specialization only.
  bool DoRunUntilStalled() override;
 
  void DoWake() override;
 
  void DoWaitForWake() override;
 
  RunnableDispatcher& native() { return native_; }
 
  Native native_;
  IdleTask idle_task_;
  bool blocking_is_allowed_ = false;
  std::atomic<int> blocking_until_released_ = false;
 
  // TODO: b/401049619 - Optionally provide metrics for production dispatchers.
  std::atomic<uint32_t> tasks_polled_ = 0u;
  std::atomic<uint32_t> tasks_completed_ = 0u;
  std::atomic<uint32_t> wake_count_ = 0u;
};
 
}  // namespace pw::async2
 
#include "pw_async2_backend/native_dispatcher_for_test.h"
 
namespace pw::async2 {
 
using DispatcherForTest =
    DispatcherForTestFacade<backend::NativeDispatcherForTest>;
 
}  // namespace pw::async2