br: use one shot session to close domain ASAP (#36558) (#36593)

ref #36546

br/pkg/glue/glue.go

@@ -28,6 +28,12 @@ type Glue interface {
 
 	// GetVersion gets BR package version to run backup/restore job
 	GetVersion() string
+
+	// UseOneShotSession temporary creates session from store when run backup job.
+	// because we don't have to own domain/session during the whole backup.
+	// we can close domain as soon as possible.
+	// and we must reuse the exists session and don't close it in SQL backup job.
+	UseOneShotSession(store kv.Storage, closeDomain bool, fn func(se Session) error) error
 }
 
 // Session is an abstraction of the session.Session interface.

br/pkg/gluetidb/glue.go

@@ -114,6 +114,45 @@ func (g Glue) GetVersion() string {
 	return g.tikvGlue.GetVersion()
 }
 
+// UseOneShotSession implements glue.Glue.
+func (g Glue) UseOneShotSession(store kv.Storage, closeDomain bool, fn func(glue.Session) error) error {
+	se, err := session.CreateSession(store)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	glueSession := &tidbSession{
+		se: se,
+	}
+	defer func() {
+		se.Close()
+		log.Info("one shot session closed")
+	}()
+	// dom will be created during session.CreateSession.
+	dom, err := session.GetDomain(store)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	// because domain was created during the whole program exists.
+	// and it will register br info to info syncer.
+	// we'd better close it as soon as possible.
+	if closeDomain {
+		defer func() {
+			dom.Close()
+			log.Info("one shot domain closed")
+		}()
+	}
+	err = fn(glueSession)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	return nil
+}
+
+// GetSessionCtx implements glue.Glue
+func (gs *tidbSession) GetSessionCtx() sessionctx.Context {
+	return gs.se
+}
+
 // Execute implements glue.Session.
 func (gs *tidbSession) Execute(ctx context.Context, sql string) error {
 	return gs.ExecuteInternal(ctx, sql)

br/pkg/gluetikv/glue.go

@@ -68,3 +68,8 @@ func (Glue) Record(name string, val uint64) {
 func (Glue) GetVersion() string {
 	return "BR\n" + build.Info()
 }
+
+// UseOneShotSession implements glue.Glue.
+func (g Glue) UseOneShotSession(store kv.Storage, closeDomain bool, fn func(glue.Session) error) error {
+	return nil
+}

br/pkg/task/backup.go

@@ -262,16 +262,19 @@ func RunBackup(c context.Context, g glue.Glue, cmdName string, cfg *BackupConfig
 		statsHandle = mgr.GetDomain().StatsHandle()
 	}
 
-	se, err := g.CreateSession(mgr.GetStorage())
-	if err != nil {
-		return errors.Trace(err)
-	}
-	newCollationEnable, err := se.GetGlobalVariable(tidbNewCollationEnabled)
+	var newCollationEnable string
+	err = g.UseOneShotSession(mgr.GetStorage(), !needDomain, func(se glue.Session) error {
+		newCollationEnable, err = se.GetGlobalVariable(tidbNewCollationEnabled)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		log.Info("get new_collations_enabled_on_first_bootstrap config from system table",
+			zap.String(tidbNewCollationEnabled, newCollationEnable))
+		return nil
+	})
 	if err != nil {
 		return errors.Trace(err)
 	}
-	log.Info("get new_collations_enabled_on_first_bootstrap config from system table",
-		zap.String(tidbNewCollationEnabled, newCollationEnable))
 
 	client, err := backup.NewBackupClient(ctx, mgr)
 	if err != nil {

br/tests/br_full_ddl/run.sh

@@ -82,6 +82,16 @@ if [ "${checksum_count}" -lt "1" ];then
     exit 1
 fi
 
+# when we have backup stats during backup, we cannot close domain during one shot session.
+# so we can check the log count of `one shot domain closed`.
+# we will call UseOneShotSession once to get the value global variable.
+one_shot_session_count=$(cat $LOG | grep "one shot session closed" | wc -l | xargs)
+one_shot_domain_count=$(cat $LOG | grep "one shot domain closed" | wc -l | xargs)
+if [ "${one_shot_session_count}" -ne "1" ] || [ "$one_shot_domain_count" -ne "0" ];then
+    echo "TEST: [$TEST_NAME] fail on one shot session check, $one_shot_session_count, $one_shot_domain_count"
+    exit 1
+fi
+
 echo "backup start without stats..."
 run_br --pd $PD_ADDR backup full -s "local://$TEST_DIR/${DB}_disable_stats" --concurrency 4
 

br/tests/br_incremental_ddl/run.sh

@@ -19,6 +19,7 @@ DB="$TEST_NAME"
 TABLE="usertable"
 ROW_COUNT=100
 PATH="tests/$TEST_NAME:bin:$PATH"
+LOG=/$TEST_DIR/backup.log
 
 echo "load data..."
 # create database
@@ -30,9 +31,24 @@ for i in $(seq $ROW_COUNT); do
     run_sql "INSERT INTO ${DB}.${TABLE}(c1) VALUES ($i);"
 done
 
+
+# Do not log to terminal
+unset BR_LOG_TO_TERM
 # full backup
 echo "full backup start..."
-run_br --pd $PD_ADDR backup table -s "local://$TEST_DIR/$DB/full" --db $DB -t $TABLE
+run_br --pd $PD_ADDR backup table -s "local://$TEST_DIR/$DB/full" --db $DB -t $TABLE --log-file $LOG
+
+# when we backup, we should close domain in one shot session.
+# so we can check the log count of `one shot domain closed` to be 1.
+# we will call UseOneShotSession once to get the value global variable.
+one_shot_session_count=$(cat $LOG | grep "one shot session closed" | wc -l | xargs)
+one_shot_domain_count=$(cat $LOG | grep "one shot domain closed" | wc -l | xargs)
+if [ "${one_shot_session_count}" -ne "1" ] || [ "$one_shot_domain_count" -ne "1" ];then
+    echo "TEST: [$TEST_NAME] fail on one shot session check during backup, $one_shot_session_count, $one_shot_domain_count"
+    exit 1
+fi
+rm -rf $LOG
+
 # run ddls
 echo "run ddls..."
 run_sql "RENAME TABLE ${DB}.${TABLE} to ${DB}.${TABLE}1;"
@@ -54,7 +70,21 @@ done
 # incremental backup
 echo "incremental backup start..."
 last_backup_ts=$(run_br validate decode --field="end-version" -s "local://$TEST_DIR/$DB/full" | grep -oE "^[0-9]+")
-run_br --pd $PD_ADDR backup db -s "local://$TEST_DIR/$DB/inc" --db $DB --lastbackupts $last_backup_ts
+run_br --pd $PD_ADDR backup db -s "local://$TEST_DIR/$DB/inc" --db $DB --lastbackupts $last_backup_ts --log-file $LOG
+
+# when we doing incremental backup, we should close domain in one shot session.
+# so we can check the log count of `one shot domain closed` to be 2.
+# we will call UseOneShotSession twice
+# 1. to get the value global variable.
+# 2. to get all ddl jobs with session.
+one_shot_session_count=$(cat $LOG | grep "one shot session closed" | wc -l | xargs)
+one_shot_domain_count=$(cat $LOG | grep "one shot domain closed" | wc -l | xargs)
+if [ "${one_shot_session_count}" -ne "2" ] || [ "$one_shot_domain_count" -ne "2" ];then
+    echo "TEST: [$TEST_NAME] fail on one shot session check during inc backup, $one_shot_session_count, $one_shot_domain_count"
+    exit 1
+fi
+rm -rf $LOG
+BR_LOG_TO_TERM=1
 
 run_sql "DROP DATABASE $DB;"
 # full restore
@@ -101,4 +131,4 @@ fi
 run_sql "INSERT INTO ${DB}.${TABLE}(c2) VALUES ('1');"
 run_sql "INSERT INTO ${DB}.${TABLE}_rename2(c) VALUES ('1');"
 
-run_sql "DROP DATABASE $DB;"
+run_sql "DROP DATABASE $DB;"

executor/brie.go

@@ -562,3 +562,9 @@ func (gs *tidbGlueSession) Record(name string, value uint64) {
 func (gs *tidbGlueSession) GetVersion() string {
 	return "TiDB\n" + printer.GetTiDBInfo()
 }
+
+// UseOneShotSession implements glue.Glue
+func (gs *tidbGlueSession) UseOneShotSession(store kv.Storage, closeDomain bool, fn func(se glue.Session) error) error {
+	// in SQL backup. we don't need to close domain.
+	return fn(gs)
+}