From 8520ee7552b079f90c7be4bf240a9dea6a03f3d9 Mon Sep 17 00:00:00 2001
From: "yihao.dai" <yihao.dai@zilliz.com>
Date: Mon, 27 Nov 2023 10:18:26 +0800
Subject: [PATCH] enhance: Print nq (#28507) (#28639)

Log nq in search path.

pr: https://github.com/milvus-io/milvus/pull/28507

---------

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 internal/proxy/impl.go           | 3 ++-
 internal/proxy/task_search.go    | 8 +++++---
 internal/querynodev2/handlers.go | 1 +
 internal/querynodev2/services.go | 1 +
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go
index 21f7f18ee5fa5..036386a7dc740 100644
--- a/internal/proxy/impl.go
+++ b/internal/proxy/impl.go
@@ -2564,7 +2564,7 @@ func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest)
 	defer func() {
 		span := tr.ElapseSpan()
 		if span >= SlowReadSpan {
-			log.Info(rpcSlow(method), zap.Duration("duration", span))
+			log.Info(rpcSlow(method), zap.Int64("nq", qt.SearchRequest.GetNq()), zap.Duration("duration", span))
 		}
 	}()
 
@@ -2596,6 +2596,7 @@ func (node *Proxy) Search(ctx context.Context, request *milvuspb.SearchRequest)
 	if err := qt.WaitToFinish(); err != nil {
 		log.Warn(
 			rpcFailedToWaitToFinish(method),
+			zap.Int64("nq", qt.SearchRequest.GetNq()),
 			zap.Error(err),
 		)
 
diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go
index 19c5498017f9a..5aaf7ff644cd8 100644
--- a/internal/proxy/task_search.go
+++ b/internal/proxy/task_search.go
@@ -210,7 +210,6 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 
 	t.Base.MsgType = commonpb.MsgType_Search
 	t.Base.SourceID = paramtable.GetNodeID()
-	log := log.Ctx(ctx)
 
 	collectionName := t.request.CollectionName
 	t.collectionName = collectionName
@@ -219,6 +218,8 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 		return err
 	}
 
+	log := log.Ctx(ctx).With(zap.Int64("collID", collID), zap.String("collName", collectionName))
+
 	t.SearchRequest.DbID = 0 // todo
 	t.SearchRequest.CollectionID = collID
 	t.schema, err = globalMetaCache.GetCollectionSchema(ctx, t.request.GetDbName(), collectionName)
@@ -270,6 +271,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 		return fmt.Errorf("%s [%d] is invalid, %w", NQKey, nq, err)
 	}
 	t.SearchRequest.Nq = nq
+	log = log.With(zap.Int64("nq", nq))
 
 	outputFieldIDs, err := getOutputFieldIDs(t.schema, t.request.GetOutputFields())
 	if err != nil {
@@ -404,7 +406,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 func (t *searchTask) Execute(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-Execute")
 	defer sp.End()
-	log := log.Ctx(ctx)
+	log := log.Ctx(ctx).With(zap.Int64("nq", t.SearchRequest.GetNq()))
 
 	tr := timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute search %d", t.ID()))
 	defer tr.CtxElapse(ctx, "done")
@@ -437,7 +439,7 @@ func (t *searchTask) PostExecute(ctx context.Context) error {
 	defer func() {
 		tr.CtxElapse(ctx, "done")
 	}()
-	log := log.Ctx(ctx)
+	log := log.Ctx(ctx).With(zap.Int64("nq", t.SearchRequest.GetNq()))
 
 	var (
 		Nq         = t.SearchRequest.GetNq()
diff --git a/internal/querynodev2/handlers.go b/internal/querynodev2/handlers.go
index 55c731310c7ae..8f734acea88a6 100644
--- a/internal/querynodev2/handlers.go
+++ b/internal/querynodev2/handlers.go
@@ -364,6 +364,7 @@ func (node *QueryNode) searchChannel(ctx context.Context, req *querypb.SearchReq
 		zap.Int64("collectionID", req.Req.GetCollectionID()),
 		zap.String("channel", channel),
 		zap.String("scope", req.GetScope().String()),
+		zap.Int64("nq", req.GetReq().GetNq()),
 	)
 	traceID := trace.SpanFromContext(ctx).SpanContext().TraceID()
 
diff --git a/internal/querynodev2/services.go b/internal/querynodev2/services.go
index b3c324aa300bc..ef5d4838131c8 100644
--- a/internal/querynodev2/services.go
+++ b/internal/querynodev2/services.go
@@ -730,6 +730,7 @@ func (node *QueryNode) Search(ctx context.Context, req *querypb.SearchRequest) (
 		zap.Int64("collectionID", req.GetReq().GetCollectionID()),
 		zap.Strings("channels", req.GetDmlChannels()),
 		zap.Bool("fromShardLeader", req.GetFromShardLeader()),
+		zap.Int64("nq", req.GetReq().GetNq()),
 	)
 
 	log.Debug("Received SearchRequest",