Speculative fix for Scheduler resuming destroyed coro_handle

couchbaselabs · Sep 29, 2023 · 9eb10ba · 9eb10ba
1 parent 5d8abd2
commit 9eb10ba
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/include/Scheduler.hh b/include/Scheduler.hh
@@ -213,12 +213,21 @@ namespace crouton {
     /** General purpose Awaitable to return from `yield_value`.
         It does nothing, just allows the Scheduler to schedule another runnable task if any. */
     struct Yielder : public CORO_NS::suspend_always {
-        explicit Yielder(coro_handle myHandle) :_handle(myHandle) { }
         coro_handle await_suspend(coro_handle h) noexcept {
+            _handle = h;
             return lifecycle::yieldingTo(h, Scheduler::current().yield(h));
         }
-        void await_resume() const noexcept {
+
+        void await_resume() noexcept {
             Scheduler::current().resumed(_handle);
+            _handle = nullptr;
+        }
+
+        ~Yielder() {
+            if (_handle) {
+                // If await_resume hasn't been called yet, then the coro is being destroyed.
+                Scheduler::current().destroying(_handle);
+            }
         }
     private:
         coro_handle _handle;

diff --git a/include/Task.hh b/include/Task.hh
@@ -78,15 +78,14 @@ namespace crouton {
         // `co_yield` returns a bool: true to keep running, false if Task has been interrupted.
         auto yield_value(bool) {
             struct yielder : public Yielder {
-                explicit yielder(coro_handle myHandle, std::shared_ptr<shared> shared)
-                :Yielder(myHandle), _shared(std::move(shared)) { }
-                [[nodiscard]] bool await_resume() const noexcept {
+                explicit yielder(std::shared_ptr<shared> shared) :_shared(std::move(shared)) { }
+                [[nodiscard]] bool await_resume() noexcept {
                     Yielder::await_resume();
                     return !_shared->interrupt;
                 }
                 std::shared_ptr<shared> _shared;
             };
-            return yielder(handle(), _shared);
+            return yielder(_shared);
         }
 
         void return_void() {