From 12687355c07b28c21e4ef628482a1a26858ac351 Mon Sep 17 00:00:00 2001
From: Jeff Tao <jhtao@taosdata.com>
Date: Tue, 6 Oct 2020 11:46:33 +0000
Subject: [PATCH] TD-1632

---
 src/connector/go      |  2 +-
 src/rpc/src/rpcMain.c | 40 +++++++++++++++++++---------------------
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/src/connector/go b/src/connector/go
index 567b7b12f3..8c58c512b6 160000
--- a/src/connector/go
+++ b/src/connector/go
@@ -1 +1 @@
-Subproject commit 567b7b12f3fd2775c718d284beffc8c38dd6c219
+Subproject commit 8c58c512b6acda8bcdfa48fdc7140227b5221766
diff --git a/src/rpc/src/rpcMain.c b/src/rpc/src/rpcMain.c
index 6a5d3b079a..b86b95b858 100644
--- a/src/rpc/src/rpcMain.c
+++ b/src/rpc/src/rpcMain.c
@@ -881,6 +881,20 @@ static int rpcProcessRspHead(SRpcConn *pConn, SRpcHead *pHead) {
   pConn->outType = 0;
   pConn->pReqMsg = NULL;
   pConn->reqMsgLen = 0;
+  SRpcReqContext *pContext = pConn->pContext;
+
+  if (pHead->code == TSDB_CODE_RPC_REDIRECT) { 
+    if (rpcContLenFromMsg(pHead->msgLen) < sizeof(SRpcEpSet)) {
+      // if EpSet is not included in the msg, treat it as NOT_READY
+      pHead->code = TSDB_CODE_RPC_NOT_READY; 
+    } else {
+      pContext->redirect++;
+      if (pContext->redirect > TSDB_MAX_REPLICA) {
+        pHead->code = TSDB_CODE_RPC_NETWORK_UNAVAIL; 
+        tWarn("%s, too many redirects, quit", pConn->info);
+      }
+    }
+  } 
 
   return TSDB_CODE_SUCCESS;
 }
@@ -950,12 +964,6 @@ static SRpcConn *rpcProcessMsgHead(SRpcInfo *pRpc, SRecvInfo *pRecv, SRpcReqCont
         SRpcReqContext *pContext = pConn->pContext;
         *ppContext = pContext;
         pConn->pContext = NULL;
-        pConn->pReqMsg = NULL;
-
-        // for UDP, port may be changed by server, the port in epSet shall be used for cache
-        if (pHead->code != TSDB_CODE_RPC_TOO_SLOW) {
-          rpcAddConnIntoCache(pRpc->pCache, pConn, pConn->peerFqdn, pContext->epSet.port[pContext->epSet.inUse], pConn->connType);    
-        } 
       }
     }
   }
@@ -1083,9 +1091,9 @@ static void rpcProcessIncomingMsg(SRpcConn *pConn, SRpcHead *pHead, SRpcReqConte
   rpcMsg.pCont = pHead->content;
   rpcMsg.msgType = pHead->msgType;
   rpcMsg.code = pHead->code; 
-  rpcMsg.ahandle = pConn->ahandle;
    
   if ( rpcIsReq(pHead->msgType) ) {
+    rpcMsg.ahandle = pConn->ahandle;
     if (rpcMsg.contLen > 0) {
       rpcMsg.handle = pConn;
       rpcAddRef(pRpc);  // add the refCount for requests
@@ -1103,25 +1111,15 @@ static void rpcProcessIncomingMsg(SRpcConn *pConn, SRpcHead *pHead, SRpcReqConte
   } else {
     // it's a response
     rpcMsg.handle = pContext;
+    rpcMsg.ahandle = pContext->ahandle;
 
     // for UDP, port may be changed by server, the port in epSet shall be used for cache
-    if (pHead->code == TSDB_CODE_RPC_TOO_SLOW) {
+    if (pHead->code != TSDB_CODE_RPC_TOO_SLOW) {
+      rpcAddConnIntoCache(pRpc->pCache, pConn, pConn->peerFqdn, pContext->epSet.port[pContext->epSet.inUse], pConn->connType);    
+    } else {
       rpcCloseConn(pConn);
     }
 
-    if (pHead->code == TSDB_CODE_RPC_REDIRECT) { 
-      if (rpcMsg.contLen < sizeof(SRpcEpSet)) {
-        // if EpSet is not included in the msg, treat it as NOT_READY
-        pHead->code = TSDB_CODE_RPC_NOT_READY; 
-      } else {
-        pContext->redirect++;
-        if (pContext->redirect > TSDB_MAX_REPLICA) {
-          pHead->code = TSDB_CODE_RPC_NETWORK_UNAVAIL; 
-          tWarn("%s, too many redirects, quit", pConn->info);
-        }
-      }
-    }
-
     if (pHead->code == TSDB_CODE_RPC_REDIRECT) {
       pContext->numOfTry = 0;
       SRpcEpSet *pEpSet = (SRpcEpSet*)pHead->content;
-- 
GitLab