diff --git a/pkg/task/restore.go b/pkg/task/restore.go index 5fffd7ce4..ff2374fa8 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 ca7504f6f..03e987456 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()