From 58ea5563a5a235c62bb20dc5ad432cd2ea9f2d82 Mon Sep 17 00:00:00 2001
From: Liu Jicong <liujicong@qq.com>
Date: Tue, 16 Aug 2022 19:46:42 +0800
Subject: [PATCH] fix(tmq): rebalance

---
 source/dnode/mnode/impl/src/mndSubscribe.c | 32 +++++++++++++++-------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/source/dnode/mnode/impl/src/mndSubscribe.c b/source/dnode/mnode/impl/src/mndSubscribe.c
index 3f310ee9c0..d517a1a1fb 100644
--- a/source/dnode/mnode/impl/src/mndSubscribe.c
+++ b/source/dnode/mnode/impl/src/mndSubscribe.c
@@ -356,7 +356,7 @@ static int32_t mndDoRebalance(SMnode *pMnode, const SMqRebInputObj *pInput, SMqR
       taosArrayPush(pConsumerEp->vgs, &pRebVg->pVgEp);
       pRebVg->newConsumerId = pConsumerEp->consumerId;
       taosArrayPush(pOutput->rebVgs, pRebVg);
-      mInfo("mq rebalance: add vgId:%d to consumer:%" PRId64 ",(second scan)", pRebVg->pVgEp->vgId,
+      mInfo("mq rebalance: add vgId:%d to consumer:%" PRId64 " (second scan) (not enough)", pRebVg->pVgEp->vgId,
             pConsumerEp->consumerId);
     }
   }
@@ -365,22 +365,34 @@ static int32_t mndDoRebalance(SMnode *pMnode, const SMqRebInputObj *pInput, SMqR
   if (taosHashGetSize(pOutput->pSub->consumerHash) != 0) {
     // if has consumer, assign all left vg
     while (1) {
+      SMqConsumerEp *pConsumerEp = NULL;
       pRemovedIter = taosHashIterate(pHash, pRemovedIter);
-      if (pRemovedIter == NULL) break;
-      pIter = taosHashIterate(pOutput->pSub->consumerHash, pIter);
-      ASSERT(pIter);
+      if (pRemovedIter == NULL) {
+        if (pIter != NULL) {
+          taosHashCancelIterate(pOutput->pSub->consumerHash, pIter);
+          pIter = NULL;
+        }
+        break;
+      }
+      while (1) {
+        pIter = taosHashIterate(pOutput->pSub->consumerHash, pIter);
+        ASSERT(pIter);
+        pConsumerEp = (SMqConsumerEp *)pIter;
+        ASSERT(pConsumerEp->consumerId > 0);
+        if (taosArrayGetSize(pConsumerEp->vgs) == minVgCnt + 1) {
+          continue;
+        }
+      }
       pRebVg = (SMqRebOutputVg *)pRemovedIter;
-      SMqConsumerEp *pConsumerEp = (SMqConsumerEp *)pIter;
-      ASSERT(pConsumerEp->consumerId > 0);
       taosArrayPush(pConsumerEp->vgs, &pRebVg->pVgEp);
       pRebVg->newConsumerId = pConsumerEp->consumerId;
       if (pRebVg->newConsumerId == pRebVg->oldConsumerId) {
-        mInfo("mq rebalance: skip vg %d for same consumer:%" PRId64 ",(second scan)", pRebVg->pVgEp->vgId,
+        mInfo("mq rebalance: skip vg %d for same consumer:%" PRId64 " (second scan)", pRebVg->pVgEp->vgId,
               pConsumerEp->consumerId);
         continue;
       }
       taosArrayPush(pOutput->rebVgs, pRebVg);
-      mInfo("mq rebalance: add vgId:%d to consumer:%" PRId64 ",(second scan)", pRebVg->pVgEp->vgId,
+      mInfo("mq rebalance: add vgId:%d to consumer:%" PRId64 " (second scan) (unassigned)", pRebVg->pVgEp->vgId,
             pConsumerEp->consumerId);
     }
   } else {
@@ -571,7 +583,7 @@ static int32_t mndProcessRebalanceReq(SRpcMsg *pMsg) {
       SMqTopicObj *pTopic = mndAcquireTopic(pMnode, topic);
       /*ASSERT(pTopic);*/
       if (pTopic == NULL) {
-        mError("rebalance %s failed since topic %s was dropped, abort", pRebInfo->key, topic);
+        mError("mq rebalance %s failed since topic %s not exist, abort", pRebInfo->key, topic);
         continue;
       }
       taosRLockLatch(&pTopic->lock);
@@ -601,7 +613,7 @@ static int32_t mndProcessRebalanceReq(SRpcMsg *pMsg) {
 
     // TODO replace assert with error check
     if (mndPersistRebResult(pMnode, pMsg, &rebOutput) < 0) {
-      mError("persist rebalance output error, possibly vnode splitted or dropped");
+      mError("mq rebalance persist rebalance output error, possibly vnode splitted or dropped");
     }
     taosArrayDestroy(pRebInfo->lostConsumers);
     taosArrayDestroy(pRebInfo->newConsumers);
-- 
GitLab