opt write snapshot

source/libs/stream/inc/streamBackendRocksdb.h

@@ -77,4 +77,12 @@ int32_t          streamStateGetParTag_rocksdb(SStreamState* pState, int64_t grou
 int32_t streamStatePutParName_rocksdb(SStreamState* pState, int64_t groupId, const char tbname[TSDB_TABLE_NAME_LEN]);
 int32_t streamStateGetParName_rocksdb(SStreamState* pState, int64_t groupId, void** pVal);
 void    streamStateDestroy_rocksdb(SStreamState* pState);
+
+void*   streamStateCreateBatch();
+int32_t streamStateGetBatchSize(void* pBatch);
+void    streamStateClearBatch(void* pBatch);
+void    streamStateDestroyBatch(void* pBatch);
+int32_t streamStatePutBatch(SStreamState* pState, const char* cfName, rocksdb_writebatch_t* pBatch, void* key,
+                            void* val, int32_t vlen);
+int32_t streamStatePutBatch_rocksdb(SStreamState* pState, void* pBatch);
 #endif

source/libs/stream/src/streamStateRocksdb.c

@@ -13,7 +13,6 @@
  * along with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 
-#include "rocksdb/c.h"
 #include "streamBackendRocksdb.h"
 #include "tcommon.h"
 #include "tlog.h"
@@ -544,6 +543,44 @@ int32_t streamStatePut_rocksdb(SStreamState* pState, const SWinKey* key, const v
   STREAM_STATE_PUT_ROCKSDB(pState, "default", &sKey, value, vLen);
   return code;
 }
+int32_t streamStatePutBatch_rocksdb(SStreamState* pState, void* pBatch) {
+  char* err = NULL;
+  rocksdb_write(pState->pTdbState->rocksdb, pState->pTdbState->writeOpts, (rocksdb_writebatch_t*)pBatch, &err);
+  if (err != NULL) {
+    qError("streamState failed to write batch, err:%s", err);
+    taosMemoryFree(err);
+    return -1;
+  }
+  return 0;
+}
+
+void* streamStateCreateBatch() {
+  rocksdb_writebatch_t* pBatch = rocksdb_writebatch_create();
+  return pBatch;
+}
+int32_t streamStateGetBatchSize(void* pBatch) {
+  if (pBatch == NULL) return -1;
+
+  return rocksdb_writebatch_count((rocksdb_writebatch_t*)pBatch);
+}
+
+void    streamStateClearBatch(void* pBatch) { rocksdb_writebatch_clear((rocksdb_writebatch_t*)pBatch); }
+void    streamStateDestroyBatch(void* pBatch) { rocksdb_writebatch_destroy((rocksdb_writebatch_t*)pBatch); }
+int32_t streamStatePutBatch(SStreamState* pState, const char* cfName, rocksdb_writebatch_t* pBatch, void* key,
+                            void* val, int32_t vlen) {
+  int i = streamGetInit(cfName);
+
+  if (i < 0) {
+    qError("streamState failed to put to cf name:%s", cfName);
+    return -1;
+  }
+  char    buf[128] = {0};
+  int32_t klen = ginitDict[i].enFunc((void*)key, buf);
+
+  rocksdb_column_family_handle_t* pCf = pState->pTdbState->pHandle[ginitDict[i].idx];
+  rocksdb_writebatch_put_cf((rocksdb_writebatch_t*)pBatch, pCf, buf, (size_t)klen, val, (size_t)vlen);
+  return 0;
+}
 int32_t streamStateGet_rocksdb(SStreamState* pState, const SWinKey* key, void** pVal, int32_t* pVLen) {
   int       code = 0;
   SStateKey sKey = {.key = *key, .opNum = pState->number};

source/libs/stream/src/tstreamFileState.c

@@ -15,11 +15,10 @@
 
 #include "tstreamFileState.h"
 
+#include "streamBackendRocksdb.h"
 #include "taos.h"
 #include "thash.h"
 #include "tsimplehash.h"
-#include "streamBackendRocksdb.h"
-
 
 #define FLUSH_RATIO                    0.2
 #define DEFAULT_MAX_STREAM_BUFFER_SIZE (128 * 1024 * 1024);
@@ -264,9 +263,7 @@ bool hasRowBuff(SStreamFileState* pFileState, void* pKey, int32_t keyLen) {
   return false;
 }
 
-void releaseRowBuffPos(SRowBuffPos* pBuff) {
-  pBuff->beUsed = false;
-}
+void releaseRowBuffPos(SRowBuffPos* pBuff) { pBuff->beUsed = false; }
 
 SStreamSnapshot* getSnapshot(SStreamFileState* pFileState) {
   clearExpiredRowBuff(pFileState, pFileState->maxTs - pFileState->deleteMark, false);
@@ -289,11 +286,24 @@ int32_t flushSnapshot(SStreamFileState* pFileState, SStreamSnapshot* pSnapshot,
   SListIter iter = {0};
   tdListInitIter(pSnapshot, &iter, TD_LIST_FORWARD);
 
+  const int32_t BATCH_LIMIT = 128;
   SListNode*    pNode = NULL;
+
+  void* batch = streamStateCreateBatch();
   while ((pNode = tdListNext(&iter)) != NULL && code == TSDB_CODE_SUCCESS) {
     SRowBuffPos* pPos = *(SRowBuffPos**)pNode->data;
-    code = streamStatePut_rocksdb(pFileState->pFileStore, pPos->pKey, pPos->pRowBuff, pFileState->rowSize);
+    if (streamStateGetBatchSize(batch) >= BATCH_LIMIT) {
+      code = streamStatePutBatch_rocksdb(pFileState->pFileStore, batch);
+      streamStateClearBatch(batch);
     }
+    code =
+        streamStatePutBatch(pFileState->pFileStore, "default", batch, pPos->pKey, pPos->pRowBuff, pFileState->rowSize);
+  }
+  if (streamStateGetBatchSize(batch) > 0) {
+    code = streamStatePutBatch_rocksdb(pFileState->pFileStore, batch);
+  }
+  streamStateDestroyBatch(batch);
+
   if (flushState) {
     int32_t len = 0;
     void*   buff = NULL;