From 8b2af4f09d68dfa183d2e3681632767dd58cd6cd Mon Sep 17 00:00:00 2001
From: kailixu <klxu@taosdata.com>
Date: Wed, 28 Aug 2024 19:12:23 +0800
Subject: [PATCH] fix: memory leak of geos

---
 include/os/osThread.h                         | 10 +++++-
 source/libs/geometry/test/geomIoFuncTest.cpp  | 24 ++++++++------
 .../geometry/test/geomRelationFuncTest.cpp    | 24 ++++++++------
 source/os/src/osThread.c                      | 32 +++++++++++++++++++
 source/util/src/tgeosctx.c                    |  6 ++++
 5 files changed, 77 insertions(+), 19 deletions(-)

diff --git a/include/os/osThread.h b/include/os/osThread.h
index 38c1b366f0..f39c01abee 100644
--- a/include/os/osThread.h
+++ b/include/os/osThread.h
@@ -22,8 +22,12 @@
 extern "C" {
 #endif
 
-#if defined(WINDOWS) && !defined(__USE_PTHREAD)
+#if defined(WINDOWS)
+#include <tlhelp32.h>
 #include <windows.h>
+#endif
+
+#if defined(WINDOWS) && !defined(__USE_PTHREAD)
 #define __USE_WIN_THREAD
 // https://learn.microsoft.com/en-us/windows/win32/winprog/using-the-windows-headers
 // #ifndef _WIN32_WINNT
@@ -275,6 +279,10 @@ int32_t  taosThreadSpinUnlock(TdThreadSpinlock *lock);
 void     taosThreadTestCancel(void);
 void     taosThreadClear(TdThread *thread);
 
+#ifdef WINDOWS
+bool taosThreadIsMain();
+#endif
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/source/libs/geometry/test/geomIoFuncTest.cpp b/source/libs/geometry/test/geomIoFuncTest.cpp
index 62fc90cb9e..2eff1bf0ea 100644
--- a/source/libs/geometry/test/geomIoFuncTest.cpp
+++ b/source/libs/geometry/test/geomIoFuncTest.cpp
@@ -320,20 +320,26 @@ void geomIoFuncTestAsTextFunction() {
   callAsTextWrapper2(TSDB_DATA_TYPE_GEOMETRY, strInput, valTypeArray, 1, TSDB_CODE_FUNC_FUNTION_PARA_VALUE);
 }
 
-static void *geomIoFuncTest(void *arg) {
+static void geomIoFuncTestImpl() {
   geomIoFuncTestMakePointFunctionTwoColumns();
   geomIoFuncTestMakePointFunctionConstant();
   geomIoFuncTestMakePointFunctionWithNull();
   geomIoFuncTestGeomFromTextFunction();
   geomIoFuncTestAsTextFunction();
-
-  return NULL;
 }
 
-TEST(GeomIoFuncTest, summary) {
-  TdThread threadId;
-  int32_t  ret = taosThreadCreate(&threadId, NULL, geomIoFuncTest, 0);
-  ASSERT_EQ(ret, 0);
-  ret = taosThreadJoin(threadId, NULL);
-  ASSERT_EQ(ret, 0);
+static void *geomIoFuncTestFunc(void *arg) {
+  geomIoFuncTestImpl();
+  return nullptr;
+}
+
+static void geomIoFuncTestInThread() {
+  TdThread thread;
+  ASSERT_EQ(taosThreadCreate(&thread, nullptr, geomIoFuncTestFunc, NULL), 0);
+  ASSERT_EQ(taosThreadJoin(thread, nullptr), 0);
+}
+
+TEST(threadGeomFuncTest, threadFuncTest) {
+  geomIoFuncTestImpl();
+  geomIoFuncTestInThread();
 }
\ No newline at end of file
diff --git a/source/libs/geometry/test/geomRelationFuncTest.cpp b/source/libs/geometry/test/geomRelationFuncTest.cpp
index 82bc3a4055..4563d443cb 100644
--- a/source/libs/geometry/test/geomRelationFuncTest.cpp
+++ b/source/libs/geometry/test/geomRelationFuncTest.cpp
@@ -238,21 +238,27 @@ void geomRelationFuncTestContainsProperlyFunction() {
   geomRelationFuncTest(containsProperlyFunction, expectedResults);
 }
 
-static void *geomRelationFuncTest(void *arg) {
+static void geomRelationFuncTestImpl() {
   geomRelationFuncTestIntersectsFunction();
   geomRelationFuncTestEqualsFunction();
   geomRelationFuncTestTouchesFunction();
   geomRelationFuncTestCoversFunction();
   geomRelationFuncTestContainsFunction();
   geomRelationFuncTestContainsProperlyFunction();
-
-  return NULL;
 }
 
-TEST(GeomRelationFuncTest, summary) {
-  TdThread threadId;
-  int32_t  ret = taosThreadCreate(&threadId, NULL, geomRelationFuncTest, 0);
-  ASSERT_EQ(ret, 0);
-  ret = taosThreadJoin(threadId, NULL);
-  ASSERT_EQ(ret, 0);
+static void *geomRelationFuncTestFunc(void *arg) {
+  geomRelationFuncTestImpl();
+  return nullptr;
 }
+
+static void geomRelationFuncTestInThread() {
+  TdThread thread;
+  ASSERT_EQ(taosThreadCreate(&thread, nullptr, geomRelationFuncTestFunc, NULL), 0);
+  ASSERT_EQ(taosThreadJoin(thread, nullptr), 0);
+}
+
+TEST(threadGeomRelationFuncTest, threadGeomRelationFuncTest) {
+  geomRelationFuncTestImpl();
+  geomRelationFuncTestInThread();
+}
\ No newline at end of file
diff --git a/source/os/src/osThread.c b/source/os/src/osThread.c
index 5a24e7775f..1fc48f145b 100644
--- a/source/os/src/osThread.c
+++ b/source/os/src/osThread.c
@@ -841,3 +841,35 @@ void taosThreadTestCancel(void) {
 void taosThreadClear(TdThread *thread) { 
   (void)memset(thread, 0, sizeof(TdThread)); 
 }
+
+#ifdef WINDOWS
+bool taosThreadIsMain() {
+  DWORD curProcessId = GetCurrentProcessId();
+  DWORD curThreadId = GetCurrentThreadId();
+  DWORD dwThreadId = -1;
+
+  HANDLE hThreadSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
+  if (hThreadSnapshot == INVALID_HANDLE_VALUE) {
+    return false;
+  }
+
+  THREADENTRY32 te32;
+  te32.dwSize = sizeof(THREADENTRY32);
+
+  if (!Thread32First(hThreadSnapshot, &te32)) {
+    CloseHandle(hThreadSnapshot);
+    return false;
+  }
+
+  do {
+    if (te32.th32OwnerProcessID == curProcessId) {
+      dwThreadId = te32.th32ThreadID;
+      break;
+    }
+  } while (Thread32Next(hThreadSnapshot, &te32));
+
+  CloseHandle(hThreadSnapshot);
+
+  return curThreadId == dwThreadId;
+}
+#endif
diff --git a/source/util/src/tgeosctx.c b/source/util/src/tgeosctx.c
index 9d8f64e01a..9c4f78684d 100644
--- a/source/util/src/tgeosctx.c
+++ b/source/util/src/tgeosctx.c
@@ -21,6 +21,10 @@ static TdThreadKey               tlGeosCtxKey = 0;
 static threadlocal SGeosContext *tlGeosCtx = NULL;
 
 static void destroyThreadLocalGeosCtx(void *param) {
+#ifdef WINDOWS
+  if (taosThreadIsMain()) return;
+#endif
+
   SGeosContext *pGeosCtx = (SGeosContext *)param;
   if (!pGeosCtx) {
     return;
@@ -72,9 +76,11 @@ int32_t getThreadLocalGeosCtx(SGeosContext **ppCtx) {
 
   SGeosContext *tlGeosCtxObj = (SGeosContext *)taosMemoryCalloc(1, sizeof(SGeosContext));
   if (!tlGeosCtxObj) {
+    taosThreadKeyDelete(tlGeosCtxKey);
     TAOS_CHECK_EXIT(TSDB_CODE_OUT_OF_MEMORY);
   }
   if ((taosThreadSetSpecific(tlGeosCtxKey, (const void *)tlGeosCtxObj)) != 0) {
+    taosThreadKeyDelete(tlGeosCtxKey);
     taosMemoryFreeClear(tlGeosCtxObj);
     TAOS_CHECK_EXIT(TAOS_SYSTEM_ERROR(errno));
   }