grpc 1.48.0-x86_64-linux → 1.49.0.pre1-x86_64-linux

data/src/core/lib/promise/arena_promise.h

@@ -19,7 +19,7 @@
 
 #include <stdlib.h>
 
-#include <type_traits>
+#include <new>
 #include <utility>
 
 #include "absl/meta/type_traits.h"
@@ -32,59 +32,58 @@ namespace grpc_core {
 
 namespace arena_promise_detail {
 
-// Type erased promise stored in the arena.
 template <typename T>
-class ImplInterface {
- public:
+struct Vtable {
   // Poll the promise, once.
-  virtual Poll<T> PollOnce() = 0;
+  Poll<T> (*poll_once)(void** arg);
   // Destroy the underlying callable object if there is one.
   // Since we don't delete (the arena owns the memory) but we may need to call a
   // destructor, we expose this for when the ArenaPromise object is destroyed.
-  virtual void Destroy() = 0;
+  void (*destroy)(void** arg);
+};
 
- protected:
-  ~ImplInterface() = default;
+template <typename T>
+struct VtableAndArg {
+  const Vtable<T>* vtable;
+  void* arg;
 };
 
-// Implementation of ImplInterface for an empty object.
+// Implementation of Vtable for an empty object.
 // Used when an empty ArenaPromise is created, or when the ArenaPromise is moved
 // from. Since in either case these objects should not be polled, we simply
 // crash if it is.
 template <typename T>
-class NullImpl final : public ImplInterface<T> {
- public:
-  Poll<T> PollOnce() override {
-    abort();
-    GPR_UNREACHABLE_CODE(return Pending{});
-  }
-  void Destroy() override {}
-
-  static ImplInterface<T>* Get() {
-    static NullImpl<T> instance;
-    return &instance;
-  }
-
- private:
-  ~NullImpl() = default;
-};
+inline const Vtable<T>* null_impl() {
+  static const Vtable<T> vtable = {[](void**) -> Poll<T> {
+                                     abort();
+                                     GPR_UNREACHABLE_CODE(return Pending{});
+                                   },
+                                   [](void**) {}};
+  return &vtable;
+}
 
 // Implementation of ImplInterface for a callable object.
 template <typename T, typename Callable>
-class CallableImpl final : public ImplInterface<T> {
- public:
-  explicit CallableImpl(Callable&& callable) : callable_(std::move(callable)) {}
-  // Forward polls to the callable object.
-  Poll<T> PollOnce() override { return poll_cast<T>(callable_()); }
-  // Destroy destructs the callable object.
-  void Destroy() override { this->~CallableImpl(); }
-
- private:
-  // Should only be called by Destroy().
-  ~CallableImpl() = default;
+inline const Vtable<T>* allocated_callable_impl() {
+  static const Vtable<T> vtable = {
+      [](void** arg) -> Poll<T> {
+        return poll_cast<T>((*static_cast<Callable*>(*arg))());
+      },
+      [](void** arg) { static_cast<Callable*>(*arg)->~Callable(); }};
+  return &vtable;
+}
 
-  Callable callable_;
-};
+// Implementation of ImplInterface for a small callable object (one that fits
+// within the void* arg)
+template <typename T, typename Callable>
+inline const Vtable<T>* inlined_callable_impl() {
+  static const Vtable<T> vtable = {
+      [](void** arg) -> Poll<T> {
+        return poll_cast<T>((*reinterpret_cast<Callable*>(arg))());
+      },
+      [](void** arg) { reinterpret_cast<Callable*>(arg)->~Callable(); }};
+  return &vtable;
+}
 
 // If a callable object is empty we can substitute any instance of that callable
 // for the one we call (for how could we tell the difference)?
@@ -94,27 +93,12 @@ class CallableImpl final : public ImplInterface<T> {
 // (this comes up often when the promise only accesses context data from the
 // containing activity).
 template <typename T, typename Callable>
-class SharedImpl final : public ImplInterface<T>, private Callable {
- public:
-  // Call the callable, or at least an exact duplicate of it - if you have no
-  // members, all your instances look the same.
-  Poll<T> PollOnce() override { return Callable::operator()(); }
-  // Nothing to destroy.
-  void Destroy() override {}
-  // Return a pointer to the shared instance - these are singletons, and are
-  // needed just to get the vtable in place.
-  static SharedImpl* Get(Callable&& callable) {
-    static_assert(sizeof(SharedImpl) == sizeof(void*),
-                  "SharedImpl should be pointer sized");
-    static SharedImpl impl(std::forward<Callable>(callable));
-    return &impl;
-  }
-
- private:
-  explicit SharedImpl(Callable&& callable)
-      : Callable(std::forward<Callable>(callable)) {}
-  ~SharedImpl() = default;
-};
+inline const Vtable<T>* shared_callable_impl(Callable&& callable) {
+  static Callable instance = std::forward<Callable>(callable);
+  static const Vtable<T> vtable = {[](void**) -> Poll<T> { return instance(); },
+                                   [](void**) {}};
+  return &vtable;
+}
 
 // Redirector type: given a callable type, expose a Make() function that creates
 // the appropriate underlying implementation.
@@ -123,26 +107,41 @@ struct ChooseImplForCallable;
 
 template <typename T, typename Callable>
 struct ChooseImplForCallable<
-    T, Callable, absl::enable_if_t<!std::is_empty<Callable>::value>> {
-  static ImplInterface<T>* Make(Callable&& callable) {
-    return GetContext<Arena>()->template New<CallableImpl<T, Callable>>(
-        std::forward<Callable>(callable));
+    T, Callable,
+    absl::enable_if_t<!std::is_empty<Callable>::value &&
+                      (sizeof(Callable) > sizeof(void*))>> {
+  static void Make(Callable&& callable, VtableAndArg<T>* out) {
+    *out = {allocated_callable_impl<T, Callable>(),
+            GetContext<Arena>()->template New<Callable>(
+                std::forward<Callable>(callable))};
+  }
+};
+
+template <typename T, typename Callable>
+struct ChooseImplForCallable<
+    T, Callable,
+    absl::enable_if_t<!std::is_empty<Callable>::value &&
+                      (sizeof(Callable) <= sizeof(void*))>> {
+  static void Make(Callable&& callable, VtableAndArg<T>* out) {
+    out->vtable = inlined_callable_impl<T, Callable>();
+    new (&out->arg) Callable(std::forward<Callable>(callable));
   }
 };
 
 template <typename T, typename Callable>
 struct ChooseImplForCallable<
     T, Callable, absl::enable_if_t<std::is_empty<Callable>::value>> {
-  static ImplInterface<T>* Make(Callable&& callable) {
-    return SharedImpl<T, Callable>::Get(std::forward<Callable>(callable));
+  static void Make(Callable&& callable, VtableAndArg<T>* out) {
+    out->vtable =
+        shared_callable_impl<T, Callable>(std::forward<Callable>(callable));
   }
 };
 
 // Wrap ChooseImplForCallable with a friend approachable syntax.
 template <typename T, typename Callable>
-ImplInterface<T>* MakeImplForCallable(Callable&& callable) {
-  return ChooseImplForCallable<T, Callable>::Make(
-      std::forward<Callable>(callable));
+void MakeImplForCallable(Callable&& callable, VtableAndArg<T>* out) {
+  ChooseImplForCallable<T, Callable>::Make(std::forward<Callable>(callable),
+                                           out);
 }
 
 }  // namespace arena_promise_detail
@@ -159,38 +158,42 @@ class ArenaPromise {
             typename Ignored =
                 absl::enable_if_t<!std::is_same<Callable, ArenaPromise>::value>>
   // NOLINTNEXTLINE(google-explicit-constructor)
-  ArenaPromise(Callable&& callable)
-      : impl_(arena_promise_detail::MakeImplForCallable<T>(
-            std::forward<Callable>(callable))) {}
+  ArenaPromise(Callable&& callable) {
+    arena_promise_detail::MakeImplForCallable<T>(
+        std::forward<Callable>(callable), &vtable_and_arg_);
+  }
 
   // ArenaPromise is not copyable.
   ArenaPromise(const ArenaPromise&) = delete;
   ArenaPromise& operator=(const ArenaPromise&) = delete;
   // ArenaPromise is movable.
-  ArenaPromise(ArenaPromise&& other) noexcept : impl_(other.impl_) {
-    other.impl_ = arena_promise_detail::NullImpl<T>::Get();
+  ArenaPromise(ArenaPromise&& other) noexcept
+      : vtable_and_arg_(other.vtable_and_arg_) {
+    other.vtable_and_arg_.vtable = arena_promise_detail::null_impl<T>();
   }
   ArenaPromise& operator=(ArenaPromise&& other) noexcept {
-    impl_->Destroy();
-    impl_ = other.impl_;
-    other.impl_ = arena_promise_detail::NullImpl<T>::Get();
+    vtable_and_arg_.vtable->destroy(&vtable_and_arg_.arg);
+    vtable_and_arg_ = other.vtable_and_arg_;
+    other.vtable_and_arg_.vtable = arena_promise_detail::null_impl<T>();
     return *this;
   }
 
   // Destruction => call Destroy on the underlying impl object.
-  ~ArenaPromise() { impl_->Destroy(); }
+  ~ArenaPromise() { vtable_and_arg_.vtable->destroy(&vtable_and_arg_.arg); }
 
   // Expose the promise interface: a call operator that returns Poll<T>.
-  Poll<T> operator()() { return impl_->PollOnce(); }
+  Poll<T> operator()() {
+    return vtable_and_arg_.vtable->poll_once(&vtable_and_arg_.arg);
+  }
 
   bool has_value() const {
-    return impl_ != arena_promise_detail::NullImpl<T>::Get();
+    return vtable_and_arg_.vtable != arena_promise_detail::null_impl<T>();
   }
 
  private:
   // Underlying impl object.
-  arena_promise_detail::ImplInterface<T>* impl_ =
-      arena_promise_detail::NullImpl<T>::Get();
+  arena_promise_detail::VtableAndArg<T> vtable_and_arg_ = {
+      arena_promise_detail::null_impl<T>(), nullptr};
 };
 
 }  // namespace grpc_core

data/src/core/lib/promise/context.h

@@ -17,7 +17,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <type_traits>
 #include <utility>
 
 #include "src/core/lib/gpr/tls.h"

data/src/core/lib/promise/detail/basic_seq.h

@@ -17,11 +17,12 @@
 
 #include <grpc/support/port_platform.h>
 
+#include <stddef.h>
+
 #include <array>
 #include <cassert>
 #include <new>
 #include <tuple>
-#include <type_traits>
 #include <utility>
 
 #include "absl/meta/type_traits.h"
@@ -31,12 +32,29 @@
 #include "src/core/lib/gprpp/construct_destruct.h"
 #include "src/core/lib/promise/detail/promise_factory.h"
 #include "src/core/lib/promise/detail/promise_like.h"
-#include "src/core/lib/promise/detail/switch.h"
 #include "src/core/lib/promise/poll.h"
 
 namespace grpc_core {
 namespace promise_detail {
 
+// Given f0, ..., fn, call function idx and return the result.
+template <typename R, typename A, R (*... f)(A* arg)>
+class JumpTable {
+ public:
+  JumpTable() = delete;
+  JumpTable(const JumpTable&) = delete;
+
+  static R Run(size_t idx, A* arg) { return fs_[idx](arg); }
+
+ private:
+  using Fn = R (*)(A* arg);
+  static const Fn fs_[sizeof...(f)];
+};
+
+template <typename R, typename A, R (*... f)(A* arg)>
+const typename JumpTable<R, A, f...>::Fn
+    JumpTable<R, A, f...>::fs_[sizeof...(f)] = {f...};
+
 // Helper for SeqState to evaluate some common types to all partial
 // specializations.
 template <template <typename> class Traits, typename FPromise, typename FNext>
@@ -335,16 +353,14 @@ class BasicSeq {
   // parameter unpacking can work.
   template <char I>
   struct RunStateStruct {
-    BasicSeq* s;
-    Poll<Result> operator()() { return s->RunState<I>(); }
+    static Poll<Result> Run(BasicSeq* s) { return s->RunState<I>(); }
   };
 
   // Similarly placate those compilers for
   // DestructCurrentPromiseAndSubsequentFactories
   template <char I>
   struct DestructCurrentPromiseAndSubsequentFactoriesStruct {
-    BasicSeq* s;
-    void operator()() {
+    static void Run(BasicSeq* s) {
       return s->DestructCurrentPromiseAndSubsequentFactories<I>();
     }
   };
@@ -357,7 +373,8 @@ class BasicSeq {
   // Duff's device like mechanic for evaluating sequences.
   template <char... I>
   Poll<Result> Run(absl::integer_sequence<char, I...>) {
-    return Switch<Poll<Result>>(state_, RunStateStruct<I>{this}...);
+    return JumpTable<Poll<Result>, BasicSeq, RunStateStruct<I>::Run...>::Run(
+        state_, this);
   }
 
   // Run the appropriate destructors for a given state.
@@ -367,8 +384,9 @@ class BasicSeq {
   // which can choose the correct instance at runtime to destroy everything.
   template <char... I>
   void RunDestruct(absl::integer_sequence<char, I...>) {
-    Switch<void>(
-        state_, DestructCurrentPromiseAndSubsequentFactoriesStruct<I>{this}...);
+    JumpTable<void, BasicSeq,
+              DestructCurrentPromiseAndSubsequentFactoriesStruct<I>::Run...>::
+        Run(state_, this);
   }
 
  public:
@@ -403,20 +421,22 @@ class BasicSeq {
 
 // As above, but models a sequence of unknown size
 // At each element, the accumulator A and the current value V is passed to some
-// function of type F as f(V, A); f is expected to return a promise that
-// resolves to Traits::WrappedType.
-template <template <typename Wrapped> class Traits, typename F, typename Arg,
-          typename Iter>
+// function of type IterTraits::Factory as f(V, IterTraits::Argument); f is
+// expected to return a promise that resolves to IterTraits::Wrapped.
+template <class IterTraits>
 class BasicSeqIter {
  private:
-  using IterValue = decltype(*std::declval<Iter>());
-  using StateCreated = decltype(std::declval<F>()(std::declval<IterValue>(),
-                                                  std::declval<Arg>()));
-  using State = PromiseLike<StateCreated>;
-  using Wrapped = typename State::Result;
+  using Traits = typename IterTraits::Traits;
+  using Iter = typename IterTraits::Iter;
+  using Factory = typename IterTraits::Factory;
+  using Argument = typename IterTraits::Argument;
+  using IterValue = typename IterTraits::IterValue;
+  using StateCreated = typename IterTraits::StateCreated;
+  using State = typename IterTraits::State;
+  using Wrapped = typename IterTraits::Wrapped;
 
  public:
-  BasicSeqIter(Iter begin, Iter end, F f, Arg arg)
+  BasicSeqIter(Iter begin, Iter end, Factory f, Argument arg)
       : cur_(begin), end_(end), f_(std::move(f)) {
     if (cur_ == end_) {
       Construct(&result_, std::move(arg));
@@ -466,7 +486,7 @@ class BasicSeqIter {
   Poll<Wrapped> PollNonEmpty() {
     Poll<Wrapped> r = state_();
     if (absl::holds_alternative<Pending>(r)) return r;
-    return Traits<Wrapped>::template CheckResultAndRunNext<Wrapped>(
+    return Traits::template CheckResultAndRunNext<Wrapped>(
         std::move(absl::get<Wrapped>(r)), [this](Wrapped arg) -> Poll<Wrapped> {
           auto next = cur_;
           ++next;
@@ -476,17 +496,17 @@ class BasicSeqIter {
           cur_ = next;
           state_.~State();
           Construct(&state_,
-                    Traits<Wrapped>::CallSeqFactory(f_, *cur_, std::move(arg)));
+                    Traits::template CallSeqFactory(f_, *cur_, std::move(arg)));
           return PollNonEmpty();
         });
   }
 
   Iter cur_;
   const Iter end_;
-  GPR_NO_UNIQUE_ADDRESS F f_;
+  GPR_NO_UNIQUE_ADDRESS Factory f_;
   union {
     GPR_NO_UNIQUE_ADDRESS State state_;
-    GPR_NO_UNIQUE_ADDRESS Arg result_;
+    GPR_NO_UNIQUE_ADDRESS Argument result_;
   };
 };
 

data/src/core/lib/promise/detail/promise_factory.h

@@ -17,7 +17,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <type_traits>
 #include <utility>
 
 #include "absl/meta/type_traits.h"

data/src/core/lib/promise/map.h

@@ -20,7 +20,6 @@
 #include <stddef.h>
 
 #include <tuple>
-#include <type_traits>
 #include <utility>
 
 #include "absl/types/variant.h"

data/src/core/lib/promise/seq.h

@@ -20,6 +20,7 @@
 #include <utility>
 
 #include "src/core/lib/promise/detail/basic_seq.h"
+#include "src/core/lib/promise/detail/promise_like.h"
 #include "src/core/lib/promise/poll.h"
 
 namespace grpc_core {
@@ -50,6 +51,26 @@ struct SeqTraits {
 template <typename... Fs>
 using Seq = BasicSeq<SeqTraits, Fs...>;
 
+template <typename I, typename F, typename Arg>
+struct SeqIterTraits {
+  using Iter = I;
+  using Factory = F;
+  using Argument = Arg;
+  using IterValue = decltype(*std::declval<Iter>());
+  using StateCreated = decltype(std::declval<F>()(std::declval<IterValue>(),
+                                                  std::declval<Arg>()));
+  using State = PromiseLike<StateCreated>;
+  using Wrapped = typename State::Result;
+
+  using Traits = SeqTraits<Wrapped>;
+};
+
+template <typename Iter, typename Factory, typename Argument>
+struct SeqIterResultTraits {
+  using IterTraits = SeqIterTraits<Iter, Factory, Argument>;
+  using Result = BasicSeqIter<IterTraits>;
+};
+
 }  // namespace promise_detail
 
 // Sequencing combinator.
@@ -75,11 +96,11 @@ F Seq(F functor) {
 //   }
 //   return argument;
 template <typename Iter, typename Factory, typename Argument>
-promise_detail::BasicSeqIter<promise_detail::SeqTraits, Factory, Argument, Iter>
+typename promise_detail::SeqIterResultTraits<Iter, Factory, Argument>::Result
 SeqIter(Iter begin, Iter end, Argument argument, Factory factory) {
-  return promise_detail::BasicSeqIter<promise_detail::SeqTraits, Factory,
-                                      Argument, Iter>(
-      begin, end, std::move(factory), std::move(argument));
+  using Result = typename promise_detail::SeqIterResultTraits<Iter, Factory,
+                                                              Argument>::Result;
+  return Result(begin, end, std::move(factory), std::move(argument));
 }
 
 }  // namespace grpc_core

data/src/core/lib/promise/sleep.cc

@@ -16,10 +16,14 @@
 
 #include "src/core/lib/promise/sleep.h"
 
+#include <utility>
+
 #include <grpc/event_engine/event_engine.h>
 
-#include "src/core/lib/event_engine/event_engine_factory.h"
+#include "src/core/lib/event_engine/default_event_engine.h"
+#include "src/core/lib/gprpp/time.h"
 #include "src/core/lib/iomgr/exec_ctx.h"
+#include "src/core/lib/promise/activity.h"
 
 namespace grpc_core {
 
@@ -28,54 +32,46 @@ using ::grpc_event_engine::experimental::GetDefaultEventEngine;
 Sleep::Sleep(Timestamp deadline) : deadline_(deadline) {}
 
 Sleep::~Sleep() {
-  if (deadline_ == Timestamp::InfPast()) return;
-  ReleasableMutexLock lock(&mu_);
-  switch (stage_) {
-    case Stage::kInitial:
-      break;
-    case Stage::kStarted:
-      if (GetDefaultEventEngine()->Cancel(timer_handle_)) {
-        lock.Release();
-        OnTimer();
-      }
-      break;
-    case Stage::kDone:
-      break;
+  if (closure_ != nullptr) closure_->Cancel();
+}
+
+Poll<absl::Status> Sleep::operator()() {
+  // Invalidate now so that we see a fresh version of the time.
+  // TODO(ctiller): the following can be safely removed when we remove ExecCtx.
+  ExecCtx::Get()->InvalidateNow();
+  // If the deadline is earlier than now we can just return.
+  if (deadline_ <= ExecCtx::Get()->Now()) return absl::OkStatus();
+  if (closure_ == nullptr) {
+    // TODO(ctiller): it's likely we'll want a pool of closures - probably per
+    // cpu? - to avoid allocating/deallocating on fast paths.
+    closure_ = new ActiveClosure(deadline_);
   }
+  return Pending{};
 }
 
-void Sleep::OnTimer() {
-  Waker tmp_waker;
-  {
-    MutexLock lock(&mu_);
-    stage_ = Stage::kDone;
-    tmp_waker = std::move(waker_);
+Sleep::ActiveClosure::ActiveClosure(Timestamp deadline)
+    : waker_(Activity::current()->MakeOwningWaker()),
+      timer_handle_(GetDefaultEventEngine()->RunAfter(
+          deadline - ExecCtx::Get()->Now(), this)) {}
+
+void Sleep::ActiveClosure::Run() {
+  ApplicationCallbackExecCtx callback_exec_ctx;
+  ExecCtx exec_ctx;
+  auto waker = std::move(waker_);
+  if (refs_.Unref()) {
+    delete this;
+  } else {
+    waker.Wakeup();
   }
-  tmp_waker.Wakeup();
 }
 
-Poll<absl::Status> Sleep::operator()() {
-  MutexLock lock(&mu_);
-  switch (stage_) {
-    case Stage::kInitial:
-      if (deadline_ <= ExecCtx::Get()->Now()) {
-        return absl::OkStatus();
-      }
-      stage_ = Stage::kStarted;
-      timer_handle_ = GetDefaultEventEngine()->RunAfter(
-          deadline_ - ExecCtx::Get()->Now(), [this] {
-            ApplicationCallbackExecCtx callback_exec_ctx;
-            ExecCtx exec_ctx;
-            OnTimer();
-          });
-      break;
-    case Stage::kStarted:
-      break;
-    case Stage::kDone:
-      return absl::OkStatus();
+void Sleep::ActiveClosure::Cancel() {
+  // If we cancel correctly then we must own both refs still and can simply
+  // delete without unreffing twice, otherwise try unreffing since this may be
+  // the last owned ref.
+  if (GetDefaultEventEngine()->Cancel(timer_handle_) || refs_.Unref()) {
+    delete this;
   }
-  waker_ = Activity::current()->MakeNonOwningWaker();
-  return Pending{};
 }
 
 }  // namespace grpc_core

data/src/core/lib/promise/sleep.h

@@ -17,14 +17,12 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <utility>
-
-#include "absl/base/thread_annotations.h"
 #include "absl/status/status.h"
 
 #include <grpc/event_engine/event_engine.h>
+#include <grpc/support/log.h>
 
-#include "src/core/lib/gprpp/sync.h"
+#include "src/core/lib/gprpp/ref_counted.h"
 #include "src/core/lib/gprpp/time.h"
 #include "src/core/lib/promise/activity.h"
 #include "src/core/lib/promise/poll.h"
@@ -32,43 +30,48 @@
 namespace grpc_core {
 
 // Promise that sleeps until a deadline and then finishes.
-class Sleep {
+class Sleep final {
  public:
   explicit Sleep(Timestamp deadline);
   ~Sleep();
 
   Sleep(const Sleep&) = delete;
   Sleep& operator=(const Sleep&) = delete;
-  Sleep(Sleep&& other) noexcept
-      : deadline_(other.deadline_), timer_handle_(other.timer_handle_) {
-    MutexLock lock2(&other.mu_);
-    stage_ = other.stage_;
-    waker_ = std::move(other.waker_);
-    other.deadline_ = Timestamp::InfPast();
+  Sleep(Sleep&& other) noexcept : deadline_(other.deadline_) {
+    // Promises can be moved only until they're polled, and since we only create
+    // the closure when first polled we can assume it's nullptr here.
+    GPR_DEBUG_ASSERT(other.closure_ == nullptr);
   };
   Sleep& operator=(Sleep&& other) noexcept {
-    if (&other == this) return *this;
-    MutexLock lock1(&mu_);
-    MutexLock lock2(&other.mu_);
+    // Promises can be moved only until they're polled, and since we only create
+    // the closure when first polled we can assume it's nullptr here.
+    GPR_DEBUG_ASSERT(closure_ == nullptr);
+    GPR_DEBUG_ASSERT(other.closure_ == nullptr);
     deadline_ = other.deadline_;
-    timer_handle_ = other.timer_handle_;
-    stage_ = other.stage_;
-    waker_ = std::move(other.waker_);
-    other.deadline_ = Timestamp::InfPast();
     return *this;
   };
 
   Poll<absl::Status> operator()();
 
  private:
-  enum class Stage { kInitial, kStarted, kDone };
-  void OnTimer();
+  class ActiveClosure final
+      : public grpc_event_engine::experimental::EventEngine::Closure {
+   public:
+    explicit ActiveClosure(Timestamp deadline);
+
+    void Run() override;
+    void Cancel();
+
+   private:
+    Waker waker_;
+    // One ref dropped by Run(), the other by Cancel().
+    RefCount refs_{2};
+    const grpc_event_engine::experimental::EventEngine::TaskHandle
+        timer_handle_;
+  };
 
   Timestamp deadline_;
-  grpc_event_engine::experimental::EventEngine::TaskHandle timer_handle_;
-  Mutex mu_;
-  Stage stage_ ABSL_GUARDED_BY(mu_) = Stage::kInitial;
-  Waker waker_ ABSL_GUARDED_BY(mu_);
+  ActiveClosure* closure_{nullptr};
 };
 
 }  // namespace grpc_core