lightning: fix panic when user cancel

br/pkg/lightning/backend/local/engine.go

@@ -32,6 +32,7 @@ import (
  "github.com/google/btree"
  "github.com/google/uuid"
  "github.com/pingcap/errors"
+ "github.com/pingcap/failpoint"
  "github.com/pingcap/tidb/br/pkg/lightning/backend"
  "github.com/pingcap/tidb/br/pkg/lightning/backend/kv"
  "github.com/pingcap/tidb/br/pkg/lightning/checkpoints"
@@ -1224,6 +1225,15 @@ func (w *Writer) flushKVs(ctx context.Context) error {
  if err != nil {
  return errors.Trace(err)
  }
+
+ failpoint.Inject("orphanWriterGoRoutine", func() {
+ _ = common.KillMySelf()
+ // mimic we meet context cancel error when `addSST`
+ <-ctx.Done()
+ time.Sleep(5 * time.Second)
+ failpoint.Return(errors.Trace(ctx.Err()))
+ })
+
  err = w.addSST(ctx, meta)
  if err != nil {
  return errors.Trace(err)

br/pkg/lightning/lightning.go

@@ -598,6 +598,11 @@ func (l *Lightning) run(taskCtx context.Context, taskCfg *config.Config, o *opti
  o.logger.Error("restore failed", log.ShortError(err))
  return errors.Trace(err)
  }
+
+ failpoint.Inject("orphanWriterGoRoutine", func() {
+ // don't exit too quickly to expose panic
+ defer time.Sleep(time.Second * 10)
+ })
  defer procedure.Close()
 
  err = procedure.Run(ctx)

br/pkg/lightning/restore/table_restore.go

@@ -501,6 +501,7 @@ func (tr *TableRestore) restoreEngine(
  metrics, _ := metric.FromContext(ctx)
 
  // Restore table data
+ChunkLoop:
  for chunkIndex, chunk := range cp.Chunks {
  if rc.status != nil && rc.status.backend == config.BackendTiDB {
  rc.status.FinishedFileSize.Add(chunk.Chunk.Offset - chunk.Key.Offset)
@@ -524,9 +525,15 @@ func (tr *TableRestore) restoreEngine(
  }
  checkFlushLock.Unlock()
 
+ failpoint.Inject("orphanWriterGoRoutine", func() {
+ if chunkIndex > 0 {
+ <-pCtx.Done()
+ }
+ })
+
  select {
  case <-pCtx.Done():
- return nil, pCtx.Err()
+ break ChunkLoop
  default:
  }
 
@@ -615,6 +622,11 @@ func (tr *TableRestore) restoreEngine(
  }
 
  wg.Wait()
+ select {
+ case <-pCtx.Done():
+ return nil, pCtx.Err()
+ default:
+ }
 
  // Report some statistics into the log for debugging.
  totalKVSize := uint64(0)

br/tests/lightning_checkpoint_chunks/run.sh

@@ -48,6 +48,11 @@ for i in $(seq "$CHUNK_COUNT"); do
  done
 done
 
+PKG="github.com/pingcap/tidb/br/pkg/lightning"
+export GO_FAILPOINTS="$PKG/backend/local/orphanWriterGoRoutine=return();$PKG/restore/orphanWriterGoRoutine=return();$PKG/orphanWriterGoRoutine=return()"
+# test won't panic
+do_run_lightning config
+
 # Set the failpoint to kill the lightning instance as soon as
 # one file (after writing totally $ROW_COUNT rows) is imported.
 # If checkpoint does work, this should kill exactly $CHUNK_COUNT instances of lightnings.