From a504007ade887f559c3ad2486d7a82de09651b5f Mon Sep 17 00:00:00 2001
From: Benguang Zhao <bgzhao@taosdata.com>
Date: Fri, 24 Nov 2023 16:13:53 +0800
Subject: [PATCH] refact: alloc just enough memory for snap data in
 tsdbSnapRAWReadData

---
 source/dnode/vnode/src/tsdb/tsdbDataFileRAW.c | 24 +++----
 source/dnode/vnode/src/tsdb/tsdbSnapshotRAW.c | 65 ++++++++++++-------
 2 files changed, 49 insertions(+), 40 deletions(-)

diff --git a/source/dnode/vnode/src/tsdb/tsdbDataFileRAW.c b/source/dnode/vnode/src/tsdb/tsdbDataFileRAW.c
index 0c3d890966..aaaaa56b0e 100644
--- a/source/dnode/vnode/src/tsdb/tsdbDataFileRAW.c
+++ b/source/dnode/vnode/src/tsdb/tsdbDataFileRAW.c
@@ -60,27 +60,21 @@ int32_t tsdbDataFileRAWReaderClose(SDataFileRAWReader **reader) {
   return 0;
 }
 
-int32_t tsdbDataFileRAWReadBlockData(SDataFileRAWReader *reader, STsdbDataRAWBlockHeader *bHdr) {
+int32_t tsdbDataFileRAWReadBlockData(SDataFileRAWReader *reader, STsdbDataRAWBlockHeader *pBlock) {
   int32_t code = 0;
   int32_t lino = 0;
 
-  bHdr->file.type = reader->config->file.type;
-  bHdr->file.fid = reader->config->file.fid;
-  bHdr->file.cid = reader->config->file.cid;
-  bHdr->file.size = reader->config->file.size;
-  bHdr->file.minVer = reader->config->file.minVer;
-  bHdr->file.maxVer = reader->config->file.maxVer;
-  bHdr->file.stt->level = reader->config->file.stt->level;
+  pBlock->file.type = reader->config->file.type;
+  pBlock->file.fid = reader->config->file.fid;
+  pBlock->file.cid = reader->config->file.cid;
+  pBlock->file.size = reader->config->file.size;
+  pBlock->file.minVer = reader->config->file.minVer;
+  pBlock->file.maxVer = reader->config->file.maxVer;
+  pBlock->file.stt->level = reader->config->file.stt->level;
 
-  int64_t size = TMIN(bHdr->dataLength, reader->config->file.size - reader->ctx->offset);
-  ASSERT(size > 0);
-  bHdr->dataLength = 0;
-  bHdr->offset = reader->ctx->offset;
-
-  code = tsdbReadFile(reader->fd, bHdr->offset, bHdr->data, size);
+  code = tsdbReadFile(reader->fd, pBlock->offset, pBlock->data, pBlock->dataLength);
   TSDB_CHECK_CODE(code, lino, _exit);
 
-  bHdr->dataLength = size;
 _exit:
   if (code) {
     TSDB_ERROR_LOG(TD_VID(reader->config->tsdb->pVnode), lino, code);
diff --git a/source/dnode/vnode/src/tsdb/tsdbSnapshotRAW.c b/source/dnode/vnode/src/tsdb/tsdbSnapshotRAW.c
index 2bcd00bfec..14c52c0bfb 100644
--- a/source/dnode/vnode/src/tsdb/tsdbSnapshotRAW.c
+++ b/source/dnode/vnode/src/tsdb/tsdbSnapshotRAW.c
@@ -183,52 +183,54 @@ static int32_t tsdbSnapRAWReadFileSetCloseIter(STsdbSnapRAWReader* reader) {
   return 0;
 }
 
-static int32_t tsdbSnapRAWReadNext(STsdbSnapRAWReader* reader, STsdbDataRAWBlockHeader* bHdr) {
+static int64_t tsdbSnapRAWReadPeek(SDataFileRAWReader* reader) {
+  int64_t size = TMIN(reader->config->file.size - reader->ctx->offset, TSDB_SNAP_RAW_PAYLOAD_SIZE);
+  return size;
+}
+
+static int32_t tsdbSnapRAWReadNext(STsdbSnapRAWReader* reader, SSnapDataHdr** ppData) {
   int32_t code = 0;
   int32_t lino = 0;
+  ppData[0] = NULL;
 
   ASSERT(reader->dataIter->offset <= reader->dataIter->size);
   ASSERT(reader->dataIter->idx <= reader->dataIter->count);
 
+  // dataReader
   if (reader->dataIter->offset == reader->dataIter->size && reader->dataIter->idx < reader->dataIter->count) {
     reader->dataIter->idx++;
   }
   if (reader->dataIter->idx == reader->dataIter->count) {
     return 0;
   }
-
+  int8_t              type = reader->type;
   SDataFileRAWReader* dataReader = TARRAY2_GET(reader->dataReaderArr, reader->dataIter->idx);
-  code = tsdbDataFileRAWReadBlockData(dataReader, bHdr);
-  TSDB_CHECK_CODE(code, lino, _exit);
 
-  reader->dataIter->offset += bHdr->dataLength;
-_exit:
-  if (code) {
-    TSDB_ERROR_LOG(TD_VID(reader->tsdb->pVnode), code, lino);
-  }
-  return code;
-}
+  // prepare
+  int64_t dataLength = tsdbSnapRAWReadPeek(dataReader);
+  ASSERT(dataLength > 0);
 
-static int32_t tsdbSnapRAWReadData(STsdbSnapRAWReader* reader, SSnapDataHdr** data) {
-  int32_t code = 0;
-  int32_t lino = 0;
-
-  void* pBuf = taosMemoryCalloc(1, sizeof(SSnapDataHdr) + sizeof(STsdbDataRAWBlockHeader) + TSDB_SNAP_RAW_PAYLOAD_SIZE);
+  void* pBuf = taosMemoryCalloc(1, sizeof(SSnapDataHdr) + sizeof(STsdbDataRAWBlockHeader) + dataLength);
   if (pBuf == NULL) {
     code = TSDB_CODE_OUT_OF_MEMORY;
     TSDB_CHECK_CODE(code, lino, _exit);
   }
-
   SSnapDataHdr* pHdr = pBuf;
-  pHdr->type = reader->type;
-  STsdbDataRAWBlockHeader* pData = (void*)pHdr->data;
-  pData->dataLength = TSDB_SNAP_RAW_PAYLOAD_SIZE;
+  pHdr->type = type;
+  pHdr->size = sizeof(STsdbDataRAWBlockHeader) + dataLength;
 
-  code = tsdbSnapRAWReadNext(reader, pData);
+  STsdbDataRAWBlockHeader* pBlock = (void*)pHdr->data;
+  pBlock->offset = dataReader->ctx->offset;
+  pBlock->dataLength = dataLength;
+
+  // read
+  code = tsdbDataFileRAWReadBlockData(dataReader, pBlock);
   TSDB_CHECK_CODE(code, lino, _exit);
 
-  ASSERT(pData->dataLength > 0 && pData->dataLength <= TSDB_SNAP_RAW_PAYLOAD_SIZE);
-  pHdr->size = sizeof(STsdbDataRAWBlockHeader) + pData->dataLength;
+  // finish
+  reader->dataIter->offset += pBlock->dataLength;
+  ppData[0] = pBuf;
+  ASSERT(reader->dataIter->offset <= reader->dataIter->size);
 
 _exit:
   if (code) {
@@ -236,7 +238,20 @@ _exit:
     pBuf = NULL;
     TSDB_ERROR_LOG(TD_VID(reader->tsdb->pVnode), code, lino);
   }
-  data[0] = pBuf;
+  return code;
+}
+
+static int32_t tsdbSnapRAWReadData(STsdbSnapRAWReader* reader, uint8_t** ppData) {
+  int32_t code = 0;
+  int32_t lino = 0;
+
+  code = tsdbSnapRAWReadNext(reader, (SSnapDataHdr**)ppData);
+  TSDB_CHECK_CODE(code, lino, _exit);
+
+_exit:
+  if (code) {
+    TSDB_ERROR_LOG(TD_VID(reader->tsdb->pVnode), code, lino);
+  }
   return code;
 }
 
@@ -288,7 +303,7 @@ int32_t tsdbSnapRAWRead(STsdbSnapRAWReader* reader, uint8_t** data) {
     }
 
     if (!reader->ctx->isDataDone) {
-      code = tsdbSnapRAWReadData(reader, (SSnapDataHdr**)data);
+      code = tsdbSnapRAWReadData(reader, data);
       TSDB_CHECK_CODE(code, lino, _exit);
       if (data[0]) {
         goto _exit;