ddl: enhance checkpoint management with chunk tracking

pkg/ddl/backfilling_operators.go

@@ -803,10 +803,6 @@ func (w *indexIngestWorker) HandleTask(ck IndexRecordChunk, send func(IndexWrite
 		return err
 	}
 	scannedCount := ck.tableScanRowCount
-	if scannedCount == 0 {
-		logutil.Logger(w.ctx).Info("finish a index ingest task", zap.Int("id", ck.ID))
-		return nil
-	}
 	if w.totalCount != nil {
 		w.totalCount.Add(scannedCount)
 	}

pkg/ddl/ingest/checkpoint.go

@@ -199,10 +199,12 @@ type CheckpointManager struct {
 
 // taskCheckpoint is the checkpoint for a single task.
 type taskCheckpoint struct {
-	totalKeys     int
-	writtenKeys   int
-	endKey        kv.Key
-	lastBatchRead bool
+	totalKeys      int
+	writtenKeys    int
+	endKey         kv.Key
+	lastBatchRead  bool
+	chunksTotal    int
+	chunksFinished int
 }
 
 // newCheckpointManagerWithStorage is the common constructor
@@ -348,14 +350,24 @@ func (s *CheckpointManager) UpdateChunk(taskID int, delta int, last bool) {
 	cp := s.checkpoints[taskID]
 	cp.totalKeys += delta
 	cp.lastBatchRead = last
+	cp.chunksTotal++
 }
 
 // FinishChunk updates the written keys of the task.
 // This is called by the writer after writing the local engine to update the current number of rows written.
 func (s *CheckpointManager) FinishChunk(taskID int, delta int) {
 	s.mu.Lock()
-	cp := s.checkpoints[taskID]
+	cp, ok := s.checkpoints[taskID]
+	if !ok {
+		s.mu.Unlock()
+		s.logger.Warn("finish chunk for unknown task", zap.Int("taskID", taskID))
+		return
+	}
 	cp.writtenKeys += delta
+	cp.chunksFinished++
+	if cp.chunksFinished == cp.chunksTotal {
+		s.logger.Info("finish a index ingest task", zap.Int("id", taskID), zap.Int("totalKeys", cp.totalKeys), zap.Int("writtenKeys", cp.writtenKeys))
+	}
 	s.mu.Unlock()
 }
 
@@ -396,7 +408,7 @@ func (s *CheckpointManager) afterFlush() {
 	defer s.mu.Unlock()
 	for {
 		cp := s.checkpoints[s.minTaskIDFinished]
-		if cp == nil || !cp.lastBatchRead || cp.writtenKeys < cp.totalKeys {
+		if cp == nil || !cp.lastBatchRead || cp.writtenKeys < cp.totalKeys || cp.chunksFinished < cp.chunksTotal {
 			break
 		}
 		delete(s.checkpoints, s.minTaskIDFinished)

pkg/ddl/ingest/checkpoint_test.go

@@ -108,6 +108,24 @@ func TestCheckpointManager(t *testing.T) {
 	require.NoError(t, mgr.AdvanceWatermark(true))
 	require.False(t, mgr.IsKeyProcessed([]byte{'2', '9'}))
 	require.False(t, mgr.IsKeyProcessed([]byte{'3', '9'}))
+
+	// Finish the remaining chunk so we can continue to new tasks.
+	mgr.FinishChunk(2, 100)
+	require.NoError(t, mgr.AdvanceWatermark(true))
+	require.True(t, mgr.IsKeyProcessed([]byte{'2', '9'}))
+	require.True(t, mgr.IsKeyProcessed([]byte{'3', '9'}))
+	require.True(t, mgr.IsKeyProcessed([]byte{'4', '9'}))
+
+	// Chunk completion counter: progress only after all chunks finish.
+	mgr.AddChunk(5, []byte{'5', '9'})
+	// mock wrong row count
+	mgr.UpdateChunk(5, 100, false)
+	mgr.UpdateChunk(5, 0, true)
+	mgr.FinishChunk(5, 100)
+	require.NoError(t, mgr.AdvanceWatermark(true))
+	mgr.FinishChunk(5, 0)
+	require.NoError(t, mgr.AdvanceWatermark(true))
+	require.True(t, mgr.IsKeyProcessed([]byte{'5', '9'}))
 }
 
 func TestCheckpointManagerUpdateReorg(t *testing.T) {

tests/realtikvtest/addindextest4/integration_test.go

@@ -451,10 +451,6 @@ func TestAddIndexResumesFromCheckpointAfterPartialImport(t *testing.T) {
 		}
 		ingest.ForceSyncFlagForTest.Store(true)
 
-		tk.Session().Close()
-		tk = testkit.NewTestKit(t, store)
-		tk.MustExec("use test")
-
 		tk.MustExec("drop table if exists t")
 		tk.MustExec("create table t (a bigint primary key, b bigint)")
 		for i := 0; i < 2000; i++ {