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/future_task.h"
#include "pw_async2/runnable_dispatcher.h"
#include "pw_async2/waker.h"
 
namespace pw::async2 {
 
 
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 T>
  auto RunInTaskUntilStalled(T& future) PW_LOCKS_EXCLUDED(internal::lock()) {
    FutureTask<T&> task(future);
    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);
      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