YT-21634: Make shutdown sequences of Invokers more robust

Basically, there are two bugs/bad behaviors:
1) Some parts of shutdown are delegated to Finalizer/Shutdown invoker. This causes a race (not a data race, just non-deterministic behaviour) between deadlocking producer (if destruction of their callback causes a deadlock) and a finalizer/shutdown thread who closes the queue triggering deadlock event. We change behaviour so that queues are closed immediately thus if any deadlocks are present, they are guaranteed to happen and be reliably detected.

2) Most queues that implement a shutdown mechanism, do not have a proper way of double checking if queue was closed after they've checked previously but before they enqueued their task. This results in a callback being stuck in queue forever. On its own this is not an issues, but given the fact that callback can and usually does hold a strong reference to the class which usually owns the queue itself, we have a reference cycle queue -owns-> callback -owns-> queue. We implement appropriate double checking mechanics which ensure such a think never occurs.
ad8c42eb67f74bc19c32c85fd8b7790b294a7144

Author: E1pp

yt/yt/core/concurrency/action_queue.cpp

@@ -54,16 +54,14 @@ class TActionQueue::TImpl
 
     void Shutdown(bool graceful)
     {
-        if (Stopped_.exchange(true)) {
+        // Proper synchronization done via Queue_->Shutdown().
+        if (Stopped_.exchange(true, std::memory_order::relaxed)) {
             return;
         }
 
-        Queue_->Shutdown();
-
-        ShutdownInvoker_->Invoke(BIND_NO_PROPAGATE([graceful, thread = Thread_, queue = Queue_] {
-            thread->Stop(graceful);
-            queue->DrainConsumer();
-        }));
+        Queue_->Shutdown(graceful);
+        Thread_->Stop(graceful);
+        Queue_->OnConsumerFinished();
     }
 
     const IInvokerPtr& GetInvoker()
@@ -79,20 +77,14 @@ class TActionQueue::TImpl
     const TMpscSingleQueueSchedulerThreadPtr Thread_;
 
     const TShutdownCookie ShutdownCookie_;
-    const IInvokerPtr ShutdownInvoker_ = GetShutdownInvoker();
 
-    std::atomic<bool> Started_ = false;
     std::atomic<bool> Stopped_ = false;
 
 
     void EnsureStarted()
     {
-        if (Started_.load(std::memory_order::relaxed)) {
-            return;
-        }
-        if (Started_.exchange(true)) {
-            return;
-        }
+        // Thread::Start already has
+        // its own short-circ mechanism.
         Thread_->Start();
     }
 };

yt/yt/core/concurrency/delayed_executor.cpp

@@ -3,6 +3,7 @@
 #include "scheduler.h"
 #include "private.h"
 
+#include <yt/yt/core/actions/invoker_util.h>
 #include <yt/yt/core/misc/relaxed_mpsc_queue.h>
 #include <yt/yt/core/misc/singleton.h>
 
@@ -60,6 +61,41 @@ DEFINE_REFCOUNTED_TYPE(TDelayedExecutorEntry)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+//! (arkady-e1ppa) MO's/Fence explanation:
+/*
+    We want for shutdown to guarantee that no callback is left in any queue
+    once it is over. Otherwise, memory leak or a deadlock of Poller/GrpcServer
+    (or someone else who blocks thread until some callback is run) will occur.
+
+    We model our queue with Enqueue being RMW^rel(Queue_, x, y) and Dequeue
+    being RMW^acq(Queue_, x, y), where x is what we have read and y is what we
+    have observed. Missing callback would imply that in Submit method enqueueing |CB|
+    we have observed Stopping_ |false| (e.g. TThread::Start (c) returned |true|)
+    but also in ThreadMain during the SubmitQueue drain (f) we have not observed the
+    |CB|. Execution is schematically listed below:
+        T1(Submit)                          T2(Shutdown)
+    RMW^rel(Queue_, empty, CB) (a)       W^rel(Stopping_, true)        (d)
+            |sequenced-before                   |simply hb
+    Fence^sc                   (b)       Fence^sc                      (e)
+            |sequenced-before                   |sequenced-before
+    R^acq(Stopping_, false)    (c)       RMW^acq(Queue_, empty, empty) (f)
+
+    Since (c) reads |false| it must be reading from Stopping_ ctor which is
+    W^na(Stopping_, false) which preceedes (d) in modification order. Thus
+    (c) must read-from some modification preceding (d) in modification order (ctor)
+    and therefore (c) -cob-> (d) (coherence ordered before).
+    Likewise, (f) reads |empty| which can only be read from Queue_ ctor or
+    prior Dequeue both of which preceede (a) in modification order (ctor is obvious;
+    former Dequeue by assumption that no one has read |CB| ever: if some (a) was
+    prior to some Dequeue in modification order, |CB| would inevitably be read).
+    So, (f) -cob-> (a). For fences we now have to relations:
+    (b) -sb-> (c) -cob-> (d) -simply hb-> (e) => (b) -S-> (e)
+    (e) -sb-> (f) -cob-> (a) -sb-> (b) => (e) -S-> (b)
+    Here sb is sequenced-before and S is sequentially-consistent total ordering.
+    We have formed a loop in S thus contradicting the assumption.
+*/
+
+
 class TDelayedExecutorImpl
 {
 public:
@@ -139,9 +175,11 @@ class TDelayedExecutorImpl
     {
         YT_VERIFY(callback);
         auto entry = New<TDelayedExecutorEntry>(std::move(callback), deadline, std::move(invoker));
-        PollerThread_->EnqueueSubmission(entry);
+        PollerThread_->EnqueueSubmission(entry); // <- (a)
+
+        std::atomic_thread_fence(std::memory_order::seq_cst); // <- (b)
 
-        if (!PollerThread_->Start()) {
+        if (!PollerThread_->Start()) { // <- (c)
             if (auto callback = TakeCallback(entry)) {
                 callback(/*aborted*/ true);
             }
@@ -213,6 +251,43 @@ class TDelayedExecutorImpl
         NProfiling::TCounter CanceledCallbacksCounter_ = ConcurrencyProfiler.Counter("/delayed_executor/canceled_callbacks");
         NProfiling::TCounter StaleCallbacksCounter_ = ConcurrencyProfiler.Counter("/delayed_executor/stale_callbacks");
 
+        class TCallbackGuard
+        {
+        public:
+            TCallbackGuard(TCallback<void(bool)> callback, bool aborted) noexcept
+                : Callback_(std::move(callback))
+                , Aborted_(aborted)
+            { }
+
+            TCallbackGuard(TCallbackGuard&& other) = default;
+
+            TCallbackGuard(const TCallbackGuard&) = delete;
+
+            TCallbackGuard& operator=(const TCallbackGuard&) = delete;
+            TCallbackGuard& operator=(TCallbackGuard&&) = delete;
+
+            void operator()()
+            {
+                auto callback = std::move(Callback_);
+                YT_VERIFY(callback);
+                callback.Run(Aborted_);
+            }
+
+            ~TCallbackGuard()
+            {
+                if (Callback_) {
+                    YT_LOG_DEBUG("Aborting delayed executor callback");
+
+                    auto callback = std::move(Callback_);
+                    callback(/*aborted*/ true);
+                }
+            }
+
+        private:
+            TCallback<void(bool)> Callback_;
+            bool Aborted_;
+        };
+
 
         void StartPrologue() override
         {
@@ -240,6 +315,13 @@ class TDelayedExecutorImpl
                 ProcessQueues();
 
                 if (IsStopping()) {
+                    // We have Stopping_.store(true) in simply happens-before relation.
+                    // Assume Stopping_.store(true, release) is (d).
+                    // NB(arkady-e1ppa): At the time of writing it is seq_cst
+                    // actually, but
+                    // 1. We don't need it to be for correctness hehe
+                    // 2. It won't help us here anyway
+                    // 3. It might be changed as it could be suboptimal.
                     break;
                 }
 
@@ -251,6 +333,8 @@ class TDelayedExecutorImpl
                 EventCount_->Wait(cookie, deadline);
             }
 
+            std::atomic_thread_fence(std::memory_order::seq_cst); // <- (e)
+
             // Perform graceful shutdown.
 
             // First run the scheduled callbacks with |aborted = true|.
@@ -267,7 +351,7 @@ class TDelayedExecutorImpl
             // Now we handle the queued callbacks similarly.
             {
                 TDelayedExecutorEntryPtr entry;
-                while (SubmitQueue_.TryDequeue(&entry)) {
+                while (SubmitQueue_.TryDequeue(&entry)) { // <- (f)
                     runAbort(entry);
                 }
             }
@@ -349,7 +433,11 @@ class TDelayedExecutorImpl
         void RunCallback(const TDelayedExecutorEntryPtr& entry, bool abort)
         {
             if (auto callback = TakeCallback(entry)) {
-                (entry->Invoker ? entry->Invoker : DelayedInvoker_)->Invoke(BIND_NO_PROPAGATE(std::move(callback), abort));
+                auto invoker = entry->Invoker
+                    ? entry->Invoker
+                    : DelayedInvoker_;
+                invoker
+                    ->Invoke(BIND_NO_PROPAGATE(TCallbackGuard(std::move(callback), abort)));
             }
         }
     };

yt/yt/core/concurrency/fair_share_action_queue.cpp

@@ -105,16 +105,14 @@ class TFairShareActionQueue
 
     void Shutdown(bool graceful)
     {
-        if (Stopped_.exchange(true)) {
+        // Syncrhonization done via Queue_->Shutdown().
+        if (Stopped_.exchange(true, std::memory_order::relaxed)) {
             return;
         }
 
-        Queue_->Shutdown();
-
-        ShutdownInvoker_->Invoke(BIND([graceful, thread = Thread_, queue = Queue_] {
-            thread->Stop(graceful);
-            queue->DrainConsumer();
-        }));
+        Queue_->Shutdown(graceful);
+        Thread_->Stop(graceful);
+        Queue_->OnConsumerFinished();
     }
 
     const IInvokerPtr& GetInvoker(int index) override
@@ -156,18 +154,13 @@ class TFairShareActionQueue
 
     std::vector<TString> BucketNames_;
 
-    std::atomic<bool> Started_ = false;
     std::atomic<bool> Stopped_ = false;
 
 
     void EnsuredStarted()
     {
-        if (Started_.load(std::memory_order::relaxed)) {
-            return;
-        }
-        if (Started_.exchange(true)) {
-            return;
-        }
+        // Thread::Start already has
+        // its own short-circ.
         Thread_->Start();
     }
 };

yt/yt/core/concurrency/fair_share_invoker_queue.cpp

@@ -54,24 +54,17 @@ const IInvokerPtr& TFairShareInvokerQueue::GetInvoker(int bucketIndex, int queue
     return bucket.Invokers[queueIndex];
 }
 
-void TFairShareInvokerQueue::Shutdown()
+void TFairShareInvokerQueue::Shutdown(bool graceful)
 {
     for (auto& bucket : Buckets_) {
-        bucket.Queue->Shutdown();
+        bucket.Queue->Shutdown(graceful);
     }
 }
 
-void TFairShareInvokerQueue::DrainProducer()
+void TFairShareInvokerQueue::OnConsumerFinished()
 {
     for (auto& bucket : Buckets_) {
-        bucket.Queue->DrainProducer();
-    }
-}
-
-void TFairShareInvokerQueue::DrainConsumer()
-{
-    for (auto& bucket : Buckets_) {
-        bucket.Queue->DrainConsumer();
+        bucket.Queue->OnConsumerFinished();
     }
 }
 

yt/yt/core/concurrency/fair_share_invoker_queue.h

@@ -40,10 +40,9 @@ class TFairShareInvokerQueue
 
     const IInvokerPtr& GetInvoker(int bucketIndex, int queueIndex) const;
 
-    void Shutdown();
-
-    void DrainProducer();
-    void DrainConsumer();
+    // See TInvokerQueue::Shutdown/OnConsumerFinished.
+    void Shutdown(bool graceful = false);
+    void OnConsumerFinished();
 
     bool IsRunning() const;
 

yt/yt/core/concurrency/fair_share_thread_pool.cpp

@@ -187,6 +187,10 @@ class TFairShareQueue
     void Invoke(TClosure callback, TBucket* bucket)
     {
         auto guard = Guard(SpinLock_);
+        // See Shutdown.
+        if (Stopping_) {
+            return;
+        }
 
         QueueSize_.fetch_add(1, std::memory_order::relaxed);
 
@@ -227,13 +231,14 @@ class TFairShareQueue
     }
 
     void Shutdown()
-    {
-        Drain();
-    }
-
-    void Drain()
     {
         auto guard = Guard(SpinLock_);
+        // Setting under spinlock because this way
+        // we have atomicity of two actions:
+        // 1) Write/read flag and 2) Drain/Enqueue callback.
+        // See two_level_fair_share_thread_pool Queue
+        // for more detailed explanation.
+        Stopping_ = true;
         for (const auto& item : Heap_) {
             item.Bucket->Drain();
         }
@@ -339,7 +344,7 @@ class TFairShareQueue
     const TIntrusivePtr<NThreading::TEventCount> CallbackEventCount_;
 
     YT_DECLARE_SPIN_LOCK(NThreading::TSpinLock, SpinLock_);
-
+    bool Stopping_ = false;
     std::vector<THeapItem> Heap_;
 
     std::atomic<int> ThreadCount_ = 0;
@@ -548,14 +553,6 @@ class TFairShareThreadPool
         TThreadPoolBase::DoShutdown();
     }
 
-    TClosure MakeFinalizerCallback() override
-    {
-        return BIND_NO_PROPAGATE([queue = Queue_, callback = TThreadPoolBase::MakeFinalizerCallback()] {
-            callback();
-            queue->Drain();
-        });
-    }
-
     void DoConfigure(int threadCount) override
     {
         Queue_->Configure(threadCount);