refactor: extract method checkKilled()

Signed-off-by: ekexium <[email protected]>
tikv · Jan 5, 2024 · 15d27cb · 15d27cb
1 parent 3521491
commit 15d27cb
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 30 deletions.
diff --git a/config/retry/backoff.go b/config/retry/backoff.go
@@ -217,15 +217,9 @@ func (b *Backoffer) BackoffWithCfgAndMaxSleep(cfg *Config, maxSleepMs int, err e
 		atomic.AddInt64(&detail.BackoffCount, 1)
 	}
 
-	if b.vars != nil && b.vars.Killed != nil {
-		killed := atomic.LoadUint32(b.vars.Killed)
-		if killed != 0 {
-			logutil.BgLogger().Info(
-				"backoff stops because a killed signal is received",
-				zap.Uint32("signal", killed),
-			)
-			return errors.WithStack(tikverr.ErrQueryInterrupted)
-		}
+	err2 := b.CheckKilled()
+	if err2 != nil {
+		return err2
 	}
 
 	var startTs interface{}
@@ -387,3 +381,17 @@ func (b *Backoffer) longestSleepCfg() (*Config, int) {
 	}
 	return nil, 0
 }
+
+func (b *Backoffer) CheckKilled() error {
+	if b.vars != nil && b.vars.Killed != nil {
+		killed := atomic.LoadUint32(b.vars.Killed)
+		if killed != 0 {
+			logutil.BgLogger().Info(
+				"backoff stops because a killed signal is received",
+				zap.Uint32("signal", killed),
+			)
+			return errors.WithStack(tikverr.ErrQueryInterrupted)
+		}
+	}
+	return nil
+}
diff --git a/internal/client/retry/backoff.go b/internal/client/retry/backoff.go
@@ -217,28 +217,38 @@ func (b *Backoffer) BackoffWithCfgAndMaxSleep(cfg *Config, maxSleepMs int, err e
 		atomic.AddInt64(&detail.BackoffCount, 1)
 	}
 
-	if b.vars != nil && b.vars.Killed != nil {
-		killed := atomic.LoadUint32(b.vars.Killed)
-		if killed != 0 {
-			logutil.BgLogger().Info(
-				"backoff stops because a killed signal is received",
-				zap.Uint32("signal", killed),
-			)
-			return errors.WithStack(tikverr.ErrQueryInterrupted)
-		}
+	err2 := b.checkKilled()
+	if err2 != nil {
+		return err2
 	}
 
 	var startTs interface{}
 	if ts := b.ctx.Value(TxnStartKey); ts != nil {
 		startTs = ts
 	}
-	logutil.Logger(b.ctx).Debug("retry later",
+	logutil.Logger(b.ctx).Debug(
+		"retry later",
 		zap.Error(err),
 		zap.Int("totalSleep", b.totalSleep),
 		zap.Int("excludedSleep", b.excludedSleep),
 		zap.Int("maxSleep", b.maxSleep),
 		zap.Stringer("type", cfg),
-		zap.Reflect("txnStartTS", startTs))
+		zap.Reflect("txnStartTS", startTs),
+	)
+	return nil
+}
+
+func (b *Backoffer) checkKilled() error {
+	if b.vars != nil && b.vars.Killed != nil {
+		killed := atomic.LoadUint32(b.vars.Killed)
+		if killed != 0 {
+			logutil.BgLogger().Info(
+				"backoff stops because a killed signal is received",
+				zap.Uint32("signal", killed),
+			)
+			return errors.WithStack(tikverr.ErrQueryInterrupted)
+		}
+	}
 	return nil
 }
 

diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
@@ -1477,16 +1477,8 @@ func (s *RegionRequestSender) SendReqCtx(
 		}
 
 		// recheck whether the session/query is killed during the Next()
-		boVars := bo.GetVars()
-		if boVars != nil && boVars.Killed != nil {
-			killed := atomic.LoadUint32(boVars.Killed)
-			if killed != 0 {
-				logutil.BgLogger().Info(
-					"backoff stops because a killed signal is received",
-					zap.Uint32("signal", killed),
-				)
-				return nil, nil, retryTimes, errors.WithStack(tikverr.ErrQueryInterrupted)
-			}
+		if err2 := bo.CheckKilled(); err != nil {
+			return nil, nil, retryTimes, err2
 		}
 		if val, err := util.EvalFailpoint("mockRetrySendReqToRegion"); err == nil {
 			if val.(bool) {