Resume a held handle only while its awaiter is alive

rclcpp_async/include/rclcpp_async/co_context.hpp

@@ -499,10 +499,15 @@ void GoalStream<ActionT>::NextAwaiter::await_suspend(std::coroutine_handle<> h)
 template <typename ActionT>
 void GoalStream<ActionT>::CancelAwaiter::await_suspend(std::coroutine_handle<> h)
 {
+  waiting_ = std::make_shared<Waiting>(h);
+  // `this` is alive exactly while waiting_ is.
   stream.client_->async_cancel_goal(
-    stream.goal_handle_, [this, h](typename CancelResponse::SharedPtr resp) {
-      response = std::move(resp);
-      stream.ctx_.resume(h);
+    stream.goal_handle_,
+    [this, w = std::weak_ptr(waiting_)](typename CancelResponse::SharedPtr resp) {
+      if (auto s = w.lock()) {
+        response = std::move(resp);
+        stream.ctx_.resume(s->h);
+      }
     });
 }
 

rclcpp_async/include/rclcpp_async/goal_stream.hpp

@@ -165,6 +165,7 @@ class GoalStream
   {
     GoalStream & stream;
     typename CancelResponse::SharedPtr response;
+    std::shared_ptr<Waiting> waiting_;
 
     bool await_ready() { return false; }
 
@@ -173,7 +174,7 @@ class GoalStream
     typename CancelResponse::SharedPtr await_resume() { return std::move(response); }
   };
 
-  CancelAwaiter cancel_goal() { return CancelAwaiter{*this, nullptr}; }
+  CancelAwaiter cancel_goal() { return CancelAwaiter{*this, nullptr, nullptr}; }
 
   void set_auto_cancel_on_stop(bool v) { auto_cancel_on_stop_ = v; }
 

rclcpp_async/include/rclcpp_async/shield.hpp

@@ -16,6 +16,7 @@
 
 #include <coroutine>
 #include <exception>
+#include <memory>
 #include <type_traits>
 #include <utility>
 
@@ -35,12 +36,14 @@ template <typename T>
 struct ShieldAwaiter
 {
   Task<T> task_;
+  std::shared_ptr<Waiting> waiting_;
 
   bool await_ready() { return task_.handle.done(); }
 
   std::coroutine_handle<> await_suspend(std::coroutine_handle<> cont)
   {
-    task_.handle.promise().continuation = cont;
+    waiting_ = std::make_shared<Waiting>(cont);
+    task_.handle.promise().continuation = waiting_;
     if (task_.started_) {
       return std::noop_coroutine();
     }
@@ -71,7 +74,7 @@ struct ShieldAwaiter
 template <typename T>
 detail::ShieldAwaiter<T> shield(Task<T> task)
 {
-  return {std::move(task)};
+  return {std::move(task), nullptr};
 }
 
 }  // namespace rclcpp_async

rclcpp_async/include/rclcpp_async/task.hpp

@@ -24,6 +24,7 @@
 #include <utility>
 
 #include "rclcpp_async/cancelled_exception.hpp"
+#include "rclcpp_async/executor.hpp"
 
 namespace rclcpp_async
 {
@@ -39,7 +40,7 @@ struct Task
   {
     std::optional<T> value;
     std::exception_ptr exception;
-    std::coroutine_handle<> continuation;
+    std::weak_ptr<Waiting> continuation;
     std::stop_source stop_source;
 
     Task get_return_object()
@@ -54,8 +55,8 @@ struct Task
       bool await_ready() noexcept { return false; }
       std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept
       {
-        if (h.promise().continuation) {
-          return h.promise().continuation;
+        if (auto c = h.promise().continuation.lock()) {
+          return c->h;
         }
         return std::noop_coroutine();
       }
@@ -91,30 +92,40 @@ struct Task
 
   void set_token(std::stop_token t) { parent_token_ = std::move(t); }
 
-  // co_await Task (coroutine chaining)
-  bool await_ready() { return handle.done(); }
-  std::coroutine_handle<> await_suspend(std::coroutine_handle<> h)
+  // co_await Task (coroutine chaining). The awaiter lives in the awaiting
+  // frame, so the continuation expires with it even if this Task outlives it.
+  struct Awaiter
   {
-    handle.promise().continuation = h;
-    if (parent_token_.stop_possible()) {
-      auto child = handle;
-      parent_cancel_cb_ = std::unique_ptr<StopCb>(
-        new StopCb(parent_token_, [child]() { child.promise().stop_source.request_stop(); }));
-    }
-    if (started_) {
-      return std::noop_coroutine();
+    Task & task;
+    std::shared_ptr<Waiting> waiting_;
+
+    bool await_ready() { return task.handle.done(); }
+    std::coroutine_handle<> await_suspend(std::coroutine_handle<> h)
+    {
+      waiting_ = std::make_shared<Waiting>(h);
+      task.handle.promise().continuation = waiting_;
+      if (task.parent_token_.stop_possible()) {
+        auto child = task.handle;
+        task.parent_cancel_cb_ = std::unique_ptr<StopCb>(new StopCb(
+          task.parent_token_, [child]() { child.promise().stop_source.request_stop(); }));
+      }
+      if (task.started_) {
+        return std::noop_coroutine();
+      }
+      task.started_ = true;
+      return task.handle;
     }
-    started_ = true;
-    return handle;
-  }
-  T await_resume()
-  {
-    parent_cancel_cb_.reset();
-    if (handle.promise().exception) {
-      std::rethrow_exception(handle.promise().exception);
+    T await_resume()
+    {
+      task.parent_cancel_cb_.reset();
+      if (task.handle.promise().exception) {
+        std::rethrow_exception(task.handle.promise().exception);
+      }
+      return std::move(*task.handle.promise().value);
     }
-    return std::move(*handle.promise().value);
-  }
+  };
+  Awaiter operator co_await() & { return Awaiter{*this, nullptr}; }
+  Awaiter operator co_await() && { return Awaiter{*this, nullptr}; }
 
   T & result()
   {
@@ -169,7 +180,7 @@ struct Task<void>
   struct promise_type
   {
     std::exception_ptr exception;
-    std::coroutine_handle<> continuation;
+    std::weak_ptr<Waiting> continuation;
     std::stop_source stop_source;
 
     Task get_return_object()
@@ -184,8 +195,8 @@ struct Task<void>
       bool await_ready() noexcept { return false; }
       std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept
       {
-        if (h.promise().continuation) {
-          return h.promise().continuation;
+        if (auto c = h.promise().continuation.lock()) {
+          return c->h;
         }
         return std::noop_coroutine();
       }
@@ -221,28 +232,38 @@ struct Task<void>
 
   void set_token(std::stop_token t) { parent_token_ = std::move(t); }
 
-  bool await_ready() { return handle.done(); }
-  std::coroutine_handle<> await_suspend(std::coroutine_handle<> h)
+  // See Task<T>::Awaiter.
+  struct Awaiter
   {
-    handle.promise().continuation = h;
-    if (parent_token_.stop_possible()) {
-      auto child = handle;
-      parent_cancel_cb_ = std::unique_ptr<StopCb>(
-        new StopCb(parent_token_, [child]() { child.promise().stop_source.request_stop(); }));
-    }
-    if (started_) {
-      return std::noop_coroutine();
+    Task & task;
+    std::shared_ptr<Waiting> waiting_;
+
+    bool await_ready() { return task.handle.done(); }
+    std::coroutine_handle<> await_suspend(std::coroutine_handle<> h)
+    {
+      waiting_ = std::make_shared<Waiting>(h);
+      task.handle.promise().continuation = waiting_;
+      if (task.parent_token_.stop_possible()) {
+        auto child = task.handle;
+        task.parent_cancel_cb_ = std::unique_ptr<StopCb>(new StopCb(
+          task.parent_token_, [child]() { child.promise().stop_source.request_stop(); }));
+      }
+      if (task.started_) {
+        return std::noop_coroutine();
+      }
+      task.started_ = true;
+      return task.handle;
     }
-    started_ = true;
-    return handle;
-  }
-  void await_resume()
-  {
-    parent_cancel_cb_.reset();
-    if (handle.promise().exception) {
-      std::rethrow_exception(handle.promise().exception);
+    void await_resume()
+    {
+      task.parent_cancel_cb_.reset();
+      if (task.handle.promise().exception) {
+        std::rethrow_exception(task.handle.promise().exception);
+      }
     }
-  }
+  };
+  Awaiter operator co_await() & { return Awaiter{*this, nullptr}; }
+  Awaiter operator co_await() && { return Awaiter{*this, nullptr}; }
 
   void cancel() { handle.promise().stop_source.request_stop(); }
 

rclcpp_async/test/test_task_destroy.cpp

@@ -394,3 +394,65 @@ TEST_F(TaskDestroyTest, OwnerDestroyedWhileItsTaskWaits)
   check([](TaskOwner & o) { return as_task(o.ctx.sleep(10s)); });
   check([](TaskOwner & o) -> Task<void> { co_await o.stream->next(); });
 }
+
+TEST_F(TaskDestroyTest, CancelGoalResponseAfterRequesterDestroyed)
+{
+  // As in GoalResponseThenRequesterDestroyed, the client answers in its own
+  // group, so the cancel response arrives only when client_executor spins.
+  auto group = node_->create_callback_group(rclcpp::CallbackGroupType::MutuallyExclusive, false);
+  rclcpp::executors::SingleThreadedExecutor client_executor;
+  client_executor.add_callback_group(group, node_->get_node_base_interface());
+  auto client = rclcpp_action::create_client<Fibonacci>(node_, "test_destroy_action", group);
+  std::shared_ptr<GoalHandle> goal_handle;  // keeps the goal active
+  action_server_ = rclcpp_action::create_server<Fibonacci>(
+    node_, "test_destroy_action",
+    [](const rclcpp_action::GoalUUID &, std::shared_ptr<const Fibonacci::Goal>) {
+      return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
+    },
+    [this](const std::shared_ptr<GoalHandle>) {
+      cancel_requested_ = true;
+      return rclcpp_action::CancelResponse::ACCEPT;
+    },
+    [&goal_handle](const std::shared_ptr<GoalHandle> gh) { goal_handle = gh; });
+  ASSERT_TRUE(client->wait_for_action_server(5s));
+
+  std::shared_ptr<GoalStream<Fibonacci>> stream;
+  auto send = [this, &client, &stream]() -> Task<void> {
+    auto r = co_await ctx_->send_goal<Fibonacci>(client, Fibonacci::Goal());
+    stream = *r.value;
+  };
+  auto sender = ctx_->create_task(send());
+  auto deadline = std::chrono::steady_clock::now() + 5s;
+  while (!stream && std::chrono::steady_clock::now() < deadline) {
+    executor_.spin_some();
+    client_executor.spin_some();
+    std::this_thread::sleep_for(1ms);
+  }
+  ASSERT_TRUE(stream);
+
+  auto cancel = [&stream]() -> Task<void> { co_await stream->cancel_goal(); };
+  {
+    auto task = ctx_->create_task(cancel());
+    deadline = std::chrono::steady_clock::now() + 5s;
+    while (!cancel_requested_ && std::chrono::steady_clock::now() < deadline) {
+      spin_for(10ms);
+    }
+    ASSERT_TRUE(cancel_requested_);
+  }
+  for (int i = 0; i < 20; i++) {
+    client_executor.spin_some();  // the cancel response finds the task gone
+    std::this_thread::sleep_for(10ms);
+  }
+  spin_for(50ms);
+}
+
+TEST_F(TaskDestroyTest, AwaitedTaskFinishesAfterAwaiterDestroyed)
+{
+  // The awaited task is owned outside the awaiting one, so it outlives it.
+  auto child = ctx_->create_task(as_task(ctx_->sleep(50ms)));
+  {
+    auto task = ctx_->create_task([&child]() -> Task<void> { co_await child; }());
+    spin_for(10ms);
+  }
+  spin_for(100ms);
+}