diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index 447629b06d1c6..491713cb2eaf6 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -3421,7 +3421,6 @@ func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (* func (node *Proxy) query(ctx context.Context, qt *queryTask) (*milvuspb.QueryResults, error) { request := qt.request method := "Query" - isProxyRequest := GetRequestLabelFromContext(ctx) if err := merr.CheckHealthy(node.GetStateCode()); err != nil { return &milvuspb.QueryResults{ @@ -3470,7 +3469,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask) (*milvuspb.QueryRes zap.Error(err), ) - if isProxyRequest { + if !qt.reQuery { metrics.ProxyFunctionCall.WithLabelValues( strconv.FormatInt(paramtable.GetNodeID(), 10), method, @@ -3493,7 +3492,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask) (*milvuspb.QueryRes rpcFailedToWaitToFinish(method), zap.Error(err)) - if isProxyRequest { + if !qt.reQuery { metrics.ProxyFunctionCall.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method, metrics.FailLabel, request.GetDbName(), request.GetCollectionName()).Inc() } @@ -3503,7 +3502,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask) (*milvuspb.QueryRes }, nil } - if isProxyRequest { + if !qt.reQuery { span := tr.CtxRecord(ctx, "wait query result") metrics.ProxyWaitForSearchResultLatency.WithLabelValues( strconv.FormatInt(paramtable.GetNodeID(), 10), @@ -3578,7 +3577,6 @@ func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (* request.GetCollectionName(), ).Inc() - ctx = SetRequestLabelForContext(ctx) res, err := node.query(ctx, qt) if err != nil || !merr.Ok(res.Status) { return res, err diff --git a/internal/proxy/util.go b/internal/proxy/util.go index 6733dd7d41c93..837304c45009e 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -1635,22 +1635,3 @@ func GetCostValue(status *commonpb.Status) int { } return value } - -type isProxyRequestKeyType struct{} - -var ctxProxyRequestKey = isProxyRequestKeyType{} - -func SetRequestLabelForContext(ctx context.Context) context.Context { - return context.WithValue(ctx, ctxProxyRequestKey, true) -} - -func GetRequestLabelFromContext(ctx context.Context) bool { - if ctx == nil { - return false - } - v := ctx.Value(ctxProxyRequestKey) - if v == nil { - return false - } - return v.(bool) -} diff --git a/internal/proxy/util_test.go b/internal/proxy/util_test.go index 778ddb095afc5..a3a3532e4e724 100644 --- a/internal/proxy/util_test.go +++ b/internal/proxy/util_test.go @@ -2472,24 +2472,3 @@ func TestGetCostValue(t *testing.T) { assert.Equal(t, 100, cost) }) } - -func TestRequestLabelWithContext(t *testing.T) { - ctx := context.Background() - - { - label := GetRequestLabelFromContext(ctx) - assert.False(t, label) - } - - ctx = SetRequestLabelForContext(ctx) - { - label := GetRequestLabelFromContext(ctx) - assert.True(t, label) - } - - { - // nolint - label := GetRequestLabelFromContext(nil) - assert.False(t, label) - } -}