impl check health (#525)

Signed-off-by: huanghaoyuan <haoyuan.huang@zilliz.com>

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
+}

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})
+	})
+
+}

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.

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
+}