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.