mcs: fix ticker stop in register

pkg/mcs/discovery/register.go

@@ -80,31 +80,7 @@
 			case _, ok := <-kresp:
 				if !ok {
 					log.Error("keep alive failed", zap.String("key", sr.key))
-					// retry
-					t := time.NewTicker(time.Duration(sr.ttl) * time.Second / 2)
-					for {
-						select {
-						case <-sr.ctx.Done():
-							log.Info("exit register process", zap.String("key", sr.key))
-							t.Stop()
-							return
-						default:
-						}
-
-						<-t.C
-						resp, err := sr.cli.Grant(sr.ctx, sr.ttl)
-						if err != nil {
-							log.Error("grant lease failed", zap.String("key", sr.key), zap.Error(err))
-							t.Stop()
-							continue
-						}
-
-						if _, err := sr.cli.Put(sr.ctx, sr.key, sr.value, clientv3.WithLease(resp.ID)); err != nil {
-							log.Error("put the key failed", zap.String("key", sr.key), zap.Error(err))
-							t.Stop()
-							continue
-						}
-					}
+					kresp = sr.renewKeepalive()
 				}
 			}
 		}
@@ -113,6 +89,36 @@
 	return nil
 }
 
+func (sr *ServiceRegister) renewKeepalive() <-chan *clientv3.LeaseKeepAliveResponse {
+	t := time.NewTicker(time.Duration(sr.ttl) * time.Second / 2)
+	defer t.Stop()
+	for {
+		select {
+		case <-sr.ctx.Done():
+			log.Info("exit register process", zap.String("key", sr.key))
+			return nil
+		case <-t.C:
+			resp, err := sr.cli.Grant(sr.ctx, sr.ttl)
+			if err != nil {
+				log.Error("grant lease failed", zap.String("key", sr.key), zap.Error(err))
+				continue
+			}
+
+			if _, err := sr.cli.Put(sr.ctx, sr.key, sr.value, clientv3.WithLease(resp.ID)); err != nil {
+				log.Error("put the key failed", zap.String("key", sr.key), zap.Error(err))
+				continue
+			}
+
+			kresp, err := sr.cli.KeepAlive(sr.ctx, resp.ID)
+			if err != nil {
+				log.Error("client keep alive failed", zap.String("key", sr.key), zap.Error(err))
+				continue
+			}
+			return kresp
+		}
+	}
+}
+
 // Deregister deregisters the service from etcd.
 func (sr *ServiceRegister) Deregister() error {
 	sr.cancel()

pkg/mcs/discovery/register_test.go

@@ -29,17 +29,13 @@ func TestRegister(t *testing.T) {
 	re := require.New(t)
 	cfg := etcdutil.NewTestSingleConfig(t)
 	etcd, err := embed.StartEtcd(cfg)
-	defer func() {
-		etcd.Close()
-	}()
 	re.NoError(err)
-
 	ep := cfg.LCUrls[0].String()
 	client, err := clientv3.NewFromURL(ep)
 	re.NoError(err)
-
 	<-etcd.Server.ReadyNotify()
-	// with http prefix
+
+	// Test register with http prefix.
 	sr := NewServiceRegister(context.Background(), client, "12345", "test_service", "http://127.0.0.1:1", "http://127.0.0.1:1", 10)
 	re.NoError(err)
 	err = sr.Register()
@@ -49,20 +45,40 @@ func TestRegister(t *testing.T) {
 	re.NoError(err)
 	re.Equal("http://127.0.0.1:1", string(resp.Kvs[0].Value))
 
+	// Test deregister.
 	err = sr.Deregister()
 	re.NoError(err)
 	resp, err = client.Get(context.Background(), sr.key)
 	re.NoError(err)
 	re.Empty(resp.Kvs)
 
+	// Test the case that ctx is canceled.
 	sr = NewServiceRegister(context.Background(), client, "12345", "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
-	re.NoError(err)
 	err = sr.Register()
 	re.NoError(err)
 	sr.cancel()
-	// ensure that the lease is expired
-	time.Sleep(3 * time.Second)
-	resp, err = client.Get(context.Background(), sr.key)
+	re.Empty(getKeyAfterLeaseExpired(re, client, sr.key))
+
+	// Test the case that keepalive is failed when the etcd is restarted.
+	sr = NewServiceRegister(context.Background(), client, "12345", "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
+	err = sr.Register()
 	re.NoError(err)
-	re.Empty(resp.Kvs)
+	re.Equal("127.0.0.1:2", getKeyAfterLeaseExpired(re, client, sr.key))
+	etcd.Close()                // close the etcd to make the keepalive failed
+	time.Sleep(3 * time.Second) // ensure that the lease is expired
+	etcd, err = embed.StartEtcd(cfg)
+	re.NoError(err)
+	<-etcd.Server.ReadyNotify()
+	defer etcd.Close()
+	re.Equal("127.0.0.1:2", getKeyAfterLeaseExpired(re, client, sr.key))
+}
+
+func getKeyAfterLeaseExpired(re *require.Assertions, client *clientv3.Client, key string) string {
+	time.Sleep(3 * time.Second) // ensure that the lease is expired
+	resp, err := client.Get(context.Background(), key)
+	re.NoError(err)
+	if len(resp.Kvs) == 0 {
+		return ""
+	}
+	return string(resp.Kvs[0].Value)
 }

server/server.go

@@ -1861,6 +1861,12 @@ func (s *Server) initTSOPrimaryWatcher() {
 		return nil
 	}
 	deleteFn := func(kv *mvccpb.KeyValue) error {
+		var oldPrimary string
+		v, ok := s.servicePrimaryMap.Load(serviceName)
+		if ok {
+			oldPrimary = v.(string)
+		}
+		log.Info("delete tso primary", zap.String("old-primary", oldPrimary))
 		s.servicePrimaryMap.Delete(serviceName)
 		return nil
 	}