Remove useCustomConfig and simpilify the session type (#23166)

Signed-off-by: yah01 <yang.cen@zilliz.com>

cmd/tools/migration/migration/runner.go

@@ -88,8 +88,13 @@ func (r *Runner) initEtcdCli() {
 func (r *Runner) init() {
 	r.initEtcdCli()
 
-	r.session = sessionutil.NewSession(r.ctx, r.cfg.EtcdCfg.MetaRootPath.GetValue(), r.etcdCli,
-		sessionutil.WithCustomConfigEnable(), sessionutil.WithTTL(60), sessionutil.WithRetryTimes(30))
+	r.session = sessionutil.NewSession(
+		r.ctx,
+		r.cfg.EtcdCfg.MetaRootPath.GetValue(),
+		r.etcdCli,
+		sessionutil.WithTTL(60),
+		sessionutil.WithRetryTimes(30),
+	)
 	// address not important here.
 	address := time.Now().String()
 	r.address = address

internal/util/sessionutil/session_util.go

@@ -107,7 +107,6 @@ type Session struct {
 	enableActiveStandBy bool
 	activeKey           string
 
-	useCustomConfig   bool
 	sessionTTL        int64
 	sessionRetryTimes int64
 	reuseNodeID       bool
@@ -115,10 +114,6 @@ type Session struct {
 
 type SessionOption func(session *Session)
 
-func WithCustomConfigEnable() SessionOption {
-	return func(session *Session) { session.useCustomConfig = true }
-}
-
 func WithTTL(ttl int64) SessionOption {
 	return func(session *Session) { session.sessionTTL = ttl }
 }
@@ -204,9 +199,8 @@ func NewSession(ctx context.Context, metaRoot string, client *clientv3.Client, o
 		Version:  common.Version,
 
 		// options
-		useCustomConfig:   false,
-		sessionTTL:        60,
-		sessionRetryTimes: 30,
+		sessionTTL:        paramtable.Get().CommonCfg.SessionTTL.GetAsInt64(),
+		sessionRetryTimes: paramtable.Get().CommonCfg.SessionRetryTimes.GetAsInt64(),
 		reuseNodeID:       true,
 	}
 
@@ -247,14 +241,11 @@ func (s *Session) Init(serverName, address string, exclusive bool, triggerKill b
 	s.Exclusive = exclusive
 	s.TriggerKill = triggerKill
 	s.checkIDExist()
-	// TO AVOID PANIC IN MIGRATION SCRIPT.
-	if !s.useCustomConfig {
-		serverID, err := s.getServerID()
-		if err != nil {
-			panic(err)
-		}
-		s.ServerID = serverID
+	serverID, err := s.getServerID()
+	if err != nil {
+		panic(err)
 	}
+	s.ServerID = serverID
 	log.Info("start server", zap.String("name", serverName), zap.String("address", address), zap.Int64("id", s.ServerID))
 }
 
@@ -374,15 +365,8 @@ func (s *Session) registerService() (<-chan *clientv3.LeaseKeepAliveResponse, er
 	var ch <-chan *clientv3.LeaseKeepAliveResponse
 	log.Debug("service begin to register to etcd", zap.String("serverName", s.ServerName), zap.Int64("ServerID", s.ServerID))
 
-	ttl := s.sessionTTL
-	retryTimes := s.sessionRetryTimes
-	if !s.useCustomConfig {
-		ttl = paramtable.Get().CommonCfg.SessionTTL.GetAsInt64()
-		retryTimes = paramtable.Get().CommonCfg.SessionRetryTimes.GetAsInt64()
-	}
-
 	registerFn := func() error {
-		resp, err := s.etcdCli.Grant(s.ctx, ttl)
+		resp, err := s.etcdCli.Grant(s.ctx, s.sessionTTL)
 		if err != nil {
 			log.Error("register service", zap.Error(err))
 			return err
@@ -424,7 +408,7 @@ func (s *Session) registerService() (<-chan *clientv3.LeaseKeepAliveResponse, er
 		log.Info("Service registered successfully", zap.String("ServerName", s.ServerName), zap.Int64("serverID", s.ServerID))
 		return nil
 	}
-	err := retry.Do(s.ctx, registerFn, retry.Attempts(uint(retryTimes)))
+	err := retry.Do(s.ctx, registerFn, retry.Attempts(uint(s.sessionRetryTimes)))
 	if err != nil {
 		return nil, err
 	}

internal/util/sessionutil/session_util_test.go

@@ -704,9 +704,8 @@ func TestSessionEventType_String(t *testing.T) {
 
 func TestSession_apply(t *testing.T) {
 	session := &Session{}
-	opts := []SessionOption{WithCustomConfigEnable(), WithTTL(100), WithRetryTimes(200)}
+	opts := []SessionOption{WithTTL(100), WithRetryTimes(200)}
 	session.apply(opts...)
-	assert.True(t, session.useCustomConfig)
 	assert.Equal(t, int64(100), session.sessionTTL)
 	assert.Equal(t, int64(200), session.sessionRetryTimes)
 }