address comment

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

@@ -30,6 +30,8 @@ type TiKVModeSwitcher interface {
 	ToImportMode(ctx context.Context)
 	// ToNormalMode switches all TiKV nodes to Normal mode.
 	ToNormalMode(ctx context.Context)
+	// Close closes the switcher's pd client.
+	Close()
 }
 
 // TiKVModeSwitcher is used to switch TiKV nodes between Import and Normal mode.
@@ -56,6 +58,13 @@ func (rc *switcher) ToNormalMode(ctx context.Context) {
 	rc.switchTiKVMode(ctx, sstpb.SwitchMode_Normal)
 }
 
+func (rc *switcher) Close() {
+	if rc.pdCli != nil {
+		rc.pdCli.Close()
+		rc.pdCli = nil
+	}
+}
+
 func (rc *switcher) switchTiKVMode(ctx context.Context, mode sstpb.SwitchMode) {
 	rc.logger.Info("switch tikv mode", zap.Stringer("mode", mode))
 

br/pkg/lightning/importer/get_pre_info.go

@@ -105,6 +105,8 @@ type TargetInfoGetter interface {
 	GetStorageInfo(ctx context.Context) (*pdtypes.StoresInfo, error)
 	// GetEmptyRegionsInfo gets the region information of all the empty regions on the target.
 	GetEmptyRegionsInfo(ctx context.Context) (*pdtypes.RegionsInfo, error)
+	// Close closes TargetInfoGetter's embed pd client.
+	Close()
 }
 
 type preInfoGetterKey string
@@ -266,6 +268,15 @@ func (g *TargetInfoGetterImpl) GetEmptyRegionsInfo(ctx context.Context) (*pdtype
 	return result, nil
 }
 
+// Close closes TargetInfoGetterImpl's embed pd client.
+// It implements the TargetInfoGetter interface.
+func (g *TargetInfoGetterImpl) Close() {
+	if g.pdCli != nil {
+		g.pdCli.Close()
+		g.pdCli = nil
+	}
+}
+
 // PreImportInfoGetterImpl implements the operations to get information used in importing preparation.
 type PreImportInfoGetterImpl struct {
 	cfg              *config.Config
@@ -833,3 +844,9 @@ func (p *PreImportInfoGetterImpl) GetTargetSysVariablesForImport(ctx context.Con
 	p.sysVarsCache = sysVars
 	return sysVars
 }
+
+// Close closes the embed target info getter.
+// It implements the PreImportInfoGetter interface.
+func (p *PreImportInfoGetterImpl) Close() {
+	p.targetInfoGetter.Close()
+}

br/pkg/lightning/importer/import.go

@@ -426,6 +426,7 @@ func NewImportControllerWithPauser(
 		db:      db,
 		tls:     tls,
 		backend: wrapper,
+		pdCli:   pdCli,
 	}
 	preInfoGetter, err := NewPreImportInfoGetter(
 		cfg,
@@ -493,6 +494,9 @@ func NewImportControllerWithPauser(
 func (rc *Controller) Close() {
 	rc.backend.Close()
 	_ = rc.db.Close()
+	if rc.pdCli != nil {
+		rc.pdCli.Close()
+	}
 }
 
 // Run starts the restore task.

br/pkg/lightning/importer/mock/mock.go

@@ -325,3 +325,7 @@ func (t *TargetInfo) IsTableEmpty(_ context.Context, schemaName string, tableNam
 func (*TargetInfo) CheckVersionRequirements(_ context.Context) error {
 	return nil
 }
+
+// Close closes the mock target.
+// It implements the TargetInfoGetter interface.
+func (*TargetInfo) Close() {}

br/pkg/lightning/precheck/precheck.go

@@ -85,4 +85,5 @@ type Checker interface {
 	// If the check is skipped, the returned `CheckResult` is nil
 	Check(ctx context.Context) (*CheckResult, error)
 	GetCheckItemID() CheckItemID
+	// Close() error
 }

disttask/importinto/dispatcher.go

@@ -171,6 +171,7 @@ func (h *flowHandle) switchTiKVMode(ctx context.Context, task *proto.Task) {
 		return
 	}
 	switcher.ToImportMode(ctx)
+	switcher.Close()
 	h.lastSwitchTime.Store(time.Now())
 }
 
@@ -341,6 +342,7 @@ func (h *flowHandle) switchTiKV2NormalMode(ctx context.Context, task *proto.Task
 		return
 	}
 	switcher.ToNormalMode(ctx)
+	switcher.Close()
 
 	// clear it, so next task can switch TiKV mode again.
 	h.lastSwitchTime.Store(time.Time{})