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 "pw_async2/runnable_dispatcher.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;
 
  void AllowBlocking() { blocking_is_allowed_ = true; }
 
  template <typename Pendable>
  Poll<internal::PendOutputOf<Pendable>> RunInTaskUntilStalled(
      Pendable& pendable) PW_LOCKS_EXCLUDED(impl::dispatcher_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();
  }
 
  uint32_t tasks_polled() const { return tasks_polled_; }
 
  uint32_t tasks_completed() const { return tasks_completed_; }
 
  uint32_t wake_count() const { return wake_count_; }
 
 private:
  // 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_;
  bool blocking_is_allowed_ = false;
 
  // TODO: b/401049619 - Optionally provide metrics for production dispatchers.
  uint32_t tasks_polled_ = 0u;
  uint32_t tasks_completed_ = 0u;
  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