From 8638d9a05ca4ba55aae71431eabe618b40454da4 Mon Sep 17 00:00:00 2001
From: kennytm <kennytm@gmail.com>
Date: Thu, 26 Mar 2020 14:58:05 +0800
Subject: [PATCH] task: do not run checksum if restore failed (#209)

---
 pkg/task/restore.go     | 12 ++++++++----
 pkg/task/restore_raw.go | 12 +++++++-----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/pkg/task/restore.go b/pkg/task/restore.go
index 5fffd7ce4c562..ff2374fa84acb 100644
--- a/pkg/task/restore.go
+++ b/pkg/task/restore.go
@@ -229,12 +229,16 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 
 	// Always run the post-work even on error, so we don't stuck in the import
 	// mode or emptied schedulers
-	err = restorePostWork(ctx, client, mgr, clusterCfg)
-	if err != nil {
-		return err
+	if errRestorePostWork := restorePostWork(ctx, client, mgr, clusterCfg); err == nil {
+		err = errRestorePostWork
 	}
 
-	if err = splitPostWork(ctx, client, newTables); err != nil {
+	if errSplitPostWork := splitPostWork(ctx, client, newTables); err == nil {
+		err = errSplitPostWork
+	}
+
+	// If any error happened, return now, don't execute checksum.
+	if err != nil {
 		return err
 	}
 
diff --git a/pkg/task/restore_raw.go b/pkg/task/restore_raw.go
index ca7504f6f2730..03e987456f318 100644
--- a/pkg/task/restore_raw.go
+++ b/pkg/task/restore_raw.go
@@ -46,7 +46,7 @@ func (cfg *RestoreRawConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 }
 
 // RunRestoreRaw starts a raw kv restore task inside the current goroutine.
-func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreRawConfig) error {
+func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreRawConfig) (err error) {
 	defer summary.Summary(cmdName)
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
@@ -113,16 +113,18 @@ func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreR
 	if err != nil {
 		return errors.Trace(err)
 	}
+	defer func() {
+		errPostWork := restorePostWork(ctx, client, mgr, removedSchedulers)
+		if err == nil {
+			err = errPostWork
+		}
+	}()
 
 	err = client.RestoreRaw(cfg.StartKey, cfg.EndKey, files, updateCh)
 	if err != nil {
 		return errors.Trace(err)
 	}
 
-	err = restorePostWork(ctx, client, mgr, removedSchedulers)
-	if err != nil {
-		return errors.Trace(err)
-	}
 	// Restore has finished.
 	updateCh.Close()