enhance: remove watch at session liveness check (#45968)

issue: #45724

---------

Signed-off-by: chyezh <[email protected]>

Author: chyezh

cmd/tools/migration/migration/runner.go

@@ -165,7 +165,6 @@ func (r *Runner) CheckSessions() error {
 
 func (r *Runner) RegisterSession() error {
 	r.session.Register()
-	r.session.LivenessCheck(r.ctx, func() {})
 	return nil
 }
 
@@ -246,7 +245,7 @@ func (r *Runner) waitUntilSessionExpired() {
 }
 
 func (r *Runner) Stop() {
-	r.session.Revoke(time.Second)
+	r.session.Stop()
 	r.waitUntilSessionExpired()
 	r.cancel()
 	r.wg.Wait()

internal/coordinator/mix_coord.go

@@ -114,10 +114,6 @@ func (s *mixCoordImpl) Register() error {
 	afterRegister := func() {
 		metrics.NumNodes.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), typeutil.MixCoordRole).Inc()
 		log.Info("MixCoord Register Finished")
-		s.session.LivenessCheck(s.ctx, func() {
-			log.Error("MixCoord disconnected from etcd, process will exit", zap.Int64("serverID", s.session.GetServerID()))
-			os.Exit(1)
-		})
 	}
 	if s.enableActiveStandBy {
 		go func() {

internal/datanode/data_node.go

@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
-	"os"
 	"sync"
 	"time"
 
@@ -157,12 +156,6 @@ func (node *DataNode) Register() error {
 
 	metrics.NumNodes.WithLabelValues(fmt.Sprint(node.GetNodeID()), typeutil.DataNodeRole).Inc()
 	log.Info("DataNode Register Finished")
-	// Start liveness check
-	node.session.LivenessCheck(node.ctx, func() {
-		log.Error("Data Node disconnected from etcd, process will exit", zap.Int64("Server Id", node.GetSession().ServerID))
-		os.Exit(1)
-	})
-
 	return nil
 }
 

internal/distributed/streamingnode/service.go

@@ -18,7 +18,6 @@ package streamingnode
 
 import (
 	"context"
-	"os"
 	"strconv"
 	"sync"
 	"time"
@@ -251,7 +250,7 @@ func (s *Server) start() (err error) {
 		return errors.Wrap(err, "StreamingNode start gRPC server fail")
 	}
 	// Register current server to etcd.
-	s.registerSessionToETCD()
+	s.session.Register()
 
 	s.componentState.OnInitialized(s.session.ServerID)
 	return nil
@@ -382,13 +381,3 @@ func (s *Server) startGPRCServer(ctx context.Context) error {
 	funcutil.CheckGrpcReady(ctx, errCh)
 	return <-errCh
 }
-
-// registerSessionToETCD registers current server to etcd.
-func (s *Server) registerSessionToETCD() {
-	s.session.Register()
-	// start liveness check
-	s.session.LivenessCheck(context.Background(), func() {
-		log.Ctx(s.ctx).Error("StreamingNode disconnected from etcd, process will exit", zap.Int64("Server Id", paramtable.GetNodeID()))
-		os.Exit(1)
-	})
-}

internal/proxy/proxy.go

@@ -20,7 +20,6 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
-	"os"
 	"sync"
 	"time"
 
@@ -157,10 +156,6 @@ func (node *Proxy) Register() error {
 	node.session.Register()
 	metrics.NumNodes.WithLabelValues(fmt.Sprint(paramtable.GetNodeID()), typeutil.ProxyRole).Inc()
 	log.Info("Proxy Register Finished")
-	node.session.LivenessCheck(node.ctx, func() {
-		log.Error("Proxy disconnected from etcd, process will exit", zap.Int64("Server Id", node.session.ServerID))
-		os.Exit(1)
-	})
 	// TODO Reset the logger
 	// Params.initLogCfg()
 	return nil

internal/querycoordv2/mocks/querynode.go

@@ -21,7 +21,6 @@ import (
 	"net"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/mock"
 	clientv3 "go.etcd.io/etcd/client/v3"
@@ -149,7 +148,7 @@ func (node *MockQueryNode) Stopping() {
 func (node *MockQueryNode) Stop() {
 	node.cancel()
 	node.server.GracefulStop()
-	node.session.Revoke(time.Second)
+	node.session.Stop()
 }
 
 func (node *MockQueryNode) getAllChannels() []*querypb.ChannelVersionInfo {

internal/querynodev2/server.go

@@ -32,7 +32,6 @@ import "C"
 import (
 	"context"
 	"fmt"
-	"os"
 	"plugin"
 	"strings"
 	"sync"
@@ -182,10 +181,6 @@ func (node *QueryNode) Register() error {
 	node.session.Register()
 	// start liveness check
 	metrics.NumNodes.WithLabelValues(fmt.Sprint(node.GetNodeID()), typeutil.QueryNodeRole).Inc()
-	node.session.LivenessCheck(node.ctx, func() {
-		log.Ctx(node.ctx).Error("Query Node disconnected from etcd, process will exit", zap.Int64("Server Id", paramtable.GetNodeID()))
-		os.Exit(1)
-	})
 	return nil
 }
 

internal/util/sessionutil/mock_session.go

@@ -16,9 +16,6 @@
 package sessionutil
 
 import (
-	"context"
-	"time"
-
 	"github.com/blang/semver/v4"
 )
 
@@ -36,9 +33,7 @@ type SessionInterface interface {
 	GoingStop() error
 	WatchServices(prefix string, revision int64, rewatch Rewatch) (watcher SessionWatcher)
 	WatchServicesWithVersionRange(prefix string, r semver.Range, revision int64, rewatch Rewatch) (watcher SessionWatcher)
-	LivenessCheck(ctx context.Context, callback func())
 	Stop()
-	Revoke(timeout time.Duration)
 	UpdateRegistered(b bool)
 	Registered() bool
 	SetDisconnected(b bool)