enhance: Add remove session command

This PR add `remove session --component [service] --sessionID [id]`
command to clean or force kill session.

Signed-off-by: Congqi Xia <[email protected]>

models/session.go

@@ -12,13 +12,23 @@ type Session struct {
 	Address    string `json:"Address,omitempty"`
 	Exclusive  bool   `json:"Exclusive,omitempty"`
 	Version    string `json:"Version,omitempty"`
+
+	key string
+}
+
+func (s *Session) SetKey(key string) {
+	s.key = key
+}
+
+func (s *Session) GetKey() string {
+	return s.key
 }
 
 func (s Session) String() string {
 	return fmt.Sprintf("Session:%s, ServerID: %d, Version: %s, Address: %s", s.ServerName, s.ServerID, s.Version, s.Address)
 }
 
-func (s Session) IP() string {
+func (s *Session) IP() string {
 	addr, err := net.ResolveTCPAddr("tcp", s.Address)
 	if err != nil {
 		return ""

states/etcd/common/session.go

@@ -37,6 +37,7 @@ func ListSessionsByPrefix(cli clientv3.KV, prefix string) ([]*models.Session, er
 		if err != nil {
 			continue
 		}
+		session.SetKey(string(kv.Key))
 
 		sessions = append(sessions, session)
 	}

states/etcd/remove/session.go

@@ -0,0 +1,54 @@
+package remove
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/milvus-io/birdwatcher/framework"
+	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	"github.com/samber/lo"
+)
+
+type RemoveSessionParam struct {
+	framework.ParamBase `use:"remove session" desc:"remove session with specified component type & node id"`
+	Component           string `name:"component" default:"" desc:"component type to remove"`
+	ID                  int64  `name:"sessionID" default:"0" desc:"session id to remove"`
+	Run                 bool   `name:"run" default:"false" desc:"actual remove session, default in dry-run mode"`
+}
+
+func (c *ComponentRemove) RemoveSessionCommand(ctx context.Context, p *RemoveSessionParam) error {
+	sessions, err := common.ListSessions(c.client, c.basePath)
+	if err != nil {
+		return err
+	}
+
+	sessions = lo.Filter(sessions, func(s *models.Session, _ int) bool {
+		return strings.EqualFold(s.ServerName, p.Component) && s.ServerID == p.ID
+	})
+
+	if len(sessions) == 0 {
+		fmt.Printf("Session component type=%s session id=%d not found", p.Component, p.ID)
+	}
+
+	fmt.Printf("%d session item found:\n", len(sessions))
+	for _, session := range sessions {
+		fmt.Println(session.String())
+		fmt.Println(session.GetKey())
+	}
+
+	if p.Run {
+		fmt.Println("Start to remove session")
+		for _, session := range sessions {
+			_, err := c.client.Delete(ctx, session.GetKey())
+			if err != nil {
+				return err
+			}
+		}
+	} else {
+		fmt.Println("[Dry-mode] skip removing")
+	}
+
+	return nil
+}