Add session revoke (#11908)

Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>

internal/datacoord/server.go

@@ -711,6 +711,7 @@ func (s *Server) Stop() error {
 	s.cluster.Close()
 	s.garbageCollector.close()
 	s.stopServerLoop()
+	s.session.Revoke(time.Second)
 
 	if Params.EnableCompaction {
 		s.stopCompactionTrigger()

internal/datanode/data_node.go

@@ -639,6 +639,8 @@ func (node *DataNode) Stop() error {
 			return err
 		}
 	}
+
+	node.session.Revoke(time.Second)
 	return nil
 }
 

internal/indexcoord/index_coord.go

@@ -285,6 +285,7 @@ func (i *IndexCoord) Stop() error {
 	for _, cb := range i.closeCallbacks {
 		cb()
 	}
+	i.session.Revoke(time.Second)
 	return nil
 }
 

internal/indexnode/indexnode.go

@@ -217,6 +217,7 @@ func (i *IndexNode) Stop() error {
 	for _, cb := range i.closeCallbacks {
 		cb()
 	}
+	i.session.Revoke(time.Second)
 	log.Debug("Index node stopped.")
 	return nil
 }

internal/proxy/proxy.go

@@ -367,6 +367,8 @@ func (node *Proxy) Stop() error {
 		cb()
 	}
 
+	node.session.Revoke(time.Second)
+
 	return nil
 }
 

internal/querycoord/query_coord.go

@@ -208,6 +208,7 @@ func (qc *QueryCoord) Stop() error {
 	qc.UpdateStateCode(internalpb.StateCode_Abnormal)
 
 	qc.loopWg.Wait()
+	qc.session.Revoke(time.Second)
 	return nil
 }
 

internal/querynode/query_node.go

@@ -298,6 +298,7 @@ func (node *QueryNode) Stop() error {
 	if node.statsService != nil {
 		node.statsService.close()
 	}
+	node.session.Revoke(time.Second)
 	return nil
 }
 

internal/rootcoord/root_coord.go

@@ -1170,6 +1170,8 @@ func (c *Core) Stop() error {
 	c.cancel()
 	c.wg.Wait()
 	c.stateCode.Store(internalpb.StateCode_Abnormal)
+	// wait at most one second to revoke
+	c.session.Revoke(time.Second)
 	return nil
 }
 

internal/util/sessionutil/session_util.go

@@ -55,7 +55,7 @@ type Session struct {
 
 	liveCh  <-chan bool
 	etcdCli *clientv3.Client
-	leaseID clientv3.LeaseID
+	leaseID *clientv3.LeaseID
 
 	metaRoot string
 }
@@ -189,7 +189,7 @@ func (s *Session) registerService() (<-chan *clientv3.LeaseKeepAliveResponse, er
 			log.Error("register service", zap.Error(err))
 			return err
 		}
-		s.leaseID = resp.ID
+		s.leaseID = &resp.ID
 
 		sessionJSON, err := json.Marshal(s)
 		if err != nil {
@@ -388,3 +388,18 @@ func (s *Session) LivenessCheck(ctx context.Context, callback func()) {
 		}
 	}
 }
+
+// Revoke revokes the internal leaseID for the session key
+func (s *Session) Revoke(timeout time.Duration) {
+	if s == nil {
+		return
+	}
+	if s.etcdCli == nil || s.leaseID == nil {
+		return
+	}
+	// can NOT use s.ctx, it may be Done here
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	// ignores resp & error, just do best effort to revoke
+	_, _ = s.etcdCli.Revoke(ctx, *s.leaseID)
+}

internal/util/sessionutil/session_util_test.go

@@ -202,3 +202,39 @@ func TestSessionLivenessCheck(t *testing.T) {
 
 	assert.False(t, flag)
 }
+
+func TestSessionRevoke(t *testing.T) {
+	s := &Session{}
+	assert.NotPanics(t, func() {
+		s.Revoke(time.Second)
+	})
+
+	s = (*Session)(nil)
+	assert.NotPanics(t, func() {
+		s.Revoke(time.Second)
+	})
+
+	ctx := context.Background()
+	Params.Init()
+
+	endpoints, err := Params.Load("_EtcdEndpoints")
+	if err != nil {
+		panic(err)
+	}
+	metaRoot := fmt.Sprintf("%d/%s", rand.Int(), DefaultServiceRoot)
+
+	etcdEndpoints := strings.Split(endpoints, ",")
+	etcdKV, err := etcdkv.NewEtcdKV(etcdEndpoints, metaRoot)
+	assert.NoError(t, err)
+	err = etcdKV.RemoveWithPrefix("")
+	assert.NoError(t, err)
+
+	defer etcdKV.Close()
+	defer etcdKV.RemoveWithPrefix("")
+
+	s = NewSession(ctx, metaRoot, etcdEndpoints)
+	s.Init("revoketest", "testAddr", false)
+	assert.NotPanics(t, func() {
+		s.Revoke(time.Second)
+	})
+}