From c1f1709d59c6b279c6d4f52bde05c16b4924e927 Mon Sep 17 00:00:00 2001
From: Benguang Zhao <bgzhao@taosdata.com>
Date: Sun, 8 Oct 2023 17:03:07 +0800
Subject: [PATCH] fix: add rsmaSnapWriterPrepareClose

---
 source/dnode/vnode/src/inc/vnodeInt.h      |  1 +
 source/dnode/vnode/src/sma/smaSnapshot.c   | 15 +++++++++++++++
 source/dnode/vnode/src/vnd/vnodeSnapshot.c |  4 ++++
 3 files changed, 20 insertions(+)

diff --git a/source/dnode/vnode/src/inc/vnodeInt.h b/source/dnode/vnode/src/inc/vnodeInt.h
index 5015202865..68334b3b63 100644
--- a/source/dnode/vnode/src/inc/vnodeInt.h
+++ b/source/dnode/vnode/src/inc/vnodeInt.h
@@ -359,6 +359,7 @@ int32_t rsmaSnapRead(SRSmaSnapReader* pReader, uint8_t** ppData);
 // SRSmaSnapWriter ========================================
 int32_t rsmaSnapWriterOpen(SSma* pSma, int64_t sver, int64_t ever, void** ppRanges, SRSmaSnapWriter** ppWriter);
 int32_t rsmaSnapWrite(SRSmaSnapWriter* pWriter, uint8_t* pData, uint32_t nData);
+int32_t rsmaSnapWriterPrepareClose(SRSmaSnapWriter* pWriter);
 int32_t rsmaSnapWriterClose(SRSmaSnapWriter** ppWriter, int8_t rollback);
 
 typedef struct {
diff --git a/source/dnode/vnode/src/sma/smaSnapshot.c b/source/dnode/vnode/src/sma/smaSnapshot.c
index 1e921b23d1..c93d9a7de6 100644
--- a/source/dnode/vnode/src/sma/smaSnapshot.c
+++ b/source/dnode/vnode/src/sma/smaSnapshot.c
@@ -165,6 +165,21 @@ _exit:
   return code;
 }
 
+int32_t rsmaSnapWriterPrepareClose(SRSmaSnapWriter* pWriter) {
+  int32_t code = 0;
+  for (int32_t i = 0; i < TSDB_RETENTION_L2; ++i) {
+    if (pWriter->pDataWriter[i]) {
+      code = tsdbSnapWriterPrepareClose(pWriter->pDataWriter[i]);
+      if (code) {
+        smaError("vgId:%d, failed to prepare close tsdbSnapWriter since %s. i: %d", SMA_VID(pWriter->pSma), terrstr(),
+                 i);
+        return -1;
+      }
+    }
+  }
+  return code;
+}
+
 int32_t rsmaSnapWriterClose(SRSmaSnapWriter** ppWriter, int8_t rollback) {
   int32_t          code = 0;
   int32_t          lino = 0;
diff --git a/source/dnode/vnode/src/vnd/vnodeSnapshot.c b/source/dnode/vnode/src/vnd/vnodeSnapshot.c
index f5a5249d1f..f5d0bf6d0f 100644
--- a/source/dnode/vnode/src/vnd/vnodeSnapshot.c
+++ b/source/dnode/vnode/src/vnd/vnodeSnapshot.c
@@ -552,6 +552,10 @@ int32_t vnodeSnapWriterClose(SVSnapWriter *pWriter, int8_t rollback, SSnapshot *
     tsdbSnapWriterPrepareClose(pWriter->pTsdbSnapWriter);
   }
 
+  if (pWriter->pRsmaSnapWriter) {
+    rsmaSnapWriterPrepareClose(pWriter->pRsmaSnapWriter);
+  }
+
   // commit json
   if (!rollback) {
     pWriter->info.state.committed = pWriter->ever;