cherry pick tikv#3322 to release-5.0-rc

Signed-off-by: Ryan Leung <rleungx@gmail.com>
ti-srebot · Jan 4, 2021 · c47efdd · c47efdd
1 parent 5fa7a08
commit c47efdd
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
+	}
+}