From f4359a6fb3c82295ac5b8b752e2804b7dd19ad84 Mon Sep 17 00:00:00 2001
From: Hongze Cheng <hzcheng@taosdata.com>
Date: Fri, 25 Mar 2022 06:02:38 +0000
Subject: [PATCH] more TDB

---
 source/libs/tdb/src/db/tdbBtree.c |  7 ++++---
 source/libs/tdb/src/inc/tdbUtil.h | 22 ++++++++++++++++++++++
 source/libs/tdb/test/tdbTest.cpp  |  2 ++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/source/libs/tdb/src/db/tdbBtree.c b/source/libs/tdb/src/db/tdbBtree.c
index c3adf7d2b6..0ef5d64d30 100644
--- a/source/libs/tdb/src/db/tdbBtree.c
+++ b/source/libs/tdb/src/db/tdbBtree.c
@@ -207,6 +207,7 @@ int tdbBtreeGet(SBTree *pBt, const void *pKey, int kLen, void **ppVal, int *vLen
   SBTC         btc;
   SCell       *pCell;
   int          cret;
+  void        *pVal;
   SCellDecoder cd;
 
   tdbBtreeCursor(&btc, pBt);
@@ -221,12 +222,12 @@ int tdbBtreeGet(SBTree *pBt, const void *pKey, int kLen, void **ppVal, int *vLen
   tdbBtreeDecodeCell(btc.pPage, pCell, &cd);
 
   *vLen = cd.vLen;
-  // TODO: here may have memory leak
-  *ppVal = realloc(*ppVal, *vLen);
-  if (*ppVal == NULL) {
+  pVal = TDB_REALLOC(*ppVal, *vLen);
+  if (pVal == NULL) {
     return -1;
   }
 
+  *ppVal = pVal;
   memcpy(*ppVal, cd.pVal, cd.vLen);
   return 0;
 }
diff --git a/source/libs/tdb/src/inc/tdbUtil.h b/source/libs/tdb/src/inc/tdbUtil.h
index 9358aae236..88fc846bf1 100644
--- a/source/libs/tdb/src/inc/tdbUtil.h
+++ b/source/libs/tdb/src/inc/tdbUtil.h
@@ -39,6 +39,28 @@ int tdbGetFileSize(const char *fname, int pgSize, SPgno *pSize);
 
 int tdbPRead(int fd, void *pData, int count, i64 offset);
 
+#define TDB_REALLOC(PTR, SIZE)                                                          \
+  ({                                                                                    \
+    void *nPtr;                                                                         \
+    if ((PTR) == NULL || ((int *)(PTR))[-1] < (SIZE)) {                                 \
+      nPtr = realloc((PTR) ? (char *)(PTR) - sizeof(int) : NULL, (SIZE) + sizeof(int)); \
+      if (nPtr) {                                                                       \
+        ((int *)nPtr)[0] = (SIZE);                                                      \
+        nPtr = (char *)nPtr + sizeof(int);                                              \
+      }                                                                                 \
+    } else {                                                                            \
+      nPtr = (PTR);                                                                     \
+    }                                                                                   \
+    nPtr;                                                                               \
+  })
+
+#define TDB_FREE(PTR)                    \
+  do {                                   \
+    if (PTR) {                           \
+      free((char *)(PTR) - sizeof(int)); \
+    }                                    \
+  } while (0)
+
 static inline void *tdbOsMalloc(void *arg, size_t size) {
   void *ptr;
 
diff --git a/source/libs/tdb/test/tdbTest.cpp b/source/libs/tdb/test/tdbTest.cpp
index 05a0f3a6ae..46437392cd 100644
--- a/source/libs/tdb/test/tdbTest.cpp
+++ b/source/libs/tdb/test/tdbTest.cpp
@@ -44,6 +44,8 @@ TEST(tdb_test, simple_test) {
         GTEST_ASSERT_EQ(vLen, strlen(val));
         GTEST_ASSERT_EQ(memcmp(val, pVal, vLen), 0);
       }
+
+      TDB_FREE(pVal);
     }
 
     {  // Loop to query the data