fix big txn dump (matrixorigin#20900)

修复了workspace 超过10M的大事务不刷盘的问题

Approved by: @triump2020, @XuPeng-SH, @sukki37

Author: Wenbin1002

pkg/vm/engine/disttae/txn.go

@@ -493,6 +493,7 @@ func (txn *Transaction) dumpBatchLocked(ctx context.Context, offset int) error {
 					"WORKSPACE-QUOTA-ACQUIRE",
 					zap.Uint64("quota", quota),
 					zap.Uint64("remaining", remaining),
+					zap.String("txn", txn.op.Txn().DebugString()),
 				)
 				txn.writeWorkspaceThreshold += quota
 				txn.extraWriteWorkspaceThreshold += quota
@@ -513,8 +514,10 @@ func (txn *Transaction) dumpBatchLocked(ctx context.Context, offset int) error {
 			"WORKSPACE-QUOTA-RELEASE",
 			zap.Uint64("quota", txn.extraWriteWorkspaceThreshold),
 			zap.Uint64("remaining", remaining),
+			zap.String("txn", txn.op.Txn().DebugString()),
 		)
 		txn.extraWriteWorkspaceThreshold = 0
+		txn.writeWorkspaceThreshold = txn.engine.config.writeWorkspaceThreshold
 	}
 
 	if dumpAll {
@@ -544,7 +547,7 @@ func (txn *Transaction) dumpBatchLocked(ctx context.Context, offset int) error {
 }
 
 func (txn *Transaction) dumpInsertBatchLocked(ctx context.Context, offset int, size *uint64, pkCount *int) error {
-	tbSize := make(map[uint64]uint64)
+	tbSize := make(map[uint64]int)
 	tbCount := make(map[uint64]int)
 	skipTable := make(map[uint64]bool)
 
@@ -556,8 +559,8 @@ func (txn *Transaction) dumpInsertBatchLocked(ctx context.Context, offset int, s
 			continue
 		}
 		if txn.writes[i].typ == INSERT && txn.writes[i].fileName == "" {
+			tbSize[txn.writes[i].tableId] += txn.writes[i].bat.Size()
 			tbCount[txn.writes[i].tableId] += txn.writes[i].bat.RowCount()
-			tbSize[txn.writes[i].tableId] += uint64(txn.writes[i].bat.Size())
 		}
 	}
 
@@ -568,12 +571,12 @@ func (txn *Transaction) dumpInsertBatchLocked(ctx context.Context, offset int, s
 	sort.Slice(keys, func(i, j int) bool {
 		return tbSize[keys[i]] < tbSize[keys[j]]
 	})
-	sum := uint64(0)
+	sum := 0
 	for _, k := range keys {
 		if tbCount[k] >= txn.engine.config.insertEntryMaxCount {
 			continue
 		}
-		if sum+tbSize[k] >= txn.writeWorkspaceThreshold {
+		if uint64(sum+tbSize[k]) >= txn.commitWorkspaceThreshold {
 			break
 		}
 		sum += tbSize[k]
@@ -612,11 +615,12 @@ func (txn *Transaction) dumpInsertBatchLocked(ctx context.Context, offset int, s
 			*size += uint64(bat.Size())
 			*pkCount += bat.RowCount()
 			// skip rowid
-			newBat := batch.NewWithSize(len(bat.Vecs) - 1)
-			newBat.SetAttributes(bat.Attrs[1:])
-			newBat.Vecs = bat.Vecs[1:]
-			newBat.SetRowCount(bat.Vecs[0].Length())
-			mp[tbKey] = append(mp[tbKey], newBat)
+			newBatch := batch.NewWithSize(len(bat.Vecs) - 1)
+			newBatch.SetAttributes(bat.Attrs[1:])
+			newBatch.Vecs = bat.Vecs[1:]
+			newBatch.SetRowCount(bat.Vecs[0].Length())
+			mp[tbKey] = append(mp[tbKey], newBatch)
+			defer bat.Clean(txn.proc.GetMPool())
 
 			keepElement = false
 		}
@@ -719,12 +723,12 @@ func (txn *Transaction) dumpDeleteBatchLocked(ctx context.Context, offset int, s
 			deleteCnt += bat.RowCount()
 			*size += uint64(bat.Size())
 
-			newBat := batch.NewWithSize(len(bat.Vecs))
-			newBat.SetAttributes(bat.Attrs)
-			newBat.Vecs = bat.Vecs
-			newBat.SetRowCount(bat.Vecs[0].Length())
+			newBatch := batch.NewWithSize(len(bat.Vecs))
+			newBatch.SetAttributes(bat.Attrs)
+			newBatch.Vecs = bat.Vecs
+			newBatch.SetRowCount(bat.Vecs[0].Length())
 
-			mp[tbKey] = append(mp[tbKey], newBat)
+			mp[tbKey] = append(mp[tbKey], newBatch)
 			defer bat.Clean(txn.proc.GetMPool())
 
 			keepElement = false
@@ -1308,6 +1312,22 @@ func (txn *Transaction) Commit(ctx context.Context) ([]txn.TxnRequest, error) {
 		return nil, nil
 	}
 
+	if err := txn.IncrStatementID(ctx, true); err != nil {
+		return nil, err
+	}
+
+	if err := txn.transferTombstonesByCommit(ctx); err != nil {
+		return nil, err
+	}
+
+	if err := txn.mergeTxnWorkspaceLocked(ctx); err != nil {
+		return nil, err
+	}
+	if err := txn.dumpBatchLocked(ctx, -1); err != nil {
+		return nil, err
+	}
+	txn.traceWorkspaceLocked(true)
+
 	if txn.workspaceSize > 10*mpool.MB {
 		logutil.Info(
 			"BIG-TXN",
@@ -1332,23 +1352,6 @@ func (txn *Transaction) Commit(ctx context.Context) ([]txn.TxnRequest, error) {
 		)
 	}
 
-	if err := txn.IncrStatementID(ctx, true); err != nil {
-		return nil, err
-	}
-
-	if err := txn.transferTombstonesByCommit(ctx); err != nil {
-		return nil, err
-	}
-
-	if err := txn.mergeTxnWorkspaceLocked(ctx); err != nil {
-		return nil, err
-	}
-	if err := txn.dumpBatchLocked(ctx, -1); err != nil {
-		return nil, err
-	}
-
-	txn.traceWorkspaceLocked(true)
-
 	if !txn.hasS3Op.Load() &&
 		txn.op.TxnOptions().CheckDupEnabled() {
 		if err := txn.checkDup(); err != nil {

pkg/vm/engine/test/disttae_engine_test.go

@@ -1304,6 +1304,7 @@ func TestWorkspaceQuota2(t *testing.T) {
 		ctx,
 		testutil.TestOptions{},
 		t,
+		testutil.WithDisttaeEngineCommitWorkspaceThreshold(1),
 		testutil.WithDisttaeEngineWriteWorkspaceThreshold(1),
 		testutil.WithDisttaeEngineQuota(800),
 	)