enhance: Use unified requery flag when set metrics (milvus-io#34632)

See also milvus-io#34601 Signed-off-by: Congqi Xia <[email protected]>
bigsheeper · Jul 15, 2024 · b5225df · b5225df
1 parent eb4bfa3
commit b5225df
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 45 deletions.
diff --git 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
@@ -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
@@ -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)
-	}
-}