loader: fix loader lost data when quickly starts/stops task

loader/loader.go

@@ -282,6 +282,14 @@ func (w *Worker) dispatchSQL(ctx context.Context, file string, offset int64, tab
 
 		tctx := tcontext.NewContext(ctx, w.logger)
 		err2 = w.checkPoint.Init(tctx, baseFile, finfo.Size())
+		failpoint.Inject("WaitLoaderStopAfterInitCheckpoint", func(v failpoint.Value) {
+			t := v.(int)
+			w.logger.Info("wait loader stop after init checkpoint")
+			w.wg.Add(1)
+			time.Sleep(time.Duration(t) * time.Second)
+			w.wg.Done()
+		})
+
 		if err2 != nil {
 			w.logger.Error("fail to initial checkpoint", zap.String("data file", file), zap.Int64("offset", offset), log.ShortError(err2))
 			return err2
@@ -649,6 +657,14 @@ func (l *Loader) Restore(ctx context.Context) error {
 		return err
 	}
 
+	failpoint.Inject("WaitLoaderStopBeforeLoadCheckpoint", func(v failpoint.Value) {
+		t := v.(int)
+		l.logger.Info("wait loader stop before load checkpoint")
+		l.wg.Add(1)
+		time.Sleep(time.Duration(t) * time.Second)
+		l.wg.Done()
+	})
+
 	// not update checkpoint in memory when restoring, so when re-Restore, we need to load checkpoint from DB
 	err := l.checkPoint.Load(tcontext.NewContext(ctx, l.logger))
 	if err != nil {
@@ -657,6 +673,7 @@ func (l *Loader) Restore(ctx context.Context) error {
 	err = l.checkPoint.CalcProgress(l.db2Tables)
 	if err != nil {
 		l.logger.Error("calc load process", log.ShortError(err))
+		return err
 	}
 	l.loadFinishedSize()
 
@@ -686,6 +703,9 @@ func (l *Loader) Restore(ctx context.Context) error {
 	if err == nil {
 		l.finish.Set(true)
 		l.logger.Info("all data files have been finished", zap.Duration("cost time", time.Since(begin)))
+		if l.cfg.CleanDumpFile && l.checkPoint.AllFinished() {
+			l.cleanDumpFiles()
+		}
 	} else if errors.Cause(err) != context.Canceled {
 		return err
 	}
@@ -714,9 +734,6 @@ func (l *Loader) Close() {
 	if err != nil {
 		l.logger.Error("close downstream DB error", log.ShortError(err))
 	}
-	if l.cfg.CleanDumpFile && l.checkPoint.AllFinished() {
-		l.cleanDumpFiles()
-	}
 	l.checkPoint.Close()
 	l.removeLabelValuesWithTaskInMetrics(l.cfg.Name)
 	l.closed.Set(true)

tests/all_mode/run.sh

@@ -105,12 +105,77 @@ function test_query_timeout(){
     export GO_FAILPOINTS=''
 }
 
+function test_stop_task_before_checkpoint(){
+    run_sql_file $cur/data/db1.prepare.sql $MYSQL_HOST1 $MYSQL_PORT1 $MYSQL_PASSWORD1
+    check_contains 'Query OK, 2 rows affected'
+    run_sql_file $cur/data/db2.prepare.sql $MYSQL_HOST2 $MYSQL_PORT2 $MYSQL_PASSWORD2
+    check_contains 'Query OK, 3 rows affected'
+
+    # start DM worker and master
+    run_dm_master $WORK_DIR/master $MASTER_PORT $cur/conf/dm-master.toml
+    check_rpc_alive $cur/../bin/check_master_online 127.0.0.1:$MASTER_PORT
+
+    export GO_FAILPOINTS='github.com/pingcap/dm/loader/WaitLoaderStopAfterInitCheckpoint=return(5)'
+    run_dm_worker $WORK_DIR/worker1 $WORKER1_PORT $cur/conf/dm-worker1.toml
+    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER1_PORT
+    run_dm_worker $WORK_DIR/worker2 $WORKER2_PORT $cur/conf/dm-worker2.toml
+    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER2_PORT
+
+    # operate mysql config to worker
+    cp $cur/conf/source1.yaml $WORK_DIR/source1.yaml
+    cp $cur/conf/source2.yaml $WORK_DIR/source2.yaml
+    sed -i "/relay-binlog-name/i\relay-dir: $WORK_DIR/worker1/relay_log" $WORK_DIR/source1.yaml
+    sed -i "/relay-binlog-name/i\relay-dir: $WORK_DIR/worker2/relay_log" $WORK_DIR/source2.yaml
+    dmctl_operate_source create $WORK_DIR/source1.yaml $SOURCE_ID1
+    dmctl_operate_source create $WORK_DIR/source2.yaml $SOURCE_ID2
+
+    # generate uncomplete checkpoint
+    dmctl_start_task "$cur/conf/dm-task.yaml" "--remove-meta"
+    check_log_contain_with_retry 'wait loader stop after init checkpoint' $WORK_DIR/worker1/log/dm-worker.log
+    check_log_contain_with_retry 'wait loader stop after init checkpoint' $WORK_DIR/worker2/log/dm-worker.log
+    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
+            "stop-task test"\
+            "\"result\": true" 3
+
+    # restart dm-worker
+    pkill -9 dm-worker.test 2>/dev/null || true
+    check_port_offline $WORKER1_PORT 20
+    check_port_offline $WORKER2_PORT 20
+
+    export GO_FAILPOINTS='github.com/pingcap/dm/loader/WaitLoaderStopBeforeLoadCheckpoint=return(5)'
+    run_dm_worker $WORK_DIR/worker1 $WORKER1_PORT $cur/conf/dm-worker1.toml
+    run_dm_worker $WORK_DIR/worker2 $WORKER2_PORT $cur/conf/dm-worker2.toml
+    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER1_PORT
+    check_rpc_alive $cur/../bin/check_worker_online 127.0.0.1:$WORKER2_PORT
+
+    # stop-task before load checkpoint
+    dmctl_start_task "$cur/conf/dm-task.yaml"
+    check_log_contain_with_retry 'wait loader stop before load checkpoint' $WORK_DIR/worker1/log/dm-worker.log
+    check_log_contain_with_retry 'wait loader stop before load checkpoint' $WORK_DIR/worker2/log/dm-worker.log
+    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
+            "stop-task test"\
+            "\"result\": true" 3
+
+    dmctl_start_task "$cur/conf/dm-task.yaml"
+    check_sync_diff $WORK_DIR $cur/conf/diff_config.toml
+    run_dm_ctl $WORK_DIR "127.0.0.1:$MASTER_PORT" \
+            "stop-task test"\
+            "\"result\": true" 3
+
+    cleanup_data all_mode
+    cleanup_process $*
+
+    export GO_FAILPOINTS=''
+}
+
 function run() {
     run_sql_both_source "SET @@GLOBAL.SQL_MODE='ANSI_QUOTES,NO_AUTO_VALUE_ON_ZERO'"
 
     test_session_config
 
     test_query_timeout
+
+    test_stop_task_before_checkpoint
 
     export GO_FAILPOINTS="github.com/pingcap/dm/dm/worker/TaskCheckInterval=return(\"500ms\")"