From 8c4747e81d63cce726c8b0c171e187040518e5de Mon Sep 17 00:00:00 2001
From: slguan <slguan@taosdata.com>
Date: Mon, 6 Apr 2020 23:52:43 +0800
Subject: [PATCH] [TD-114] add balance interface, so it can create database
 with replica 2 option

---
 src/dnode/CMakeLists.txt    |  4 ++++
 src/mnode/inc/mgmtBalance.h |  2 +-
 src/mnode/src/mgmtBalance.c | 33 ++++++++++++++++++++++++++++++---
 src/mnode/src/mgmtDnode.c   |  2 +-
 src/mnode/src/mgmtVgroup.c  | 12 +++++++-----
 5 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/src/dnode/CMakeLists.txt b/src/dnode/CMakeLists.txt
index f0be3f9044..a41b3734eb 100644
--- a/src/dnode/CMakeLists.txt
+++ b/src/dnode/CMakeLists.txt
@@ -25,6 +25,10 @@ IF ((TD_LINUX_64) OR (TD_LINUX_32 AND TD_ARM))
   IF (TD_CLUSTER)
     TARGET_LINK_LIBRARIES(taosd cluster)
   ENDIF ()
+  
+  IF (TD_VPEER)
+    TARGET_LINK_LIBRARIES(taosd balance)
+  ENDIF ()
 
   SET(PREPARE_ENV_CMD "prepare_env_cmd")
   SET(PREPARE_ENV_TARGET "prepare_env_target")
diff --git a/src/mnode/inc/mgmtBalance.h b/src/mnode/inc/mgmtBalance.h
index 401074d171..05f2ed94a7 100644
--- a/src/mnode/inc/mgmtBalance.h
+++ b/src/mnode/inc/mgmtBalance.h
@@ -23,7 +23,7 @@ extern "C" {
 
 int32_t mgmtInitBalance();
 void    mgmtCleanupBalance();
-void    mgmtStartBalanceTimer(int32_t afterMs) ;
+void    mgmtBalanceNotify() ;
 int32_t mgmtAllocVnodes(SVgObj *pVgroup);
 
 #ifdef __cplusplus
diff --git a/src/mnode/src/mgmtBalance.c b/src/mnode/src/mgmtBalance.c
index 3a7ae8dcde..e3d7f61708 100644
--- a/src/mnode/src/mgmtBalance.c
+++ b/src/mnode/src/mgmtBalance.c
@@ -18,11 +18,37 @@
 #include "mgmtBalance.h"
 #include "mgmtDnode.h"
 
-int32_t mgmtInitBalance() { return 0; }
-void    mgmtCleanupBalance() {}
-void    mgmtStartBalanceTimer(int32_t afterMs) {}
+extern int32_t balanceInit();
+extern void    balanceCleanUp();
+extern void    balanceNotify();
+extern int32_t balanceAllocVnodes(SVgObj *pVgroup);
+
+int32_t mgmtInitBalance() {
+#ifdef _VPEER
+  return balanceInit();
+#else
+  return 0;
+#endif
+}
+
+void mgmtCleanupBalance() {
+#ifdef _VPEER
+  balanceCleanUp();
+#endif
+}
+
+void mgmtBalanceNotify() {
+#ifdef _VPEER
+  balanceNotify();
+#else
+  return 0;
+#endif
+}
 
 int32_t mgmtAllocVnodes(SVgObj *pVgroup) {
+#ifdef _VPEER
+  return balanceAllocVnodes(pVgroup);
+#else
   void *     pNode = NULL;
   SDnodeObj *pDnode = NULL;
   SDnodeObj *pSelDnode = NULL;
@@ -53,4 +79,5 @@ int32_t mgmtAllocVnodes(SVgObj *pVgroup) {
 
   mTrace("dnode:%d, alloc one vnode to vgroup, openVnodes:%d", pSelDnode->dnodeId, pSelDnode->openVnodes);
   return TSDB_CODE_SUCCESS;
+#endif  
 }
diff --git a/src/mnode/src/mgmtDnode.c b/src/mnode/src/mgmtDnode.c
index 14e23bad9f..ada0bce2e9 100644
--- a/src/mnode/src/mgmtDnode.c
+++ b/src/mnode/src/mgmtDnode.c
@@ -251,7 +251,7 @@ void mgmtProcessDnodeStatusMsg(SRpcMsg *rpcMsg) {
   if (pDnode->status != TSDB_DN_STATUS_READY) {
     mTrace("dnode:%d, from offline to online", pDnode->dnodeId);
     pDnode->status = TSDB_DN_STATUS_READY;
-    mgmtStartBalanceTimer(200);
+    mgmtBalanceNotify();
   }
 
   mgmtDecDnodeRef(pDnode);
diff --git a/src/mnode/src/mgmtVgroup.c b/src/mnode/src/mgmtVgroup.c
index 9f6edb17e5..bbf960203a 100644
--- a/src/mnode/src/mgmtVgroup.c
+++ b/src/mnode/src/mgmtVgroup.c
@@ -32,7 +32,7 @@
 #include "mgmtVgroup.h"
 
 void   *tsVgroupSdb = NULL;
-static int32_t tsVgUpdateSize = 0;
+int32_t tsVgUpdateSize = 0;
 
 static int32_t mgmtGetVgroupMeta(STableMetaMsg *pMeta, SShowObj *pShow, void *pConn);
 static int32_t mgmtRetrieveVgroups(SShowObj *pShow, char *data, int32_t rows, void *pConn);
@@ -93,10 +93,12 @@ static int32_t mgmtVgroupActionInsert(SSdbOperDesc *pOper) {
 
   for (int32_t i = 0; i < pVgroup->numOfVnodes; ++i) {
     SDnodeObj *pDnode = mgmtGetDnode(pVgroup->vnodeGid[i].dnodeId);
-    pVgroup->vnodeGid[i].privateIp = pDnode->privateIp;
-    pVgroup->vnodeGid[i].publicIp = pDnode->publicIp;
-    atomic_add_fetch_32(&pDnode->openVnodes, 1);
-    mgmtDecDnodeRef(pDnode);
+    if (pDnode != NULL) {
+      pVgroup->vnodeGid[i].privateIp = pDnode->privateIp;
+      pVgroup->vnodeGid[i].publicIp = pDnode->publicIp;
+      atomic_add_fetch_32(&pDnode->openVnodes, 1);    
+      mgmtDecDnodeRef(pDnode);
+    }     
   }
 
   mgmtAddVgroupIntoDb(pVgroup);
-- 
GitLab