From 4e4b3661ea1bbd809f3c8407d0b77163d0cb11eb Mon Sep 17 00:00:00 2001
From: Minglei Jin <mljin@taosdata.com>
Date: Mon, 19 Dec 2022 11:39:41 +0800
Subject: [PATCH] fix(tdb): rollback in-memory pages

---
 source/libs/tdb/src/db/tdbPCache.c | 28 +++++++++++++++++++++++++++-
 source/libs/tdb/src/db/tdbPager.c  |  1 +
 source/libs/tdb/src/inc/tdbInt.h   |  2 ++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/source/libs/tdb/src/db/tdbPCache.c b/source/libs/tdb/src/db/tdbPCache.c
index b67fe562eb..30dcb05f80 100644
--- a/source/libs/tdb/src/db/tdbPCache.c
+++ b/source/libs/tdb/src/db/tdbPCache.c
@@ -192,6 +192,28 @@ SPage *tdbPCacheFetch(SPCache *pCache, const SPgid *pPgid, TXN *pTxn) {
   return pPage;
 }
 
+void tdbPCacheMarkFree(SPCache *pCache, SPage *pPage) {
+  tdbPCacheLock(pCache);
+  tdbPCacheRemovePageFromHash(pCache, pPage);
+  pPage->isFree = 1;
+  tdbPCacheUnlock(pCache);
+}
+
+static void tdbPCacheFreePage(SPCache *pCache, SPage *pPage) {
+  if (pPage->id < pCache->nPages) {
+    pPage->pFreeNext = pCache->pFree;
+    pCache->pFree = pPage;
+    pPage->isFree = 0;
+    ++pCache->nFree;
+    tdbTrace("pcache/free page %p/%d/%d", pPage, TDB_PAGE_PGNO(pPage), pPage->id);
+  } else {
+    tdbTrace("pcache destroy page: %p/%d/%d", pPage, TDB_PAGE_PGNO(pPage), pPage->id);
+
+    tdbPCacheRemovePageFromHash(pCache, pPage);
+    tdbPageDestroy(pPage, tdbDefaultFree, NULL);
+  }
+}
+
 void tdbPCacheRelease(SPCache *pCache, SPage *pPage, TXN *pTxn) {
   i32 nRef;
 
@@ -209,7 +231,11 @@ void tdbPCacheRelease(SPCache *pCache, SPage *pPage, TXN *pTxn) {
     // nRef = tdbGetPageRef(pPage);
     // if (nRef == 0) {
     if (pPage->isLocal) {
-      tdbPCacheUnpinPage(pCache, pPage);
+      if (!pPage->isFree) {
+        tdbPCacheUnpinPage(pCache, pPage);
+      } else {
+        tdbPCacheFreePage(pCache, pPage);
+      }
     } else {
       if (TDB_TXN_IS_WRITE(pTxn)) {
         // remove from hash
diff --git a/source/libs/tdb/src/db/tdbPager.c b/source/libs/tdb/src/db/tdbPager.c
index 648e99d6a5..23ab0004b9 100644
--- a/source/libs/tdb/src/db/tdbPager.c
+++ b/source/libs/tdb/src/db/tdbPager.c
@@ -524,6 +524,7 @@ int tdbPagerAbort(SPager *pPager, TXN *pTxn) {
 
     tRBTreeDrop(&pPager->rbt, (SRBTreeNode *)pPage);
     hashset_remove(pTxn->jPageSet, (void *)((long)TDB_PAGE_PGNO(pPage)));
+    tdbPCacheMarkFree(pPager->pCache, pPage);
     tdbPCacheRelease(pPager->pCache, pPage, pTxn);
   }
 
diff --git a/source/libs/tdb/src/inc/tdbInt.h b/source/libs/tdb/src/inc/tdbInt.h
index 055a8a1062..ab9c5995ec 100644
--- a/source/libs/tdb/src/inc/tdbInt.h
+++ b/source/libs/tdb/src/inc/tdbInt.h
@@ -205,6 +205,7 @@ int  tdbPagerRollback(SPager *pPager);
   u8           isAnchor;   \
   u8           isLocal;    \
   u8           isDirty;    \
+  u8           isFree;     \
   volatile i32 nRef;       \
   i32          id;         \
   SPage       *pFreeNext;  \
@@ -222,6 +223,7 @@ int    tdbPCacheClose(SPCache *pCache);
 int    tdbPCacheAlter(SPCache *pCache, int32_t nPage);
 SPage *tdbPCacheFetch(SPCache *pCache, const SPgid *pPgid, TXN *pTxn);
 void   tdbPCacheRelease(SPCache *pCache, SPage *pPage, TXN *pTxn);
+void   tdbPCacheMarkFree(SPCache *pCache, SPage *pPage);
 int    tdbPCacheGetPageSize(SPCache *pCache);
 
 // tdbPage.c ====================================