From 0c74539d3a16c0471803e76badad6d34a5de7fbd Mon Sep 17 00:00:00 2001
From: kqzh <35095889+kqzh@users.noreply.github.com>
Date: Sun, 8 Oct 2023 11:35:52 +0800
Subject: [PATCH] feat: add close

---
 pkg/client/client.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pkg/client/client.go b/pkg/client/client.go
index fc1a3a6..c7ffcf9 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -34,6 +34,7 @@ type Client interface {
 	StopAgent(req *pb.StopAgentRequest) (*pb.StopAgentResponse, error)
 	HealthCheck(req *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error)
 	GetSpaceUsages(req *pb.GetSpaceUsagesRequest) (*pb.GetSpaceUsagesResponse, error)
+	Close() error
 }
 
 func genSessionId() string {
@@ -43,6 +44,7 @@ func genSessionId() string {
 type client struct {
 	ctx     context.Context
 	addr    *nebula.HostAddr
+	conn    *grpc.ClientConn
 	storage pb.StorageServiceClient
 	agent   pb.AgentServiceClient
 }
@@ -62,6 +64,7 @@ func New(ctx context.Context, cfg *Config) (Client, error) {
 	c := &client{
 		ctx:     ctx,
 		addr:    cfg.Addr,
+		conn:    conn,
 		storage: pb.NewStorageServiceClient(conn),
 		agent:   pb.NewAgentServiceClient(conn),
 	}
@@ -69,6 +72,10 @@ func New(ctx context.Context, cfg *Config) (Client, error) {
 	return c, nil
 }
 
+func (c *client) Close() error {
+	return c.conn.Close()
+}
+
 func (c *client) GetAddr() *nebula.HostAddr {
 	return c.addr
 }