enhance: [2.5] Add management api to check querycoord balance status (#…

…37784) (#39909)

issue: #37783
pr: #37784

Signed-off-by: Wei Liu <wei.liu@zilliz.com>

internal/distributed/querycoord/client/client.go

@@ -472,6 +472,17 @@ func (c *Client) ResumeBalance(ctx context.Context, req *querypb.ResumeBalanceRe
 	})
 }
 
+func (c *Client) CheckBalanceStatus(ctx context.Context, req *querypb.CheckBalanceStatusRequest, opts ...grpc.CallOption) (*querypb.CheckBalanceStatusResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.grpcClient.GetNodeID())),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryCoordClient) (*querypb.CheckBalanceStatusResponse, error) {
+		return client.CheckBalanceStatus(ctx, req)
+	})
+}
+
 func (c *Client) SuspendNode(ctx context.Context, req *querypb.SuspendNodeRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
 	req = typeutil.Clone(req)
 	commonpbutil.UpdateMsgBase(

internal/distributed/querycoord/client/client_test.go

@@ -185,6 +185,9 @@ func Test_NewClient(t *testing.T) {
 
 		r39, err := client.CheckQueryNodeDistribution(ctx, nil)
 		retCheck(retNotNil, r39, err)
+
+		r40, err := client.CheckBalanceStatus(ctx, nil)
+		retCheck(retNotNil, r40, err)
 	}
 
 	client.(*Client).grpcClient = &mock.GRPCClientBase[querypb.QueryCoordClient]{

internal/distributed/querycoord/service.go

@@ -457,6 +457,10 @@ func (s *Server) ResumeBalance(ctx context.Context, req *querypb.ResumeBalanceRe
 	return s.queryCoord.ResumeBalance(ctx, req)
 }
 
+func (s *Server) CheckBalanceStatus(ctx context.Context, req *querypb.CheckBalanceStatusRequest) (*querypb.CheckBalanceStatusResponse, error) {
+	return s.queryCoord.CheckBalanceStatus(ctx, req)
+}
+
 func (s *Server) SuspendNode(ctx context.Context, req *querypb.SuspendNodeRequest) (*commonpb.Status, error) {
 	return s.queryCoord.SuspendNode(ctx, req)
 }

internal/distributed/querycoord/service_test.go

@@ -311,6 +311,14 @@ func Test_NewServer(t *testing.T) {
 			assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
 		})
 
+		t.Run("CheckBalanceStatus", func(t *testing.T) {
+			req := &querypb.CheckBalanceStatusRequest{}
+			mqc.EXPECT().CheckBalanceStatus(mock.Anything, req).Return(&querypb.CheckBalanceStatusResponse{Status: merr.Success()}, nil)
+			resp, err := server.CheckBalanceStatus(ctx, req)
+			assert.NoError(t, err)
+			assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+		})
+
 		t.Run("SuspendNode", func(t *testing.T) {
 			req := &querypb.SuspendNodeRequest{}
 			mqc.EXPECT().SuspendNode(mock.Anything, req).Return(merr.Success(), nil)

internal/http/router.go

@@ -53,6 +53,7 @@ const (
 
 	RouteSuspendQueryCoordBalance = "/management/querycoord/balance/suspend"
 	RouteResumeQueryCoordBalance  = "/management/querycoord/balance/resume"
+	RouteQueryCoordBalanceStatus  = "/management/querycoord/balance/status"
 	RouteTransferSegment          = "/management/querycoord/transfer/segment"
 	RouteTransferChannel          = "/management/querycoord/transfer/channel"
 

internal/proxy/management.go

@@ -82,6 +82,10 @@ func RegisterMgrRoute(proxy *Proxy) {
 			Path:        management.RouteCheckQueryNodeDistribution,
 			HandlerFunc: proxy.CheckQueryNodeDistribution,
 		})
+		management.Register(&management.Handler{
+			Path:        management.RouteQueryCoordBalanceStatus,
+			HandlerFunc: proxy.CheckQueryCoordBalanceStatus,
+		})
 	})
 }
 
@@ -249,6 +253,29 @@ func (node *Proxy) ResumeQueryCoordBalance(w http.ResponseWriter, req *http.Requ
 	w.Write([]byte(`{"msg": "OK"}`))
 }
 
+func (node *Proxy) CheckQueryCoordBalanceStatus(w http.ResponseWriter, req *http.Request) {
+	resp, err := node.queryCoord.CheckBalanceStatus(req.Context(), &querypb.CheckBalanceStatusRequest{
+		Base: commonpbutil.NewMsgBase(),
+	})
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(fmt.Sprintf(`{"msg": "failed to check balance status, %s"}`, err.Error())))
+		return
+	}
+
+	if !merr.Ok(resp.GetStatus()) {
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte(fmt.Sprintf(`{"msg": "failed to check balance status, %s"}`, resp.GetStatus().GetReason())))
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+	balanceStatus := "suspended"
+	if resp.IsActive {
+		balanceStatus = "active"
+	}
+	w.Write([]byte(fmt.Sprintf(`{"msg": "OK", "status": "%v"}`, balanceStatus)))
+}
+
 func (node *Proxy) SuspendQueryNode(w http.ResponseWriter, req *http.Request) {
 	err := req.ParseForm()
 	if err != nil {

internal/proxy/management_test.go

@@ -371,6 +371,68 @@ func (s *ProxyManagementSuite) TestResumeQueryCoordBalance() {
 	})
 }
 
+func (s *ProxyManagementSuite) TestCheckBalanceStatus() {
+	s.Run("normal", func() {
+		s.SetupTest()
+		defer s.TearDownTest()
+
+		s.querycoord.EXPECT().CheckBalanceStatus(mock.Anything, mock.Anything).Return(&querypb.CheckBalanceStatusResponse{
+			Status:   merr.Success(),
+			IsActive: true,
+		}, nil).Times(1)
+
+		req, err := http.NewRequest(http.MethodPost, management.RouteQueryCoordBalanceStatus, nil)
+		s.Require().NoError(err)
+
+		recorder := httptest.NewRecorder()
+		s.proxy.CheckQueryCoordBalanceStatus(recorder, req)
+		s.Equal(http.StatusOK, recorder.Code)
+		s.Equal(`{"msg": "OK", "status": "active"}`, recorder.Body.String())
+
+		s.querycoord.EXPECT().CheckBalanceStatus(mock.Anything, mock.Anything).Return(&querypb.CheckBalanceStatusResponse{
+			Status:   merr.Success(),
+			IsActive: false,
+		}, nil).Times(1)
+
+		req, err = http.NewRequest(http.MethodPost, management.RouteQueryCoordBalanceStatus, nil)
+		s.Require().NoError(err)
+		recorder = httptest.NewRecorder()
+		s.proxy.CheckQueryCoordBalanceStatus(recorder, req)
+		s.Equal(http.StatusOK, recorder.Code)
+		s.Equal(`{"msg": "OK", "status": "suspended"}`, recorder.Body.String())
+	})
+
+	s.Run("return_error", func() {
+		s.SetupTest()
+		defer s.TearDownTest()
+
+		s.querycoord.EXPECT().CheckBalanceStatus(mock.Anything, mock.Anything).Return(nil, errors.New("mocked error"))
+
+		req, err := http.NewRequest(http.MethodPost, management.RouteQueryCoordBalanceStatus, nil)
+		s.Require().NoError(err)
+
+		recorder := httptest.NewRecorder()
+		s.proxy.CheckQueryCoordBalanceStatus(recorder, req)
+		s.Equal(http.StatusInternalServerError, recorder.Code)
+	})
+
+	s.Run("return_failure", func() {
+		s.SetupTest()
+		defer s.TearDownTest()
+
+		req, err := http.NewRequest(http.MethodPost, management.RouteQueryCoordBalanceStatus, nil)
+		s.Require().NoError(err)
+
+		s.querycoord.EXPECT().CheckBalanceStatus(mock.Anything, mock.Anything).Return(&querypb.CheckBalanceStatusResponse{
+			Status: merr.Status(merr.ErrServiceNotReady),
+		}, nil)
+
+		recorder := httptest.NewRecorder()
+		s.proxy.CheckQueryCoordBalanceStatus(recorder, req)
+		s.Equal(http.StatusInternalServerError, recorder.Code)
+	})
+}
+
 func (s *ProxyManagementSuite) TestSuspendQueryNode() {
 	s.Run("normal", func() {
 		s.SetupTest()

internal/querycoordv2/ops_service_test.go

@@ -402,10 +402,18 @@ func (suite *OpsServiceSuite) TestSuspendAndResumeBalance() {
 	suite.True(merr.Ok(resp))
 	suite.False(suite.checkerController.IsActive(utils.BalanceChecker))
 
+	resp1, err := suite.server.CheckBalanceStatus(ctx, &querypb.CheckBalanceStatusRequest{})
+	suite.NoError(err)
+	suite.Equal(false, resp1.GetIsActive())
+
 	resp, err = suite.server.ResumeBalance(ctx, &querypb.ResumeBalanceRequest{})
 	suite.NoError(err)
 	suite.True(merr.Ok(resp))
 	suite.True(suite.checkerController.IsActive(utils.BalanceChecker))
+
+	resp2, err := suite.server.CheckBalanceStatus(ctx, &querypb.CheckBalanceStatusRequest{})
+	suite.NoError(err)
+	suite.Equal(true, resp2.GetIsActive())
 }
 
 func (suite *OpsServiceSuite) TestSuspendAndResumeNode() {

internal/querycoordv2/ops_services.go

@@ -194,6 +194,33 @@ func (s *Server) ResumeBalance(ctx context.Context, req *querypb.ResumeBalanceRe
 	return merr.Success(), nil
 }
 
+// CheckBalanceStatus checks whether balance is active or suspended
+func (s *Server) CheckBalanceStatus(ctx context.Context, req *querypb.CheckBalanceStatusRequest) (*querypb.CheckBalanceStatusResponse, error) {
+	log := log.Ctx(ctx)
+	log.Info("CheckBalanceStatus request received")
+
+	errMsg := "failed to check balance status"
+	if err := merr.CheckHealthy(s.State()); err != nil {
+		log.Warn(errMsg, zap.Error(err))
+		return &querypb.CheckBalanceStatusResponse{
+			Status: merr.Status(err),
+		}, nil
+	}
+
+	isActive, err := s.checkerController.IsActive(utils.BalanceChecker)
+	if err != nil {
+		log.Warn(errMsg, zap.Error(err))
+		return &querypb.CheckBalanceStatusResponse{
+			Status: merr.Status(err),
+		}, nil
+	}
+
+	return &querypb.CheckBalanceStatusResponse{
+		Status:   merr.Success(),
+		IsActive: isActive,
+	}, nil
+}
+
 // suspend node from resource operation, for given node, suspend load_segment/sub_channel operations
 func (s *Server) SuspendNode(ctx context.Context, req *querypb.SuspendNodeRequest) (*commonpb.Status, error) {
 	log := log.Ctx(ctx)

internal/util/mock/grpc_querycoord_client.go

@@ -162,6 +162,10 @@ func (m *GrpcQueryCoordClient) ResumeBalance(ctx context.Context, req *querypb.R
 	return &commonpb.Status{}, m.Err
 }
 
+func (m *GrpcQueryCoordClient) CheckBalanceStatus(ctx context.Context, req *querypb.CheckBalanceStatusRequest, opts ...grpc.CallOption) (*querypb.CheckBalanceStatusResponse, error) {
+	return &querypb.CheckBalanceStatusResponse{}, m.Err
+}
+
 func (m *GrpcQueryCoordClient) SuspendNode(ctx context.Context, req *querypb.SuspendNodeRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
 	return &commonpb.Status{}, m.Err
 }