br: use online to disable import mode (#70221)

close #63505

Author: Leavrth

br/pkg/restore/BUILD.bazel

@@ -60,7 +60,7 @@ go_test(
     ],
     embed = [":restore"],
     flaky = True,
-    shard_count = 21,
+    shard_count = 22,
     deps = [
         "//br/pkg/conn",
         "//br/pkg/mock",

br/pkg/restore/import_mode_switcher.go

@@ -242,14 +242,17 @@ func RestorePostWork(
 	ctx context.Context,
 	switcher *ImportModeSwitcher,
 	restoreSchedulers pdutil.UndoFunc,
+	isOnline bool,
 ) {
 	if ctx.Err() != nil {
 		log.Warn("context canceled, try shutdown")
 		ctx = context.Background()
 	}
 
-	if err := switcher.SwitchToNormalMode(ctx); err != nil {
-		log.Warn("fail to switch to normal mode", zap.Error(err))
+	if !isOnline {
+		if err := switcher.SwitchToNormalMode(ctx); err != nil {
+			log.Warn("fail to switch to normal mode", zap.Error(err))
+		}
 	}
 	if err := restoreSchedulers(ctx); err != nil {
 		log.Warn("failed to restore PD schedulers", zap.Error(err))

br/pkg/restore/import_mode_switcher_test.go

@@ -35,18 +35,29 @@ import (
 type mockImportServer struct {
 	import_sstpb.ImportSSTServer
 
+	mu    sync.Mutex
+	modes []import_sstpb.SwitchMode
 	count int
 	ch    chan struct{}
 }
 
 func (s *mockImportServer) SwitchMode(_ context.Context, req *import_sstpb.SwitchModeRequest) (*import_sstpb.SwitchModeResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.modes = append(s.modes, req.GetMode())
 	s.count -= 1
-	if s.count == 0 {
+	if s.count == 0 && s.ch != nil {
 		s.ch <- struct{}{}
 	}
 	return &import_sstpb.SwitchModeResponse{}, nil
 }
 
+func (s *mockImportServer) Modes() []import_sstpb.SwitchMode {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return append([]import_sstpb.SwitchMode(nil), s.modes...)
+}
+
 func TestRestorePreWork(t *testing.T) {
 	ctx := context.Background()
 	lis, err := net.Listen("tcp", ":0")
@@ -62,7 +73,9 @@ func TestRestorePreWork(t *testing.T) {
 	go func() {
 		defer wg.Done()
 		err := s.Serve(lis)
-		require.NoError(t, err)
+		if err != nil && err != grpc.ErrServerStopped {
+			require.NoError(t, err)
+		}
 	}()
 
 	pdClient := split.NewFakePDClient([]*metapb.Store{
@@ -109,7 +122,7 @@ func TestRestorePreWork(t *testing.T) {
 		}
 	}
 	<-ch
-	restore.RestorePostWork(ctx, switcher, undo)
+	restore.RestorePostWork(ctx, switcher, undo, false)
 	// check the cfg done
 	{
 		cfgs, err := pdHTTPCli.GetConfig(context.TODO())
@@ -126,4 +139,53 @@ func TestRestorePreWork(t *testing.T) {
 
 	s.Stop()
 	lis.Close()
+	wg.Wait()
+}
+
+func TestRestorePostWorkOnlineSkipsNormalMode(t *testing.T) {
+	ctx := context.Background()
+	lis, err := net.Listen("tcp", ":0")
+	require.NoError(t, err)
+	addr := lis.Addr().String()
+
+	s := grpc.NewServer()
+	importServer := &mockImportServer{}
+	import_sstpb.RegisterImportSSTServer(s, importServer)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err := s.Serve(lis)
+		if err != nil && err != grpc.ErrServerStopped {
+			require.NoError(t, err)
+		}
+	}()
+	defer func() {
+		s.Stop()
+		lis.Close()
+		wg.Wait()
+	}()
+
+	pdClient := split.NewFakePDClient([]*metapb.Store{
+		{
+			Id:      1,
+			Address: addr,
+		},
+	}, false, nil)
+	switcher := restore.NewImportModeSwitcher(pdClient, time.Hour, nil)
+	require.NoError(t, switcher.GoSwitchToImportMode(ctx))
+	defer func() {
+		require.NoError(t, switcher.SwitchToNormalMode(ctx))
+	}()
+	require.Equal(t, []import_sstpb.SwitchMode{import_sstpb.SwitchMode_Import}, importServer.Modes())
+
+	restoredSchedulers := false
+	restore.RestorePostWork(ctx, switcher, func(context.Context) error {
+		restoredSchedulers = true
+		return nil
+	}, true)
+
+	require.True(t, restoredSchedulers)
+	require.Equal(t, []import_sstpb.SwitchMode{import_sstpb.SwitchMode_Import}, importServer.Modes())
 }

br/pkg/restore/log_client/client.go

@@ -372,6 +372,7 @@ func (rc *LogClient) RestoreSSTFileSets(
 	ctx context.Context,
 	backupFileSets restore.BatchBackupFileSet,
 	importModeSwitcher *restore.ImportModeSwitcher,
+	online bool,
 	snapshotRestoreDataSize uint64,
 	checkpointCompactedSSTSize uint64,
 	onProgress func(int64),
@@ -390,16 +391,18 @@ func (rc *LogClient) RestoreSSTFileSets(
 		return errors.Trace(err)
 	}
 
-	err := importModeSwitcher.GoSwitchToImportMode(ctx)
-	if err != nil {
-		return errors.Trace(err)
-	}
-	defer func() {
-		switchErr := importModeSwitcher.SwitchToNormalMode(ctx)
-		if switchErr != nil {
-			log.Warn("[Compacted SST Restore] Failed to switch back to normal mode after restoration.", zap.Error(switchErr))
+	if !online {
+		err := importModeSwitcher.GoSwitchToImportMode(ctx)
+		if err != nil {
+			return errors.Trace(err)
 		}
-	}()
+		defer func() {
+			switchErr := importModeSwitcher.SwitchToNormalMode(ctx)
+			if switchErr != nil {
+				log.Warn("[Compacted SST Restore] Failed to switch back to normal mode after restoration.", zap.Error(switchErr))
+			}
+		}()
+	}
 
 	log.Info("[Compacted SST Restore] Start to restore SST files",
 		zap.Int("sst-file-count", len(backupFileSets)))
@@ -413,21 +416,24 @@ func (rc *LogClient) RestoreSSTFileSets(
 	// where batch processing may lead to increased complexity and potential inefficiencies.
 	// TODO: Future enhancements may explore the feasibility of reintroducing batch restoration
 	// while maintaining optimal performance and resource utilization.
-	err = rc.sstRestoreManager.restorer.GoRestore(onProgress, backupFileSets)
-	if err != nil {
+	if err := rc.sstRestoreManager.restorer.GoRestore(onProgress, backupFileSets); err != nil {
 		return errors.Trace(err)
 	}
-	err = rc.sstRestoreManager.restorer.WaitUntilFinish()
+	err := rc.sstRestoreManager.restorer.WaitUntilFinish()
 
+	var totalKvs, totalBytes, totalSize uint64
 	for _, files := range backupFileSets {
 		for _, f := range files.SSTFiles {
-			log.Info("Collected file.", zap.Uint64("total_kv", f.TotalKvs), zap.Uint64("total_bytes", f.TotalBytes), zap.Uint64("size", f.Size_))
+			totalKvs += f.TotalKvs
+			totalBytes += f.TotalBytes
+			totalSize += f.Size_
 			atomic.AddUint64(&rc.restoreStat.restoreSSTKVCount, f.TotalKvs)
 			atomic.AddUint64(&rc.restoreStat.restoreSSTKVSize, f.TotalBytes)
 			atomic.AddUint64(&rc.restoreStat.restoreSSTPhySize, f.Size_)
 		}
 	}
 	atomic.AddUint64(&rc.restoreStat.restoreSSTTakes, uint64(time.Since(begin)))
+	log.Info("Collected files", zap.Uint64("total_kv", totalKvs), zap.Uint64("total_bytes", totalBytes), zap.Uint64("total_size", totalSize))
 	return err
 }
 

br/pkg/task/restore.go

@@ -1659,7 +1659,7 @@ func runSnapshotRestore(c context.Context, mgr *conn.Mgr, g glue.Glue, cmdName s
 		log.Info("start to restore pd scheduler")
 		// run the post-work to avoid being stuck in the import
 		// mode or emptied schedulers.
-		restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulersFunc)
+		restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulersFunc, cfg.Online)
 		log.Info("finish restoring pd scheduler")
 	}()
 

br/pkg/task/restore_raw.go

@@ -161,7 +161,7 @@ func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreR
 	if err != nil {
 		return errors.Trace(err)
 	}
-	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulers)
+	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulers, cfg.Online)
 
 	start := time.Now()
 	err = client.GetRestorer(nil).GoRestore(onProgress, restore.CreateUniqueFileSets(files))

br/pkg/task/restore_txn.go

@@ -101,7 +101,7 @@ func RunRestoreTxn(c context.Context, g glue.Glue, cmdName string, cfg *Config)
 	if err != nil {
 		return errors.Trace(err)
 	}
-	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulers)
+	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulers, false)
 
 	err = client.GetRestorer(nil).GoRestore(onProgress, restore.CreateUniqueFileSets(files))
 	if err != nil {

br/pkg/task/stream.go

@@ -1742,7 +1742,7 @@ func restoreStream(
 
 	// Always run the post-work even on error, so we don't stuck in the import
 	// mode or emptied schedulers
-	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulersFunc)
+	defer restore.RestorePostWork(ctx, importModeSwitcher, restoreSchedulersFunc, cfg.Online)
 
 	migs, err := client.GetLockedMigrations(ctx)
 	if err != nil {
@@ -1878,6 +1878,7 @@ func restoreStream(
 			ctx,
 			sstFileSets,
 			importModeSwitcher,
+			cfg.Online,
 			cfg.snapshotRestoreDataSize,
 			checkpointSSTSize,
 			p.IncBy,