Fix potential receiver blocked in deconstruction

dbms/src/Flash/Mpp/MPPReceiverSet.h

@@ -30,6 +30,11 @@ class MPPReceiverSet
     explicit MPPReceiverSet(const String & req_id)
         : log(Logger::get(req_id))
     {}
+    ~MPPReceiverSet()
+    {
+        /// close will close every receiver's internal MPMC queue and avoid blocking risk in waitAllConnectionsDone
+        close();
+    }
     void addExchangeReceiver(const String & executor_id, const ExchangeReceiverPtr & exchange_receiver);
     void addCoprocessorReader(const CoprocessorReaderPtr & coprocessor_reader);
     ExchangeReceiverPtr getExchangeReceiver(const String & executor_id) const;

dbms/src/Flash/Mpp/MPPTask.cpp

@@ -258,35 +258,46 @@ void MPPTask::registerTunnels(const mpp::DispatchTaskRequest & task_request)
 void MPPTask::initExchangeReceivers()
 {
     auto receiver_set_local = std::make_shared<MPPReceiverSet>(log->identifier());
-    dag_context->dag_request.traverse([&](const tipb::Executor & executor) {
-        if (executor.tp() == tipb::ExecType::TypeExchangeReceiver)
-        {
-            assert(executor.has_executor_id());
-            const auto & executor_id = executor.executor_id();
-            // In order to distinguish different exchange receivers.
-            auto exchange_receiver = std::make_shared<ExchangeReceiver>(
-                std::make_shared<GRPCReceiverContext>(
-                    executor.exchange_receiver(),
-                    dag_context->getMPPTaskMeta(),
-                    context->getTMTContext().getKVCluster(),
-                    context->getTMTContext().getMPPTaskManager(),
-                    context->getSettingsRef().enable_local_tunnel,
-                    context->getSettingsRef().enable_async_grpc_client),
-                executor.exchange_receiver().encoded_task_meta_size(),
-                context->getMaxStreams(),
-                log->identifier(),
-                executor_id,
-                executor.fine_grained_shuffle_stream_count(),
-                context->getSettingsRef());
-
-            if (status != RUNNING)
-                throw Exception(
-                    "exchange receiver map can not be initialized, because the task is not in running state");
-
-            receiver_set_local->addExchangeReceiver(executor_id, exchange_receiver);
-        }
-        return true;
-    });
+    try
+    {
+        dag_context->dag_request.traverse([&](const tipb::Executor & executor) {
+            if (executor.tp() == tipb::ExecType::TypeExchangeReceiver)
+            {
+                assert(executor.has_executor_id());
+                const auto & executor_id = executor.executor_id();
+                // In order to distinguish different exchange receivers.
+                auto exchange_receiver = std::make_shared<ExchangeReceiver>(
+                    std::make_shared<GRPCReceiverContext>(
+                        executor.exchange_receiver(),
+                        dag_context->getMPPTaskMeta(),
+                        context->getTMTContext().getKVCluster(),
+                        context->getTMTContext().getMPPTaskManager(),
+                        context->getSettingsRef().enable_local_tunnel,
+                        context->getSettingsRef().enable_async_grpc_client),
+                    executor.exchange_receiver().encoded_task_meta_size(),
+                    context->getMaxStreams(),
+                    log->identifier(),
+                    executor_id,
+                    executor.fine_grained_shuffle_stream_count(),
+                    context->getSettingsRef());
+
+                receiver_set_local->addExchangeReceiver(executor_id, exchange_receiver);
+
+                if (status != RUNNING)
+                    throw Exception(
+                        "exchange receiver map can not be initialized, because the task is not in running state");
+            }
+            return true;
+        });
+    }
+    catch (...)
+    {
+        std::unique_lock lock(mtx);
+        if (status != RUNNING)
+            throw Exception("exchange receiver map can not be initialized, because the task is not in running state");
+        receiver_set = std::move(receiver_set_local);
+        throw;
+    }
     {
         std::unique_lock lock(mtx);
         if (status != RUNNING)