From f78a31806c4a8c0e35bba2e0a873a0fdd3211848 Mon Sep 17 00:00:00 2001
From: Liu Jicong <liujicong@qq.com>
Date: Thu, 28 Apr 2022 15:19:06 +0800
Subject: [PATCH] enh(tmq): add rebalance global lock

---
 include/common/tmsg.h                      |  1 -
 source/dnode/mnode/impl/inc/mndConsumer.h  |  5 +++++
 source/dnode/mnode/impl/inc/mndTrans.h     |  4 ++--
 source/dnode/mnode/impl/src/mndConsumer.c  | 19 ++++++++++++++-----
 source/dnode/mnode/impl/src/mndSubscribe.c |  7 +++++--
 source/dnode/mnode/impl/src/mndTrans.c     |  5 +++++
 6 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/include/common/tmsg.h b/include/common/tmsg.h
index 89f7cd6c39..145fc90a77 100644
--- a/include/common/tmsg.h
+++ b/include/common/tmsg.h
@@ -1474,7 +1474,6 @@ _err:
 // this message is sent from mnode to mnode(read thread to write thread), so there is no need for serialization or
 // deserialization
 typedef struct {
-  int8_t*   mqInReb;
   SHashObj* rebSubHash;  // SHashObj<key, SMqRebSubscribe>
 } SMqDoRebalanceMsg;
 
diff --git a/source/dnode/mnode/impl/inc/mndConsumer.h b/source/dnode/mnode/impl/inc/mndConsumer.h
index 19e202dddc..a8bfe91cbf 100644
--- a/source/dnode/mnode/impl/inc/mndConsumer.h
+++ b/source/dnode/mnode/impl/inc/mndConsumer.h
@@ -44,6 +44,11 @@ SSdbRow *mndConsumerActionDecode(SSdbRaw *pRaw);
 
 int32_t mndSetConsumerCommitLogs(SMnode *pMnode, STrans *pTrans, SMqConsumerObj *pConsumer);
 
+bool mndRebTryStart();
+void mndRebEnd();
+void mndRebCntInc();
+void mndRebCntDec();
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/source/dnode/mnode/impl/inc/mndTrans.h b/source/dnode/mnode/impl/inc/mndTrans.h
index 5b5aff7c86..7644ec3c4c 100644
--- a/source/dnode/mnode/impl/inc/mndTrans.h
+++ b/source/dnode/mnode/impl/inc/mndTrans.h
@@ -36,8 +36,8 @@ typedef struct {
 typedef enum {
   TEST_TRANS_START_FUNC = 1,
   TEST_TRANS_STOP_FUNC = 2,
-  CONSUME_TRANS_START_FUNC = 3,
-  CONSUME_TRANS_STOP_FUNC = 4,
+  MQ_REB_TRANS_START_FUNC = 3,
+  MQ_REB_TRANS_STOP_FUNC = 4,
 } ETrnFuncType;
 
 typedef void (*TransCbFp)(SMnode *pMnode, void *param, int32_t paramLen);
diff --git a/source/dnode/mnode/impl/src/mndConsumer.c b/source/dnode/mnode/impl/src/mndConsumer.c
index 025f61dc87..b17772bdb2 100644
--- a/source/dnode/mnode/impl/src/mndConsumer.c
+++ b/source/dnode/mnode/impl/src/mndConsumer.c
@@ -35,7 +35,7 @@
 
 #define MND_CONSUMER_LOST_HB_CNT 3
 
-static int8_t mqInRebFlag = 0;
+static int8_t mqRebLock = 0;
 
 static const char *mndConsumerStatusName(int status);
 
@@ -75,6 +75,17 @@ int32_t mndInitConsumer(SMnode *pMnode) {
 
 void mndCleanupConsumer(SMnode *pMnode) {}
 
+bool mndRebTryStart() {
+  int8_t old = atomic_val_compare_exchange_8(&mqRebLock, 0, 1);
+  return old == 0;
+}
+
+void mndRebEnd() { atomic_sub_fetch_8(&mqRebLock, 1); }
+
+void mndRebCntInc() { atomic_add_fetch_8(&mqRebLock, 1); }
+
+void mndRebCntDec() { atomic_sub_fetch_8(&mqRebLock, 1); }
+
 static int32_t mndProcessConsumerLostMsg(SNodeMsg *pMsg) {
   SMnode             *pMnode = pMsg->pNode;
   SMqConsumerLostMsg *pLostMsg = pMsg->rpcMsg.pCont;
@@ -143,8 +154,7 @@ static int32_t mndProcessMqTimerMsg(SNodeMsg *pMsg) {
   void           *pIter = NULL;
 
   // rebalance cannot be parallel
-  int8_t old = atomic_val_compare_exchange_8(&mqInRebFlag, 0, 1);
-  if (old != 0) {
+  if (!mndRebTryStart()) {
     mInfo("mq rebalance already in progress, do nothing");
     return 0;
   }
@@ -152,7 +162,6 @@ static int32_t mndProcessMqTimerMsg(SNodeMsg *pMsg) {
   SMqDoRebalanceMsg *pRebMsg = rpcMallocCont(sizeof(SMqDoRebalanceMsg));
   pRebMsg->rebSubHash = taosHashInit(64, MurmurHash3_32, true, HASH_NO_LOCK);
   // TODO set cleanfp
-  pRebMsg->mqInReb = &mqInRebFlag;
 
   // iterate all consumers, find all modification
   while (1) {
@@ -223,7 +232,7 @@ static int32_t mndProcessMqTimerMsg(SNodeMsg *pMsg) {
     taosHashCleanup(pRebMsg->rebSubHash);
     rpcFreeCont(pRebMsg);
     mTrace("mq rebalance finished, no modification");
-    atomic_store_8(&mqInRebFlag, 0);
+    mndRebEnd();
   }
   return 0;
 }
diff --git a/source/dnode/mnode/impl/src/mndSubscribe.c b/source/dnode/mnode/impl/src/mndSubscribe.c
index aeecab34cb..2b3af85066 100644
--- a/source/dnode/mnode/impl/src/mndSubscribe.c
+++ b/source/dnode/mnode/impl/src/mndSubscribe.c
@@ -452,7 +452,10 @@ static int32_t mndPersistRebResult(SMnode *pMnode, SNodeMsg *pMsg, const SMqRebO
   }
   // 4. TODO commit log: modification log
 
-  // 5. execution
+  // 5. set cb
+  mndTransSetCb(pTrans, MQ_REB_TRANS_START_FUNC, MQ_REB_TRANS_STOP_FUNC, NULL, 0);
+
+  // 6. execution
   if (mndTransPrepare(pMnode, pTrans) != 0) goto REB_FAIL;
 
   mndTransDrop(pTrans);
@@ -518,9 +521,9 @@ static int32_t mndProcessRebalanceReq(SNodeMsg *pMsg) {
   }
 
   // reset flag
-  atomic_store_8(pReq->mqInReb, 0);
   mInfo("mq rebalance completed successfully");
   taosHashCleanup(pReq->rebSubHash);
+  mndRebEnd();
 
   return 0;
 }
diff --git a/source/dnode/mnode/impl/src/mndTrans.c b/source/dnode/mnode/impl/src/mndTrans.c
index c08b0f6db9..2d10c4a7a5 100644
--- a/source/dnode/mnode/impl/src/mndTrans.c
+++ b/source/dnode/mnode/impl/src/mndTrans.c
@@ -16,6 +16,7 @@
 #define _DEFAULT_SOURCE
 #include "mndTrans.h"
 #include "mndAuth.h"
+#include "mndConsumer.h"
 #include "mndDb.h"
 #include "mndShow.h"
 #include "mndSync.h"
@@ -442,6 +443,10 @@ static TransCbFp mndTransGetCbFp(ETrnFuncType ftype) {
       return mndTransTestStartFunc;
     case TEST_TRANS_STOP_FUNC:
       return mndTransTestStopFunc;
+    case MQ_REB_TRANS_START_FUNC:
+      return mndRebCntInc;
+    case MQ_REB_TRANS_STOP_FUNC:
+      return mndRebCntDec;
     default:
       return NULL;
   }