fix: assertion failure in JobQueue::stop

include/xrpl/core/Coro.ipp

@@ -3,50 +3,76 @@
 
 #include <xrpl/basics/ByteUtilities.h>
 
+#include <mutex>
+
 namespace xrpl {
 
 template <class F>
 JobQueue::Coro::Coro(Coro_create_t, JobQueue& jq, JobType type, std::string const& name, F&& f)
     : jq_(jq)
     , type_(type)
     , name_(name)
-    , running_(false)
     , coro_(
-          [this, fn = std::forward<F>(f)](boost::coroutines::asymmetric_coroutine<void>::push_type& do_yield) {
+          [this, fn = std::forward<F>(f)](boost::coroutines::asymmetric_coroutine<int>::push_type& do_yield) {
+              state_ = CoroState::Running;
               yield_ = &do_yield;
-              yield();
-              fn(shared_from_this());
-#ifndef NDEBUG
-              finished_ = true;
-#endif
+              try
+              {
+                  yield();
+                  // self makes Coro alive until this function returns
+                  std::shared_ptr<Coro> self;
+                  self = shared_from_this();
+                  fn(self);
+                  state_ = CoroState::Finished;
+              }
+              catch (...)
+              {
+                  state_ = CoroState::Finished;
+                  throw;
+              }
           },
           boost::coroutines::attributes(megabytes(1)))
 {
 }
 
 inline JobQueue::Coro::~Coro()
 {
-#ifndef NDEBUG
-    XRPL_ASSERT(finished_, "xrpl::JobQueue::Coro::~Coro : is finished");
-#endif
+    cancel();
+    XRPL_ASSERT(state_ != CoroState::Running, "xrpl::JobQueue::Coro::~Coro : is not running");
 }
 
 inline void
-JobQueue::Coro::yield() const
+JobQueue::Coro::yield()
 {
+    CoroState expected = CoroState::Running;
+    if (!state_.compare_exchange_strong(expected, CoroState::Suspended))
+    {
+        return;
+    }
+
     {
         std::lock_guard lock(jq_.m_mutex);
+
         ++jq_.nSuspend_;
+        jq_.cv_.notify_all();
+        jq_.m_suspendedCoros[this] = weak_from_this();
     }
-    (*yield_)();
+
+    state_.notify_all();
+    (*yield_)(++yieldCount_);
 }
 
 inline bool
 JobQueue::Coro::post()
 {
+    if (state_ != CoroState::Suspended)
     {
-        std::lock_guard lk(mutex_run_);
-        running_ = true;
+        // The coroutine will run until it finishes if the JobQueue has stopped.
+        // In the case where make_shared<Coro>() succeeds and then the JobQueue
+        // stops before coro_ gets executed, post() will still be called and
+        // state_ will be Finished. We should return false and avoid XRPL_ASSERT
+        // as it's a valid edge case.
+        return false;
     }
 
     // sp keeps 'this' alive
@@ -55,68 +81,85 @@
         return true;
     }
 
-    // The coroutine will not run.  Clean up running_.
-    std::lock_guard lk(mutex_run_);
-    running_ = false;
-    cv_.notify_all();
     return false;
 }
 
 inline void
 JobQueue::Coro::resume()
 {
+    auto suspended = CoroState::Suspended;
+    if (!state_.compare_exchange_strong(suspended, CoroState::Running))
     {
-        std::lock_guard lk(mutex_run_);
-        running_ = true;
+        return;
     }
+
+    state_.notify_all();
+
+    // There's a small chance that the coroutine has not yielded yet and is
+    // still running. We need to wait for it to yield before we can resume it.
+    waitForYield();
+
     {
         std::lock_guard lock(jq_.m_mutex);
+        jq_.m_suspendedCoros.erase(this);
         --jq_.nSuspend_;
+        jq_.cv_.notify_all();
     }
+
     auto saved = detail::getLocalValues().release();
     detail::getLocalValues().reset(&lvs_);
-    std::lock_guard lock(mutex_);
     XRPL_ASSERT(static_cast<bool>(coro_), "xrpl::JobQueue::Coro::resume : is runnable");
     coro_();
     detail::getLocalValues().release();
     detail::getLocalValues().reset(saved);
-    std::lock_guard lk(mutex_run_);
-    running_ = false;
-    cv_.notify_all();
 }
 
 inline bool
 JobQueue::Coro::runnable() const
 {
-    return static_cast<bool>(coro_);
+    // There's an edge case where the coroutine has updated the status
+    // to Finished but the function hasn't exited and therefore, coro_ is
+    // still valid. However, the coroutine is not technically runnable in this
+    // case, because the coroutine is about to exit and static_cast<bool>(coro_)
+    // is going to be false.
+    return static_cast<bool>(coro_) && state_ != CoroState::Finished;
+}
+
+inline void
+JobQueue::Coro::join()
+{
+    state_.wait(CoroState::Running);
 }
 
 inline void
-JobQueue::Coro::expectEarlyExit()
+JobQueue::Coro::cancel()
 {
-#ifndef NDEBUG
-    if (!finished_)
-#endif
+    auto suspended = CoroState::Suspended;
+    if (!state_.compare_exchange_strong(suspended, CoroState::Running))
+    {
+        return;
+    }
+
+    waitForYield();
+
+    coro_ = {};
+
     {
-        // expectEarlyExit() must only ever be called from outside the
-        // Coro's stack.  It you're inside the stack you can simply return
-        // and be done.
-        //
-        // That said, since we're outside the Coro's stack, we need to
-        // decrement the nSuspend that the Coro's call to yield caused.
         std::lock_guard lock(jq_.m_mutex);
+        jq_.m_suspendedCoros.erase(this);
         --jq_.nSuspend_;
-#ifndef NDEBUG
-        finished_ = true;
-#endif
+        jq_.cv_.notify_all();
     }
+
+    XRPL_ASSERT(state_ == CoroState::Finished, "ripple::JobQueue::Coro::cancel : should have finished");
 }
 
 inline void
-JobQueue::Coro::join()
+JobQueue::Coro::waitForYield() const
 {
-    std::unique_lock<std::mutex> lk(mutex_run_);
-    cv_.wait(lk, [this]() { return running_ == false; });
+    // Busy-wait for the coroutine to yield so that it's safe to destroy it.
+    while (yieldCount_ != coro_.get())
+        ;
 }
 
 }  // namespace xrpl

include/xrpl/core/JobQueue.h

@@ -37,23 +37,26 @@
 class JobQueue : private Workers::Callback
 {
 public:
+    enum class QueueState { Accepting, Stopping, Stopped };
+
     /** Coroutines must run to completion. */
     class Coro : public std::enable_shared_from_this<Coro>
     {
+        friend class JobQueue;
+
+    public:
+        enum class CoroState { None, Suspended, Running, Finished };
+
     private:
         detail::LocalValues lvs_;
         JobQueue& jq_;
         JobType type_;
         std::string name_;
-        bool running_;
-        std::mutex mutex_;
-        std::mutex mutex_run_;
-        std::condition_variable cv_;
-        boost::coroutines::asymmetric_coroutine<void>::pull_type coro_;
-        boost::coroutines::asymmetric_coroutine<void>::push_type* yield_;
-#ifndef NDEBUG
-        bool finished_ = false;
-#endif
+        std::atomic<CoroState> state_ = CoroState::None;
+        std::atomic_int yieldCount_ = 0;
+
+        boost::coroutines::asymmetric_coroutine<int>::pull_type coro_;
+        boost::coroutines::asymmetric_coroutine<int>::push_type* yield_;
 
     public:
         // Private: Used in the implementation
@@ -74,10 +77,12 @@
             Note:
               The associated Job function returns.
               Undefined behavior if called consecutively without a corresponding
-           post.
+              post.
+              It may not suspend at all if the JobQueue is stopping, and returns
+              false in such a case.
         */
         void
-        yield() const;
+        yield();
 
         /** Schedule coroutine execution.
             Effects:
@@ -107,17 +112,28 @@
         void
         resume();
 
+        CoroState
+        state() const
+        {
+            return state_;
+        }
+
         /** Returns true if the Coro is still runnable (has not returned). */
         bool
         runnable() const;
 
-        /** Once called, the Coro allows early exit without an assert. */
-        void
-        expectEarlyExit();
-
         /** Waits until coroutine returns from the user function. */
         void
         join();
+
+        /** Cancel the coroutine. */
+        void
+        cancel();
+
+    private:
+        /** Wait for the coroutine to yield. */
+        void
+        waitForYield() const;
     };
 
     using JobFunction = std::function<void()>;
@@ -144,12 +160,13 @@
         typename = std::enable_if_t<std::is_same<decltype(std::declval<JobHandler&&>()()), void>::value>>
     bool
     addJob(JobType type, std::string const& name, JobHandler&& jobHandler)
+        requires std::is_void_v<std::invoke_result_t<JobHandler>>
     {
-        if (auto optionalCountedJob = jobCounter_.wrap(std::forward<JobHandler>(jobHandler)))
+        if (queueState_ != QueueState::Accepting)
         {
-            return addRefCountedJob(type, name, std::move(*optionalCountedJob));
+            return false;
         }
-        return false;
+        return addJobNoStatusCheck(type, name, std::forward<JobHandler>(jobHandler));
     }
 
     /** Creates a coroutine and adds a job to the queue which will run it.
@@ -208,13 +225,16 @@
     bool
     isStopping() const
     {
-        return stopping_;
+        return queueState_ == QueueState::Stopping;
     }
 
     // We may be able to move away from this, but we can keep it during the
     // transition.
     bool
-    isStopped() const;
+    isStopped() const
+    {
+        return queueState_ == QueueState::Stopped;
+    }
 
 private:
     friend class Coro;
@@ -226,8 +246,7 @@
     std::uint64_t m_lastJob;
     std::set<Job> m_jobSet;
     JobCounter jobCounter_;
-    std::atomic_bool stopping_{false};
-    std::atomic_bool stopped_{false};
+    std::atomic<QueueState> queueState_{QueueState::Accepting};
     JobDataMap m_jobData;
     JobTypeData m_invalidJobData;
 
@@ -237,6 +256,8 @@
     // The number of suspended coroutines
     int nSuspend_ = 0;
 
+    std::map<void*, std::weak_ptr<Coro>> m_suspendedCoros;
+
     Workers m_workers;
 
     // Statistics tracking
@@ -252,6 +273,18 @@
     JobTypeData&
     getJobTypeData(JobType type);
 
+    template <typename JobHandler>
+    bool
+    addJobNoStatusCheck(JobType type, std::string const& name, JobHandler&& jobHandler)
+        requires std::is_void_v<std::invoke_result_t<JobHandler>>
+    {
+        if (auto optionalCountedJob = jobCounter_.wrap(std::forward<JobHandler>(jobHandler)))
+        {
+            return addRefCountedJob(type, name, std::move(*optionalCountedJob));
+        }
+        return false;
+    }
+
     // Adds a reference counted job to the JobQueue.
     //
     //    param type The type of job.
@@ -386,6 +419,10 @@
 std::shared_ptr<JobQueue::Coro>
 JobQueue::postCoro(JobType t, std::string const& name, F&& f)
 {
+    if (queueState_ != QueueState::Accepting)
+    {
+        return nullptr;
+    }
     /*  First param is a detail type to make construction private.
         Last param is the function the coroutine runs. Signature of
         void(std::shared_ptr<Coro>).
@@ -395,7 +432,6 @@
     {
         // The Coro was not successfully posted.  Disable it so it's destructor
         // can run with no negative side effects.  Then destroy it.
-        coro->expectEarlyExit();
         coro.reset();
     }
     return coro;