From ba11c4e56ac1fe849b37c11ca29bb23f7bf5ce57 Mon Sep 17 00:00:00 2001
From: Congqi Xia <congqi.xia@zilliz.com>
Date: Mon, 12 Aug 2024 11:43:08 +0800
Subject: [PATCH] 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 <congqi.xia@zilliz.com>
---
 models/session.go             | 12 +++++++-
 states/etcd/common/session.go |  1 +
 states/etcd/remove/session.go | 54 +++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 states/etcd/remove/session.go

diff --git a/models/session.go b/models/session.go
index 123442b3..f1b7ebe6 100644
--- a/models/session.go
+++ b/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 ""
diff --git a/states/etcd/common/session.go b/states/etcd/common/session.go
index e4c4374c..0b3b6407 100644
--- a/states/etcd/common/session.go
+++ b/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)
 	}
diff --git a/states/etcd/remove/session.go b/states/etcd/remove/session.go
new file mode 100644
index 00000000..a0846a64
--- /dev/null
+++ b/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
+}