fix tests

Signed-off-by: Ryan Leung <rleungx@gmail.com>
rleungx · Dec 11, 2024 · dddb935 · dddb935
1 parent 11fbaa4
commit dddb935
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 16 deletions.
diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
@@ -419,32 +419,40 @@ func (c *RaftCluster) checkSchedulingService() {
 // checkTSOService checks the TSO service.
 func (c *RaftCluster) checkTSOService() {
 	if !c.opt.GetMicroServiceConfig().IsTSODynamicSwitchingEnabled() {
-		if err := c.startTSOJobsIfNeeded(); err != nil {
+		if err := c.switchToInternalTSO(); err != nil {
 			log.Error("failed to start TSO jobs", errs.ZapError(err))
 			return
 		}
-		if c.IsServiceIndependent(constant.TSOServiceName) {
-			c.UnsetServiceIndependent(constant.TSOServiceName)
-		}
 		return
 	}
 
 	servers, err := discovery.Discover(c.etcdClient, constant.TSOServiceName)
 	if err != nil || len(servers) == 0 {
-		if err := c.startTSOJobsIfNeeded(); err != nil {
-			log.Error("failed to start TSO jobs", errs.ZapError(err))
+		if err := c.switchToInternalTSO(); err != nil {
+			log.Error("failed to switch to internal TSO", errs.ZapError(err))
 			return
 		}
-		if c.IsServiceIndependent(constant.TSOServiceName) {
-			log.Info("TSO is provided by PD")
-			c.UnsetServiceIndependent(constant.TSOServiceName)
-		}
-	} else {
-		c.stopTSOJobsIfNeeded()
-		if !c.IsServiceIndependent(constant.TSOServiceName) {
-			log.Info("TSO is provided by TSO server")
-			c.SetServiceIndependent(constant.TSOServiceName)
-		}
+	} else if len(servers) > 0 {
+		c.switchToExternalTSO()
+	}
+}
+
+func (c *RaftCluster) switchToInternalTSO() error {
+	if err := c.startTSOJobsIfNeeded(); err != nil {
+		return err
+	}
+	if c.IsServiceIndependent(constant.TSOServiceName) {
+		c.UnsetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to internal TSO")
+	}
+	return nil
+}
+
+func (c *RaftCluster) switchToExternalTSO() {
+	c.stopTSOJobsIfNeeded()
+	if !c.IsServiceIndependent(constant.TSOServiceName) {
+		c.SetServiceIndependent(constant.TSOServiceName)
+		log.Info("successfully switched to external TSO")
 	}
 }
 

diff --git a/tests/server/tso/tso_test.go b/tests/server/tso/tso_test.go
@@ -132,6 +132,7 @@ func checkAndReturnTimestampResponse(re *require.Assertions, req *pdpb.TsoReques
 	re.GreaterOrEqual(uint32(timestamp.GetLogical())>>timestamp.GetSuffixBits(), req.GetCount())
 	return timestamp
 }
+
 func TestLogicalOverflow(t *testing.T) {
 	re := require.New(t)