tso: fix the corner case that may cause TSO fallback (#4885)

server/server.go

@@ -1238,7 +1238,14 @@ func (s *Server) campaignLeader() {
 		log.Error("failed to initialize the global TSO allocator", errs.ZapError(err))
 		return
 	}
-	defer s.tsoAllocatorManager.ResetAllocatorGroup(tso.GlobalDCLocation)
+	defer func() {
+		s.tsoAllocatorManager.ResetAllocatorGroup(tso.GlobalDCLocation)
+		failpoint.Inject("updateAfterResetTSO", func() {
+			if err = alllocator.UpdateTSO(); err != nil {
+				panic(err)
+			}
+		})
+	}()
 
 	if err := s.reloadConfigFromKV(); err != nil {
 		log.Error("failed to reload configuration", errs.ZapError(err))

server/tso/tso.go

@@ -70,9 +70,13 @@ type timestampOracle struct {
 	dcLocation    string
 }
 
-func (t *timestampOracle) setTSOPhysical(next time.Time) {
+func (t *timestampOracle) setTSOPhysical(next time.Time, force bool) {
 	t.tsoMux.Lock()
 	defer t.tsoMux.Unlock()
+	// Do not update the zero physical time if the `force` flag is false.
+	if t.tsoMux.physical == typeutil.ZeroTime && !force {
+		return
+	}
 	// make sure the ts won't fall back
 	if typeutil.SubTSOPhysicalByWallClock(next, t.tsoMux.physical) > 0 {
 		t.tsoMux.physical = next
@@ -190,7 +194,7 @@ func (t *timestampOracle) SyncTimestamp(leadership *election.Leadership) error {
 	tsoCounter.WithLabelValues("sync_ok", t.dcLocation).Inc()
 	log.Info("sync and save timestamp", zap.Time("last", last), zap.Time("save", save), zap.Time("next", next))
 	// save into memory
-	t.setTSOPhysical(next)
+	t.setTSOPhysical(next, true)
 	return nil
 }
 
@@ -266,6 +270,9 @@ func (t *timestampOracle) resetUserTimestamp(leadership *election.Leadership, ts
 // 1. The saved time is monotonically increasing.
 // 2. The physical time is monotonically increasing.
 // 3. The physical time is always less than the saved timestamp.
+//
+// NOTICE: this function should be called after the TSO in memory has been initialized
+// and should not be called when the TSO in memory has been reset anymore.
 func (t *timestampOracle) UpdateTimestamp(leadership *election.Leadership) error {
 	prevPhysical, prevLogical := t.getTSO()
 	tsoGauge.WithLabelValues("tso", t.dcLocation).Set(float64(prevPhysical.Unix()))
@@ -315,7 +322,7 @@ func (t *timestampOracle) UpdateTimestamp(leadership *election.Leadership) error
 		}
 	}
 	// save into memory
-	t.setTSOPhysical(next)
+	t.setTSOPhysical(next, false)
 
 	return nil
 }

tests/client/client_test.go

@@ -182,6 +182,44 @@ func (s *clientTestSuite) TestLeaderTransfer(c *C) {
 	wg.Wait()
 }
 
+// More details can be found in this issue: https://github.com/tikv/pd/issues/4884
+func (s *clientTestSuite) TestUpdateAfterResetTSO(c *C) {
+	cluster, err := tests.NewTestCluster(s.ctx, 2)
+	c.Assert(err, IsNil)
+	defer cluster.Destroy()
+
+	endpoints := s.runServer(c, cluster)
+	cli := s.setupCli(c, endpoints, false)
+
+	testutil.WaitUntil(c, func(c *C) bool {
+		_, _, err := cli.GetTS(context.TODO())
+		return err == nil
+	})
+	// Transfer leader to trigger the TSO resetting.
+	c.Assert(failpoint.Enable("github.com/tikv/pd/server/updateAfterResetTSO", "return(true)"), IsNil)
+	oldLeaderName := cluster.WaitLeader()
+	err = cluster.GetServer(oldLeaderName).ResignLeader()
+	c.Assert(err, IsNil)
+	c.Assert(failpoint.Disable("github.com/tikv/pd/server/updateAfterResetTSO"), IsNil)
+	newLeaderName := cluster.WaitLeader()
+	c.Assert(newLeaderName, Not(Equals), oldLeaderName)
+	// Request a new TSO.
+	testutil.WaitUntil(c, func(c *C) bool {
+		_, _, err := cli.GetTS(context.TODO())
+		return err == nil
+	})
+	// Transfer leader back.
+	c.Assert(failpoint.Enable("github.com/tikv/pd/server/tso/delaySyncTimestamp", `return(true)`), IsNil)
+	err = cluster.GetServer(newLeaderName).ResignLeader()
+	c.Assert(err, IsNil)
+	// Should NOT panic here.
+	testutil.WaitUntil(c, func(c *C) bool {
+		_, _, err := cli.GetTS(context.TODO())
+		return err == nil
+	})
+	c.Assert(failpoint.Disable("github.com/tikv/pd/server/tso/delaySyncTimestamp"), IsNil)
+}
+
 func (s *clientTestSuite) TestTSOAllocatorLeader(c *C) {
 	dcLocationConfig := map[string]string{
 		"pd1": "dc-1",