cherry pick #3322 to release-4.0 (#3323)

Signed-off-by: ti-srebot <ti-srebot@pingcap.com> Co-authored-by: 混沌DM <hundundm@gmail.com>
tikv · Jan 5, 2021 · b720955 · b720955
1 parent 7b36ffc
commit b720955
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 8 deletions.
diff --git a/server/id/id.go b/server/id/id.go
@@ -58,8 +58,6 @@ func (alloc *AllocatorImpl) Alloc() (uint64, error) {
 		if err := alloc.generateLocked(); err != nil {
 			return 0, err
 		}
-
-		alloc.base = alloc.end - allocStep
 	}
 
 	alloc.base++
@@ -116,6 +114,7 @@ func (alloc *AllocatorImpl) generateLocked() error {
 	log.Info("idAllocator allocates a new id", zap.Uint64("alloc-id", end))
 	idGauge.WithLabelValues("idalloc").Set(float64(end))
 	alloc.end = end
+	alloc.base = end - allocStep
 	return nil
 }
 

diff --git a/tests/server/id/id_test.go b/tests/server/id/id_test.go
@@ -51,11 +51,11 @@ func (s *testAllocIDSuite) SetUpSuite(c *C) {
 func (s *testAllocIDSuite) TearDownSuite(c *C) {
 	s.cancel()
 }
+
 func (s *testAllocIDSuite) TestID(c *C) {
-	var err error
 	cluster, err := tests.NewTestCluster(s.ctx, 1)
-	defer cluster.Destroy()
 	c.Assert(err, IsNil)
+	defer cluster.Destroy()
 
 	err = cluster.RunInitialServers()
 	c.Assert(err, IsNil)
@@ -96,10 +96,9 @@ func (s *testAllocIDSuite) TestID(c *C) {
 }
 
 func (s *testAllocIDSuite) TestCommand(c *C) {
-	var err error
 	cluster, err := tests.NewTestCluster(s.ctx, 1)
-	defer cluster.Destroy()
 	c.Assert(err, IsNil)
+	defer cluster.Destroy()
 
 	err = cluster.RunInitialServers()
 	c.Assert(err, IsNil)
@@ -121,10 +120,9 @@ func (s *testAllocIDSuite) TestCommand(c *C) {
 }
 
 func (s *testAllocIDSuite) TestMonotonicID(c *C) {
-	var err error
 	cluster, err := tests.NewTestCluster(s.ctx, 2)
-	defer cluster.Destroy()
 	c.Assert(err, IsNil)
+	defer cluster.Destroy()
 
 	err = cluster.RunInitialServers()
 	c.Assert(err, IsNil)
@@ -164,3 +162,33 @@ func (s *testAllocIDSuite) TestMonotonicID(c *C) {
 		last3 = id
 	}
 }
+
+func (s *testAllocIDSuite) TestPDRestart(c *C) {
+	cluster, err := tests.NewTestCluster(s.ctx, 1)
+	c.Assert(err, IsNil)
+	defer cluster.Destroy()
+
+	err = cluster.RunInitialServers()
+	c.Assert(err, IsNil)
+	cluster.WaitLeader()
+	leaderServer := cluster.GetServer(cluster.GetLeader())
+
+	var last uint64
+	for i := uint64(0); i < 10; i++ {
+		id, err := leaderServer.GetAllocator().Alloc()
+		c.Assert(err, IsNil)
+		c.Assert(id, Greater, last)
+		last = id
+	}
+
+	c.Assert(leaderServer.Stop(), IsNil)
+	c.Assert(leaderServer.Run(), IsNil)
+	cluster.WaitLeader()
+
+	for i := uint64(0); i < 10; i++ {
+		id, err := leaderServer.GetAllocator().Alloc()
+		c.Assert(err, IsNil)
+		c.Assert(id, Greater, last)
+		last = id
+	}
+}