From 91782c5b99279e2b05846bbb9dfa8aad9ec3a673 Mon Sep 17 00:00:00 2001
From: yihaoDeng <yhdeng@taosdata.com>
Date: Sun, 18 Oct 2020 01:04:17 +0000
Subject: [PATCH] TD-1733

---
 src/query/inc/qExecutor.h |  7 +++++++
 src/tsdb/src/tsdbRead.c   | 27 +++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/query/inc/qExecutor.h b/src/query/inc/qExecutor.h
index 5d570821cb..33b724e434 100644
--- a/src/query/inc/qExecutor.h
+++ b/src/query/inc/qExecutor.h
@@ -185,11 +185,18 @@ enum {
   QUERY_RESULT_READY     = 2,
 };
 
+typedef struct SMemRef {
+  int32_t ref; 
+  void *mem;
+  void *imem;
+} SMemRef;
+
 typedef struct SQInfo {
   void*            signature;
   int32_t          code;   // error code to returned to client
   int64_t          owner; // if it is in execution
   void*            tsdb;
+  SMemRef          memRef; 
   int32_t          vgId;
   STableGroupInfo  tableGroupInfo;       // table <tid, last_key> list  SArray<STableKeyInfo>
   STableGroupInfo  tableqinfoGroupInfo;  // this is a group array list, including SArray<STableQueryInfo*> structure
diff --git a/src/tsdb/src/tsdbRead.c b/src/tsdb/src/tsdbRead.c
index b8ccbd903c..ac6c2e0c5a 100644
--- a/src/tsdb/src/tsdbRead.c
+++ b/src/tsdb/src/tsdbRead.c
@@ -20,6 +20,7 @@
 #include "exception.h"
 
 #include "../../query/inc/qAst.h"  // todo move to common module
+#include "../../query/inc/qExecutor.h"  // todo move to common module
 #include "tlosertree.h"
 #include "tsdb.h"
 #include "tsdbMain.h"
@@ -143,6 +144,7 @@ static int     tsdbReadRowsFromCache(STableCheckInfo* pCheckInfo, TSKEY maxKey,
                                      STsdbQueryHandle* pQueryHandle);
 static int     tsdbCheckInfoCompar(const void* key1, const void* key2);
 
+
 static void tsdbInitDataBlockLoadInfo(SDataBlockLoadInfo* pBlockLoadInfo) {
   pBlockLoadInfo->slot = -1;
   pBlockLoadInfo->tid = -1;
@@ -182,6 +184,27 @@ static SArray* getDefaultLoadColumns(STsdbQueryHandle* pQueryHandle, bool loadTS
   return pLocalIdList;
 }
 
+static void tsdbMayTakeMemSnapshot(TsdbQueryHandleT pHandle) { 
+  STsdbQueryHandle* pSecQueryHandle = (STsdbQueryHandle*) pHandle;
+  SQInfo *pQInfo = (SQInfo *)(pSecQueryHandle->qinfo);
+
+  if (pQInfo->memRef.ref++ == 0) {
+    tsdbTakeMemSnapshot(pSecQueryHandle->pTsdb, &pSecQueryHandle->mem, &pSecQueryHandle->imem);
+    pQInfo->memRef.mem = pSecQueryHandle->mem;
+    pQInfo->memRef.imem = pSecQueryHandle->imem;
+  } else {
+    pSecQueryHandle->mem = (SMemTable *)(pQInfo->memRef.mem);
+    pSecQueryHandle->imem = (SMemTable *)(pQInfo->memRef.imem);
+  }
+}
+static void tsdbMayUnTakeMemSnapshot(TsdbQueryHandleT pHandle) {
+  STsdbQueryHandle* pSecQueryHandle = (STsdbQueryHandle*) pHandle; 
+  SQInfo *pQInfo = (SQInfo *)(pSecQueryHandle->qinfo);
+
+  if (--pQInfo->memRef.ref == 0) {
+    tsdbUnTakeMemSnapShot(pSecQueryHandle->pTsdb, pSecQueryHandle->mem, pSecQueryHandle->imem);
+  }
+}
 static SArray* createCheckInfoFromTableGroup(STsdbQueryHandle* pQueryHandle, STableGroupInfo* pGroupList, STsdbMeta* pMeta) {
   size_t sizeOfGroup = taosArrayGetSize(pGroupList->pGroupList);
   assert(sizeOfGroup >= 1 && pMeta != NULL);
@@ -270,7 +293,7 @@ static STsdbQueryHandle* tsdbQueryTablesImpl(TSDB_REPO_T* tsdb, STsdbQueryCond*
     goto out_of_memory;
   }
 
-  tsdbTakeMemSnapshot(pQueryHandle->pTsdb, &pQueryHandle->mem, &pQueryHandle->imem);
+  tsdbMayTakeMemSnapshot(pQueryHandle);
   assert(pCond != NULL && pCond->numOfCols > 0);
 
   if (ASCENDING_TRAVERSE(pCond->order)) {
@@ -2701,7 +2724,7 @@ void tsdbCleanupQueryHandle(TsdbQueryHandleT queryHandle) {
   taosTFree(pQueryHandle->statis);
 
   // todo check error
-  tsdbUnTakeMemSnapShot(pQueryHandle->pTsdb, pQueryHandle->mem, pQueryHandle->imem);
+  tsdbMayUnTakeMemSnapshot(pQueryHandle);
 
   tsdbDestroyHelper(&pQueryHandle->rhelper);
 
-- 
GitLab