Fix schedule_from task cleanup before start

include/nvexec/stream/schedule_from.cuh

@@ -19,10 +19,9 @@
 #pragma once
 
 #include "../../stdexec/execution.hpp"
+#include <memory>
 #include <utility>
 
-#include "../detail/cuda_atomic.cuh"  // IWYU pragma: keep
-
 #include "common.cuh"
 
 namespace nv::execution
@@ -71,21 +70,34 @@ namespace nv::execution
           opstate& opstate_;
         };
 
+        struct task_deleter
+        {
+          void operator()(task_t* task) const noexcept
+          {
+            if (task != nullptr)
+            {
+              task->free_(task);
+            }
+          }
+        };
+
+        using task_ptr_t = std::unique_ptr<task_t, task_deleter>;
+
         opstate(Sender&& sndr, Receiver&& rcvr, context ctx)
           : _strm::opstate_base<Receiver>(static_cast<Receiver&&>(rcvr), ctx)
           , ctx_(ctx)
           , storage_(host_allocate<variant_t>(this->status_, ctx.pinned_resource_))
-          , task_(host_allocate<task_t>(this->status_,
-                                        ctx.pinned_resource_,
-                                        receiver{*this},
-                                        storage_.get(),
-                                        this->get_stream(),
-                                        ctx.pinned_resource_)
-                    .release())
+          , task_(task_ptr_t{host_allocate<task_t>(this->status_,
+                                                   ctx.pinned_resource_,
+                                                   receiver{*this},
+                                                   storage_.get(),
+                                                   this->get_stream(),
+                                                   ctx.pinned_resource_)
+                               .release()})
           , env_(host_allocate(this->status_, ctx_.pinned_resource_, this->make_env()))
-          , inner_op_{
-              connect(static_cast<Sender&&>(sndr),
-                      enqueue_receiver_t{env_.get(), storage_.get(), task_, ctx_.hub_->producer()})}
+          , inner_op_{connect(
+              static_cast<Sender&&>(sndr),
+              enqueue_receiver_t{env_.get(), storage_.get(), task_.get(), ctx_.hub_->producer()})}
         {
           if (this->status_ == cudaSuccess)
           {
@@ -97,23 +109,21 @@ namespace nv::execution
 
         void start() & noexcept
         {
-          started_.test_and_set(::cuda::std::memory_order::relaxed);
-
           if (status_ != cudaSuccess)
           {
             // Couldn't allocate memory for operation state, complete with error
             STDEXEC::set_error(std::move(this->rcvr_), std::move(status_));
             return;
           }
 
+          task_.release();
           STDEXEC::start(inner_op_);
         }
 
         cudaError_t                  status_{cudaSuccess};
         context                      ctx_;
         host_ptr_t<variant_t>        storage_;
-        task_t*                      task_;
-        ::cuda::std::atomic_flag     started_{};
+        task_ptr_t                   task_;
         host_ptr_t<__decay_t<env_t>> env_{};
         inner_opstate_t              inner_op_;
       };

test/nvexec/continues_on.cpp

@@ -5,27 +5,96 @@
 #include "nvexec/stream_context.cuh"
 
 #include <memory_resource>
+#include <new>
 
 namespace
 {
   class pinned_memory_resource_t : public std::pmr::memory_resource
   {
+    std::size_t allocations_{};
+    std::size_t deallocations_{};
+    std::size_t fail_after_{};
+    bool        fail_enabled_{};
+
     void* do_allocate(std::size_t bytes, std::size_t) override
     {
+      if (fail_enabled_ && allocations_ == fail_after_)
+      {
+        throw std::bad_alloc();
+      }
+
       void* storage{};
       STDEXEC_TRY_CUDA_API(cudaMallocHost(&storage, bytes));
+      ++allocations_;
       return storage;
     }
 
     void do_deallocate(void* storage, std::size_t, std::size_t) override
     {
       STDEXEC_ASSERT_CUDA_API(cudaFreeHost(storage));
+      ++deallocations_;
     }
 
     auto do_is_equal(std::pmr::memory_resource const & other) const noexcept -> bool override
     {
       return this == &other;
     }
+
+   public:
+    auto allocations() const noexcept -> std::size_t
+    {
+      return allocations_;
+    }
+
+    auto deallocations() const noexcept -> std::size_t
+    {
+      return deallocations_;
+    }
+
+    void fail_after(std::size_t allocation) noexcept
+    {
+      fail_after_   = allocation;
+      fail_enabled_ = true;
+    }
+  };
+
+  struct noop_receiver
+  {
+    using receiver_concept = STDEXEC::receiver_tag;
+
+    auto get_env() const noexcept -> STDEXEC::env<>
+    {
+      return {};
+    }
+
+    void set_value() noexcept {}
+
+    template <class Error>
+    void set_error(Error&&) noexcept
+    {}
+
+    void set_stopped() noexcept {}
+  };
+
+  struct error_receiver
+  {
+    using receiver_concept = STDEXEC::receiver_tag;
+
+    cudaError_t* error_;
+
+    auto get_env() const noexcept -> STDEXEC::env<>
+    {
+      return {};
+    }
+
+    void set_value() noexcept {}
+
+    void set_error(cudaError_t error) noexcept
+    {
+      *error_ = error;
+    }
+
+    void set_stopped() noexcept {}
   };
 
   class destruction_probe_t
@@ -68,6 +137,46 @@ namespace
     STDEXEC::sync_wait(std::move(sndr));
   }
 
+  TEST_CASE("continues_on frees its task when the operation is not started",
+            "[cuda][stream][adaptors][continues_on]")
+  {
+    pinned_memory_resource_t pinned_memory;
+    nvexec::stream_context   ctx;
+    auto                     scheduler = ctx.get_scheduler();
+    scheduler.ctx_.pinned_resource_    = &pinned_memory;
+
+    auto sndr = STDEXEC::just() | STDEXEC::continues_on(scheduler);
+    {
+      auto op = STDEXEC::connect(std::move(sndr), noop_receiver{});
+      (void) op;
+    }
+
+    REQUIRE(pinned_memory.allocations() > 0);
+    REQUIRE(pinned_memory.allocations() == pinned_memory.deallocations());
+  }
+
+  TEST_CASE("schedule_from frees its task when setup fails",
+            "[cuda][stream][adaptors][schedule_from]")
+  {
+    pinned_memory_resource_t pinned_memory;
+    pinned_memory.fail_after(2);
+
+    nvexec::stream_context ctx;
+    auto                   scheduler = ctx.get_scheduler();
+    scheduler.ctx_.pinned_resource_  = &pinned_memory;
+
+    cudaError_t error = cudaSuccess;
+    auto        sndr  = STDEXEC::schedule_from(STDEXEC::schedule(scheduler));
+    {
+      auto op = STDEXEC::connect(std::move(sndr), error_receiver{&error});
+      STDEXEC::start(op);
+    }
+
+    REQUIRE(error == cudaErrorMemoryAllocation);
+    REQUIRE(pinned_memory.allocations() == 2);
+    REQUIRE(pinned_memory.allocations() == pinned_memory.deallocations());
+  }
+
   TEST_CASE("continues on after schedule", "[cuda][stream][adaptors][continues_on]")
   {
     nvexec::stream_context ctx;