From 869ef9fb9ba6ff0c45a9f244fbc2028650a055ed Mon Sep 17 00:00:00 2001
From: huanghaoyuan <haoyuan.huang@zilliz.com>
Date: Wed, 26 Jul 2023 10:28:57 +0800
Subject: [PATCH] impl check health

Signed-off-by: huanghaoyuan <haoyuan.huang@zilliz.com>
---
 client/admin.go        | 25 +++++++++++++++++-
 client/admin_test.go   | 57 ++++++++++++++++++++++++++++++++++++++++++
 client/client.go       |  2 ++
 entity/check_health.go | 22 ++++++++++++++++
 4 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 client/admin_test.go
 create mode 100644 entity/check_health.go

diff --git a/client/admin.go b/client/admin.go
index 69ebd9f6..3246d65d 100644
--- a/client/admin.go
+++ b/client/admin.go
@@ -15,9 +15,10 @@ import (
 	"context"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
-// GetVersion returns milvus milvuspb.version information.
+// GetVersion returns milvus server version information.
 func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
 	if c.Service == nil {
 		return "", ErrClientNotReady
@@ -28,3 +29,25 @@ func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
 	}
 	return resp.GetVersion(), nil
 }
+
+// CheckHealth returns milvus state
+func (c *GrpcClient) CheckHealth(ctx context.Context) (*entity.MilvusState, error) {
+	if c.Service == nil {
+		return nil, ErrClientNotReady
+	}
+	resp, err := c.Service.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
+	if err != nil {
+		return nil, err
+	}
+
+	states := make([]entity.QuotaState, 0, len(resp.GetQuotaStates()))
+	for _, state := range resp.GetQuotaStates() {
+		states = append(states, entity.QuotaState(state))
+	}
+
+	return &entity.MilvusState{
+		IsHealthy:   resp.GetIsHealthy(),
+		Reasons:     resp.GetReasons(),
+		QuotaStates: states,
+	}, nil
+}
diff --git a/client/admin_test.go b/client/admin_test.go
new file mode 100644
index 00000000..ca70d1bc
--- /dev/null
+++ b/client/admin_test.go
@@ -0,0 +1,57 @@
+package client
+
+import (
+	"context"
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+	common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+	server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCheckHealth(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(ctx, t)
+	defer c.Close()
+
+	t.Run("test milvus healthy", func(t *testing.T) {
+		mockServer.SetInjection(MCheckHealth, func(ctx context.Context, message proto.Message) (proto.Message, error) {
+			resp := &server.CheckHealthResponse{
+				Status:      &common.Status{ErrorCode: common.ErrorCode_Success},
+				IsHealthy:   true,
+				Reasons:     nil,
+				QuotaStates: nil,
+			}
+			return resp, nil
+		})
+		defer mockServer.DelInjection(MCheckHealth)
+
+		resp, err := c.CheckHealth(ctx)
+		assert.Nil(t, err)
+		assert.True(t, resp.IsHealthy)
+		assert.Empty(t, resp.Reasons)
+		assert.Empty(t, resp.QuotaStates)
+	})
+
+	t.Run("test milvus unhealthy", func(t *testing.T) {
+		mockServer.SetInjection(MCheckHealth, func(ctx context.Context, message proto.Message) (proto.Message, error) {
+			resp := &server.CheckHealthResponse{
+				Status:      &common.Status{ErrorCode: common.ErrorCode_Success},
+				IsHealthy:   false,
+				Reasons:     []string{"some reason"},
+				QuotaStates: []server.QuotaState{server.QuotaState_DenyToRead, server.QuotaState_DenyToWrite},
+			}
+			return resp, nil
+		})
+		defer mockServer.DelInjection(MCheckHealth)
+
+		resp, err := c.CheckHealth(ctx)
+		assert.Nil(t, err)
+		assert.False(t, resp.IsHealthy)
+		assert.ElementsMatch(t, resp.Reasons, []string{"some reason"})
+		assert.ElementsMatch(t, resp.QuotaStates, []entity.QuotaState{entity.QuotaStateDenyToRead, entity.QuotaStateDenyToWrite})
+	})
+
+}
diff --git a/client/client.go b/client/client.go
index b974b53d..5f066b02 100644
--- a/client/client.go
+++ b/client/client.go
@@ -209,6 +209,8 @@ type Client interface {
 
 	// GetVersion get milvus version
 	GetVersion(ctx context.Context) (string, error)
+	// CheckHealth returns milvus state
+	CheckHealth(ctx context.Context) (*entity.MilvusState, error)
 }
 
 // NewClient create a client connected to remote milvus cluster.
diff --git a/entity/check_health.go b/entity/check_health.go
new file mode 100644
index 00000000..c821d8c9
--- /dev/null
+++ b/entity/check_health.go
@@ -0,0 +1,22 @@
+package entity
+
+type QuotaState int32
+
+const (
+	// QuotaStateUnknown zero value placeholder
+	QuotaStateUnknown QuotaState = 0
+	// QuotaStateReadLimited too many read tasks, read requests are limited
+	QuotaStateReadLimited QuotaState = 2
+	// QuotaStateWriteLimited too many write tasks, write requests are limited
+	QuotaStateWriteLimited QuotaState = 3
+	// QuotaStateDenyToRead too many read tasks, temporarily unable to process read requests
+	QuotaStateDenyToRead QuotaState = 4
+	// QuotaStateDenyToWrite too many write tasks, temporarily unable to process write requests
+	QuotaStateDenyToWrite QuotaState = 5
+)
+
+type MilvusState struct {
+	IsHealthy   bool
+	Reasons     []string
+	QuotaStates []QuotaState
+}