ddl: release scan chunks earlier to cap the memory usage (pingcap#67452)

close pingcap#64911

Author: joechenrh

pkg/ddl/backfilling_operators.go

@@ -532,6 +532,16 @@ func (w *tableScanWorker) scanRecords(task TableScanTask, sender func(IndexRecor
 		idxResults  []IndexRecordChunk
 		execDetails kvutil.ExecDetails
 	)
+	// Local ingest may trigger partial import/reset while the scan transaction is
+	// still open, so only the global-sort path can stream results immediately.
+	enableStreaming := w.reorgMeta.UseCloudStorage
+	sendResult := func(idxResult IndexRecordChunk) {
+		sender(idxResult)
+		if w.cpOp != nil {
+			w.cpOp.UpdateChunk(task.ID, int(idxResult.tableScanRowCount), idxResult.Done)
+		}
+		w.totalCount.Add(idxResult.tableScanRowCount)
+	}
 	var scanCtx context.Context = w.ctx
 	if scanCtx.Value(kvutil.ExecDetailsKey) == nil {
 		scanCtx = context.WithValue(w.ctx, kvutil.ExecDetailsKey, &execDetails)
@@ -565,6 +575,11 @@ func (w *tableScanWorker) scanRecords(task TableScanTask, sender func(IndexRecor
 			failpoint.InjectCall("beforeGetChunk")
 			srcChk := w.getChunk()
 			done, err = fetchTableScanResult(scanCtx, w.copCtx.GetBase(), rs, srcChk)
+			failpoint.Inject("mockScanRecordPartialError", func(val failpoint.Value) {
+				if shouldFail, _ := val.(bool); shouldFail {
+					err = errors.New("mock partial scan error")
+				}
+			})
 			if err != nil || scanCtx.Err() != nil {
 				w.recycleChunk(srcChk)
 				terror.Call(rs.Close)
@@ -574,19 +589,21 @@ func (w *tableScanWorker) scanRecords(task TableScanTask, sender func(IndexRecor
 			execDetails = kvutil.ExecDetails{}
 
 			_, tableScanRowCount := distsqlCtx.RuntimeStatsColl.GetCopCountAndRows(tableScanCopID)
-			idxResults = append(idxResults, IndexRecordChunk{ID: task.ID, Chunk: srcChk, Done: done, ctx: w.ctx, tableScanRowCount: tableScanRowCount - lastTableScanRowCount, conditionPushed: conditionPushed})
+			idxResult := IndexRecordChunk{ID: task.ID, Chunk: srcChk, Done: done, ctx: w.ctx, tableScanRowCount: tableScanRowCount - lastTableScanRowCount, conditionPushed: conditionPushed}
 			lastTableScanRowCount = tableScanRowCount
+			if enableStreaming {
+				sendResult(idxResult)
+			} else {
+				idxResults = append(idxResults, idxResult)
+			}
 		}
 		return rs.Close()
 	})
 
-	for i, idxResult := range idxResults {
-		sender(idxResult)
-		if w.cpOp != nil {
-			done := i == len(idxResults)-1
-			w.cpOp.UpdateChunk(task.ID, int(idxResult.tableScanRowCount), done)
+	if !enableStreaming {
+		for _, idxResult := range idxResults {
+			sendResult(idxResult)
 		}
-		w.totalCount.Add(idxResult.tableScanRowCount)
 	}
 
 	return err

tests/realtikvtest/addindextest4/integration_test.go

@@ -443,21 +443,28 @@ func TestCancelAfterReorgTimeout(t *testing.T) {
 	}, 10*time.Second, 300*time.Millisecond)
 }
 
+// setupAddIndexClassicVars enables fast-reorg and toggles the dist-task mode
+// in classic kernel. It is a no-op in next-gen.
+func setupAddIndexClassicVars(tk *testkit.TestKit, distTaskOn bool) {
+	if !kerneltype.IsClassic() {
+		return
+	}
+	tk.MustExec("set global tidb_ddl_enable_fast_reorg = 1")
+	if distTaskOn {
+		tk.MustExec("set global tidb_enable_dist_task = 1")
+	} else {
+		tk.MustExec("set global tidb_enable_dist_task = 0")
+	}
+}
+
 func TestAddIndexResumesFromCheckpointAfterPartialImport(t *testing.T) {
 	runCase := func(t *testing.T, distTaskOn bool) {
 		store := realtikvtest.CreateMockStoreAndSetup(t)
 
 		tk := testkit.NewTestKit(t, store)
 		tk.MustExec("use test")
 
-		if kerneltype.IsClassic() {
-			tk.MustExec("set global tidb_ddl_enable_fast_reorg = 1")
-			if distTaskOn {
-				tk.MustExec("set global tidb_enable_dist_task = 1")
-			} else {
-				tk.MustExec("set global tidb_enable_dist_task = 0")
-			}
-		}
+		setupAddIndexClassicVars(tk, distTaskOn)
 		ingest.ForceSyncFlagForTest.Store(true)
 
 		tk.MustExec("drop table if exists t")
@@ -489,3 +496,46 @@ func TestAddIndexResumesFromCheckpointAfterPartialImport(t *testing.T) {
 	t.Run("dist_task_off", func(t *testing.T) { runCase(t, false) })
 	t.Run("dist_task_on", func(t *testing.T) { runCase(t, true) })
 }
+
+func TestAddIndexResumesFromCheckpointAfterPartialScan(t *testing.T) {
+	runCase := func(t *testing.T, distTaskOn bool) {
+		store := realtikvtest.CreateMockStoreAndSetup(t)
+
+		tk := testkit.NewTestKit(t, store)
+		tk.MustExec("use test")
+
+		setupAddIndexClassicVars(tk, distTaskOn)
+		ingest.ForceSyncFlagForTest.Store(true)
+
+		tk.MustExec("drop table if exists t")
+		tk.MustExec("create table t (a bigint primary key, b bigint)")
+		for i := range 2000 {
+			tk.MustExec("insert into t values (?, ?)", i, i)
+		}
+
+		// Let the first chunk fetch in scanRecords succeed, then inject an
+		// error on the second fetch. The local-ingest path buffers scan
+		// results until the scan finishes, so this exercises the partial-scan
+		// checkpoint path where the retry must restart without a flushed chunk
+		// advancing the checkpoint. Subsequent fetches (in the retry) see the
+		// failpoint as already consumed and proceed normally.
+		testfailpoint.Enable(t,
+			"github.com/pingcap/tidb/pkg/ddl/mockScanRecordPartialError",
+			"1*return(false)->1*return(true)")
+
+		tk.MustExec("alter table t add unique index idx_b(b)")
+
+		tblCntStr := tk.MustQuery("select count(*) from t").Rows()[0][0].(string)
+		idxCntStr := tk.MustQuery("select count(*) from t use index(idx_b)").Rows()[0][0].(string)
+		tblCnt, err := strconv.Atoi(tblCntStr)
+		require.NoError(t, err)
+		idxCnt, err := strconv.Atoi(idxCntStr)
+		require.NoError(t, err)
+		require.Equal(t, tblCnt, idxCnt)
+
+		tk.MustExec("admin check table t")
+	}
+
+	t.Run("dist_task_off", func(t *testing.T) { runCase(t, false) })
+	t.Run("dist_task_on", func(t *testing.T) { runCase(t, true) })
+}