diff --git a/internal/querynode/impl.go b/internal/querynode/impl.go index bb553f48ea..d897c73080 100644 --- a/internal/querynode/impl.go +++ b/internal/querynode/impl.go @@ -21,6 +21,7 @@ import ( "go.uber.org/zap" + "github.com/milvus-io/milvus/internal/common" "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/proto/commonpb" "github.com/milvus-io/milvus/internal/proto/internalpb" @@ -38,17 +39,21 @@ func (node *QueryNode) GetComponentStates(ctx context.Context) (*internalpb.Comp ErrorCode: commonpb.ErrorCode_Success, }, } - code := node.stateCode.Load().(internalpb.StateCode) - if code != internalpb.StateCode_Healthy { - err := fmt.Errorf("query node %d is not ready", Params.QueryNodeID) + code, ok := node.stateCode.Load().(internalpb.StateCode) + if !ok { + errMsg := "unexpected error in type assertion" stats.Status = &commonpb.Status{ ErrorCode: commonpb.ErrorCode_UnexpectedError, - Reason: err.Error(), + Reason: errMsg, } - return stats, err + return stats, errors.New(errMsg) + } + nodeID := common.NotRegisteredID + if node.session != nil && node.session.Registered() { + nodeID = node.session.ServerID } info := &internalpb.ComponentInfo{ - NodeID: Params.QueryNodeID, + NodeID: nodeID, Role: typeutil.QueryNodeRole, StateCode: code, } diff --git a/internal/querynode/impl_test.go b/internal/querynode/impl_test.go index ac58f80534..b706f1dae5 100644 --- a/internal/querynode/impl_test.go +++ b/internal/querynode/impl_test.go @@ -15,6 +15,7 @@ import ( "context" "encoding/json" "math/rand" + "sync/atomic" "testing" "github.com/stretchr/testify/assert" @@ -33,12 +34,20 @@ func TestImpl_GetComponentStates(t *testing.T) { node, err := genSimpleQueryNode(ctx) assert.NoError(t, err) + node.session.UpdateRegistered(true) + rsp, err := node.GetComponentStates(ctx) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, rsp.Status.ErrorCode) node.UpdateStateCode(internalpb.StateCode_Abnormal) rsp, err = node.GetComponentStates(ctx) + assert.NoError(t, err) + assert.Equal(t, commonpb.ErrorCode_Success, rsp.Status.ErrorCode) + + node.stateCode = atomic.Value{} + node.stateCode.Store("invalid") + rsp, err = node.GetComponentStates(ctx) assert.Error(t, err) assert.Equal(t, commonpb.ErrorCode_UnexpectedError, rsp.Status.ErrorCode) }