From 800e7c4e7acd19ec33c2524d1fdde61e8ce64495 Mon Sep 17 00:00:00 2001
From: kailixu <klxu@taosdata.com>
Date: Mon, 19 Aug 2024 17:27:15 +0800
Subject: [PATCH] fix: memory leak of geos

---
 source/client/src/clientMain.c           |  2 +
 source/dnode/mgmt/node_mgmt/src/dmMgmt.c |  2 +
 source/util/src/tgeosctx.c               | 92 +++++++++++++++++-------
 source/util/src/tsched.c                 |  2 +-
 source/util/src/tworker.c                |  4 +-
 5 files changed, 75 insertions(+), 27 deletions(-)

diff --git a/source/client/src/clientMain.c b/source/client/src/clientMain.c
index 12702a93f3..a403f9d1c2 100644
--- a/source/client/src/clientMain.c
+++ b/source/client/src/clientMain.c
@@ -23,6 +23,7 @@
 #include "query.h"
 #include "scheduler.h"
 #include "tdatablock.h"
+#include "tgeosctx.h"
 #include "tglobal.h"
 #include "tmsg.h"
 #include "tref.h"
@@ -94,6 +95,7 @@ void taos_cleanup(void) {
   tmqMgmtClose();
 
   DestroyRegexCache();
+  destroyThreadLocalGeosCtx();
 
   tscInfo("all local resources released");
   taosCleanupCfg();
diff --git a/source/dnode/mgmt/node_mgmt/src/dmMgmt.c b/source/dnode/mgmt/node_mgmt/src/dmMgmt.c
index fdce9fd4c9..1d62d4bd90 100644
--- a/source/dnode/mgmt/node_mgmt/src/dmMgmt.c
+++ b/source/dnode/mgmt/node_mgmt/src/dmMgmt.c
@@ -19,6 +19,7 @@
 #include "index.h"
 #include "qworker.h"
 #include "tcompression.h"
+#include "tgeosctx.h"
 #include "tglobal.h"
 #include "tgrant.h"
 #include "tstream.h"
@@ -121,6 +122,7 @@ void dmCleanupDnode(SDnode *pDnode) {
   streamMetaCleanup();
   indexCleanup();
   taosConvDestroy();
+  destroyThreadLocalGeosCtx();
 
   // compress destroy
   tsCompressExit();
diff --git a/source/util/src/tgeosctx.c b/source/util/src/tgeosctx.c
index a05734c911..47d5cc992b 100644
--- a/source/util/src/tgeosctx.c
+++ b/source/util/src/tgeosctx.c
@@ -15,38 +15,82 @@
 
 #include "tgeosctx.h"
 #include "tdef.h"
+#include "tlockfree.h"
+#include "tlog.h"
 
-static threadlocal SGeosContext tlGeosCtx = {0};
+typedef struct {
+  SGeosContext *pool;
+  int32_t       capacity;
+  int32_t       size;
+  SRWLatch      lock;
+} SGeosContextPool;
 
-SGeosContext* getThreadLocalGeosCtx() { return &tlGeosCtx; }
+static SGeosContextPool sGeosPool = {0};
 
-void destroyThreadLocalGeosCtx() {
-  if (tlGeosCtx.WKTReader) {
-    GEOSWKTReader_destroy_r(tlGeosCtx.handle, tlGeosCtx.WKTReader);
-    tlGeosCtx.WKTReader = NULL;
+static threadlocal SGeosContext *tlGeosCtx = NULL;
+
+SGeosContext *getThreadLocalGeosCtx() {
+  if (tlGeosCtx) return tlGeosCtx;
+
+  taosWLockLatch(&sGeosPool.lock);
+  if (sGeosPool.size >= sGeosPool.capacity) {
+    sGeosPool.capacity += 64;
+    void *tmp = taosMemoryRealloc(sGeosPool.pool, sGeosPool.capacity * sizeof(SGeosContext));
+    if (!tmp) {
+      taosWUnLockLatch(&sGeosPool.lock);
+      terrno = TSDB_CODE_OUT_OF_MEMORY;
+      return NULL;
+    }
+    sGeosPool.pool = tmp;
+    TAOS_MEMSET(sGeosPool.pool + sGeosPool.size, 0, (sGeosPool.capacity - sGeosPool.size) * sizeof(SGeosContext));
   }
+  tlGeosCtx = sGeosPool.pool + sGeosPool.size;
+  ++sGeosPool.size;
+  taosWUnLockLatch(&sGeosPool.lock);
 
-  if (tlGeosCtx.WKTWriter) {
-    GEOSWKTWriter_destroy_r(tlGeosCtx.handle, tlGeosCtx.WKTWriter);
-    tlGeosCtx.WKTWriter = NULL;
-  }
+  return tlGeosCtx;
+}
 
-  if (tlGeosCtx.WKBReader) {
-    GEOSWKBReader_destroy_r(tlGeosCtx.handle, tlGeosCtx.WKBReader);
-    tlGeosCtx.WKBReader = NULL;
-  }
+static void destroyGeosCtx(SGeosContext *pCtx) {
+  if (pCtx) {
+    if (pCtx->WKTReader) {
+      GEOSWKTReader_destroy_r(pCtx->handle, pCtx->WKTReader);
+      pCtx->WKTReader = NULL;
+    }
 
-  if (tlGeosCtx.WKBWriter) {
-    GEOSWKBWriter_destroy_r(tlGeosCtx.handle, tlGeosCtx.WKBWriter);
-    tlGeosCtx.WKBWriter = NULL;
-  }
+    if (pCtx->WKTWriter) {
+      GEOSWKTWriter_destroy_r(pCtx->handle, pCtx->WKTWriter);
+      pCtx->WKTWriter = NULL;
+    }
 
-  if (tlGeosCtx.WKTRegex) {
-    destroyRegexes(tlGeosCtx.WKTRegex, tlGeosCtx.WKTMatchData);
-  }
+    if (pCtx->WKBReader) {
+      GEOSWKBReader_destroy_r(pCtx->handle, pCtx->WKBReader);
+      pCtx->WKBReader = NULL;
+    }
 
-  if (tlGeosCtx.handle) {
-    GEOS_finish_r(tlGeosCtx.handle);
-    tlGeosCtx.handle = NULL;
+    if (pCtx->WKBWriter) {
+      GEOSWKBWriter_destroy_r(pCtx->handle, pCtx->WKBWriter);
+      pCtx->WKBWriter = NULL;
+    }
+
+    if (pCtx->WKTRegex) {
+      destroyRegexes(pCtx->WKTRegex, pCtx->WKTMatchData);
+      pCtx->WKTRegex = NULL;
+      pCtx->WKTMatchData = NULL;
+    }
+
+    if (pCtx->handle) {
+      GEOS_finish_r(pCtx->handle);
+      pCtx->handle = NULL;
+    }
   }
 }
+
+void destroyThreadLocalGeosCtx() {
+  uInfo("geos ctx is cleaned up");
+  if (!sGeosPool.pool) return;
+  for (int32_t i = 0; i < sGeosPool.size; ++i) {
+    destroyGeosCtx(sGeosPool.pool + i);
+  }
+  taosMemoryFreeClear(sGeosPool.pool);
+}
\ No newline at end of file
diff --git a/source/util/src/tsched.c b/source/util/src/tsched.c
index 6779e8dee5..509dba0890 100644
--- a/source/util/src/tsched.c
+++ b/source/util/src/tsched.c
@@ -178,7 +178,7 @@ void *taosProcessSchedQueue(void *scheduler) {
       (*(msg.tfp))(msg.ahandle, msg.thandle);
   }
 
-  destroyThreadLocalGeosCtx();
+  // destroyThreadLocalGeosCtx();
 
   return NULL;
 }
diff --git a/source/util/src/tworker.c b/source/util/src/tworker.c
index b2064d6787..2da1abed78 100644
--- a/source/util/src/tworker.c
+++ b/source/util/src/tworker.c
@@ -105,7 +105,7 @@ static void *tQWorkerThreadFp(SQueueWorker *worker) {
     taosUpdateItemSize(qinfo.queue, 1);
   }
 
-  destroyThreadLocalGeosCtx();
+  // destroyThreadLocalGeosCtx();
   DestoryThreadLocalRegComp();
 
   return NULL;
@@ -665,7 +665,7 @@ static void *tQueryAutoQWorkerThreadFp(SQueryAutoQWorker *worker) {
     }
   }
 
-  destroyThreadLocalGeosCtx();
+  // destroyThreadLocalGeosCtx();
   DestoryThreadLocalRegComp();
 
   return NULL;