[BugFix] Fix SenderQueue hang when cancel (#35380)

Signed-off-by: stdpain <[email protected]> (cherry picked from commit 571b9e5)
StarRocks · Nov 24, 2023 · 8b17325 · 8b17325
1 parent 90a24c0
commit 8b17325
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 5 deletions.
diff --git a/be/src/runtime/sender_queue.cpp b/be/src/runtime/sender_queue.cpp
@@ -15,6 +15,7 @@
 #include "util/logging.h"
 #include "util/runtime_profile.h"
 #include "util/time.h"
+#include "util/uid_util.h"
 
 namespace starrocks {
 
@@ -566,6 +567,33 @@ void DataStreamRecvr::PipelineSenderQueue::clean_buffer_queues() {
     }
 }
 
+void DataStreamRecvr::PipelineSenderQueue::check_leak_closure() {
+    std::lock_guard<Mutex> l(_lock);
+    for (size_t i = 0; i < _chunk_queues.size(); i++) {
+        auto& chunk_queue = _chunk_queues[i];
+        ChunkItem item;
+        while (chunk_queue.size_approx() > 0) {
+            if (chunk_queue.try_dequeue(item)) {
+                if (item.closure != nullptr) {
+                    DCHECK(false) << "leak closure detected";
+                    LOG(WARNING) << "leak closure detected in fragment:" << print_id(_recvr->fragment_instance_id());
+                }
+            }
+        }
+    }
+
+    for (auto& [_, chunk_queues] : _buffered_chunk_queues) {
+        for (auto& [_, chunk_queue] : chunk_queues) {
+            for (auto& item : chunk_queue) {
+                if (item.closure != nullptr) {
+                    DCHECK(false) << "leak closure detected";
+                    LOG(WARNING) << "leak closure detected in fragment:" << print_id(_recvr->fragment_instance_id());
+                }
+            }
+        }
+    }
+}
+
 Status DataStreamRecvr::PipelineSenderQueue::try_to_build_chunk_meta(const PTransmitChunkParams& request,
                                                                      Metrics& metrics) {
     ScopedTimer<MonotonicStopWatch> wait_timer(metrics.wait_lock_timer);
@@ -750,14 +778,15 @@ Status DataStreamRecvr::PipelineSenderQueue::add_chunks(const PTransmitChunkPara
                 // We cannot early-return for short circuit, it may occur for parts of parallelism,
                 // and the other parallelism may need to proceed.
             }
-            if (_is_cancelled) {
-                clean_buffer_queues();
-                return Status::OK();
-            }
             _recvr->_num_buffered_bytes += chunk_bytes;
         }
     }
 
+    // if senderqueue is cancelled clear all closure buffers
+    if (_is_cancelled) {
+        clean_buffer_queues();
+    }
+
     return Status::OK();
 }
 

diff --git a/be/src/runtime/sender_queue.h b/be/src/runtime/sender_queue.h
@@ -137,7 +137,10 @@ class DataStreamRecvr::NonPipelineSenderQueue final : public DataStreamRecvr::Se
 class DataStreamRecvr::PipelineSenderQueue final : public DataStreamRecvr::SenderQueue {
 public:
     PipelineSenderQueue(DataStreamRecvr* parent_recvr, int num_senders, int degree_of_parallism);
-    ~PipelineSenderQueue() override { close(); }
+    ~PipelineSenderQueue() override {
+        check_leak_closure();
+        close();
+    }
 
     Status get_chunk(vectorized::Chunk** chunk, const int32_t driver_sequence = -1) override;
 
@@ -198,6 +201,8 @@ class DataStreamRecvr::PipelineSenderQueue final : public DataStreamRecvr::Sende
 
     void clean_buffer_queues();
 
+    void check_leak_closure();
+
     StatusOr<ChunkList> get_chunks_from_pass_through(const int32_t sender_id, size_t& total_chunk_bytes);
 
     template <bool need_deserialization>