Avoid data race in grpc alarm

dbms/src/Flash/EstablishCall.cpp

@@ -216,6 +216,11 @@ void EstablishCallData::initRpc()
     }
 }
 
+grpc::Alarm & EstablishCallData::getAlarm()
+{
+    return alarm;
+}
+
 void EstablishCallData::tryConnectTunnel()
 {
     auto * task_manager = service->getContext()->getTMTContext().getMPPTaskManager().get();

dbms/src/Flash/EstablishCall.h

@@ -19,6 +19,7 @@
 #include <Common/Stopwatch.h>
 #include <Flash/FlashService.h>
 #include <Flash/Mpp/MPPTaskId.h>
+#include <grpcpp/alarm.h>
 #include <kvproto/tikvpb.grpc.pb.h>
 #include <prometheus/gauge.h>
 
@@ -95,6 +96,8 @@ class EstablishCallData final
     grpc::ServerContext * getGrpcContext() { return &ctx; }
 
     String getResourceGroupName() const { return resource_group_name; }
+    grpc::Alarm & getAlarm();
+
 
 private:
     /// WARNING: Since a event from one grpc completion queue may be handled by different
@@ -142,5 +145,6 @@ class EstablishCallData final
     String resource_group_name;
     String connection_id;
     double waiting_task_time_ms = 0;
+    grpc::Alarm alarm{};
 };
 } // namespace DB

dbms/src/Flash/Mpp/MPPTaskManager.cpp

@@ -54,7 +54,7 @@ void MPPGatherTaskSet::cancelAlarmsBySenderTaskId(const MPPTaskId & task_id)
     if (alarm_it != alarms.end())
     {
         for (auto & alarm : alarm_it->second)
-            alarm.second.Cancel();
+            alarm.second.get().Cancel();
         alarms.erase(alarm_it);
     }
 }
@@ -185,10 +185,11 @@ std::pair<MPPTunnelPtr, String> MPPTaskManager::findAsyncTunnel(
                     context.getSettingsRef().auto_spill_check_min_interval_ms.get());
             if (gather_task_set == nullptr)
                 gather_task_set = query->addMPPGatherTaskSet(id.gather_id);
-            auto & alarm = gather_task_set->alarms[sender_task_id][receiver_task_id];
+            auto & alarm = call_data->getAlarm();
             call_data->setCallStateAndUpdateMetrics(
                 EstablishCallData::WAIT_TUNNEL,
                 GET_METRIC(tiflash_establish_calldata_count, type_wait_tunnel_calldata));
+            gather_task_set->alarms[sender_task_id].emplace(receiver_task_id, std::ref(alarm));
             if likely (cq != nullptr)
             {
                 alarm.Set(cq, Clock::now() + std::chrono::seconds(10), call_data->asGRPCKickTag());
@@ -220,7 +221,6 @@ std::pair<MPPTunnelPtr, String> MPPTaskManager::findAsyncTunnel(
         }
     }
     /// don't need to delete the alarm here because registerMPPTask will delete all the related alarm
-
     return task->getTunnel(request);
 }
 
@@ -315,7 +315,7 @@ void MPPTaskManager::abortMPPGather(const MPPGatherId & gather_id, const String
         for (auto & alarms_per_task : gather_task_set->alarms)
         {
             for (auto & alarm : alarms_per_task.second)
-                alarm.second.Cancel();
+                alarm.second.get().Cancel();
         }
         gather_task_set->alarms.clear();
         if (!gather_task_set->hasMPPTask())

dbms/src/Flash/Mpp/MPPTaskManager.h

@@ -24,6 +24,7 @@
 
 #include <chrono>
 #include <condition_variable>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <unordered_map>
@@ -42,7 +43,7 @@ struct MPPGatherTaskSet
     State state = Normal;
     String error_message;
     /// <sender_task_id, <receiver_task_id, alarm>>
-    std::unordered_map<Int64, std::unordered_map<Int64, grpc::Alarm>> alarms;
+    std::unordered_map<Int64, std::unordered_map<Int64, std::reference_wrapper<grpc::Alarm>>> alarms;
     /// only used in scheduler
     std::queue<MPPTaskId> waiting_tasks;
     bool isInNormalState() const { return state == Normal; }