From f2c0ead51af67b2dd52affac4f35d1ed92d15a87 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Thu, 25 Jan 2024 16:59:00 +0800
Subject: [PATCH] enhance: Add trace span for scheduling read tasks in
 QueryNode (#30265)

This PR adds a trace span for search/query task scheduling duration

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/querynodev2/tasks/query_task.go | 10 ++++++++++
 internal/querynodev2/tasks/task.go       | 11 ++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/internal/querynodev2/tasks/query_task.go b/internal/querynodev2/tasks/query_task.go
index 0460be5223..b495c2b1a8 100644
--- a/internal/querynodev2/tasks/query_task.go
+++ b/internal/querynodev2/tasks/query_task.go
@@ -6,6 +6,9 @@ import (
 	"strconv"
 	"time"
 
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
@@ -16,6 +19,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/metricsinfo"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/timerecord"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 var _ Task = &QueryTask{}
@@ -25,6 +29,7 @@ func NewQueryTask(ctx context.Context,
 	manager *segments.Manager,
 	req *querypb.QueryRequest,
 ) *QueryTask {
+	ctx, span := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, "schedule")
 	return &QueryTask{
 		ctx:            ctx,
 		collection:     collection,
@@ -32,6 +37,7 @@ func NewQueryTask(ctx context.Context,
 		req:            req,
 		notifier:       make(chan error, 1),
 		tr:             timerecord.NewTimeRecorderWithTrace(ctx, "queryTask"),
+		scheduleSpan:   span,
 	}
 }
 
@@ -43,6 +49,7 @@ type QueryTask struct {
 	result         *internalpb.RetrieveResults
 	notifier       chan error
 	tr             *timerecord.TimeRecorder
+	scheduleSpan   trace.Span
 }
 
 // Return the username which task is belong to.
@@ -81,6 +88,9 @@ func (t *QueryTask) PreExecute() error {
 
 // Execute the task, only call once.
 func (t *QueryTask) Execute() error {
+	if t.scheduleSpan != nil {
+		t.scheduleSpan.End()
+	}
 	tr := timerecord.NewTimeRecorderWithTrace(t.ctx, "QueryTask")
 
 	retrievePlan, err := segments.NewRetrievePlan(
diff --git a/internal/querynodev2/tasks/task.go b/internal/querynodev2/tasks/task.go
index 83498157a2..774678dd39 100644
--- a/internal/querynodev2/tasks/task.go
+++ b/internal/querynodev2/tasks/task.go
@@ -9,6 +9,8 @@ import (
 	"strconv"
 
 	"github.com/golang/protobuf/proto"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@@ -23,6 +25,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/metricsinfo"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/timerecord"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 var (
@@ -46,7 +49,8 @@ type SearchTask struct {
 	others           []*SearchTask
 	notifier         chan error
 
-	tr *timerecord.TimeRecorder
+	tr           *timerecord.TimeRecorder
+	scheduleSpan trace.Span
 }
 
 func NewSearchTask(ctx context.Context,
@@ -54,6 +58,7 @@ func NewSearchTask(ctx context.Context,
 	manager *segments.Manager,
 	req *querypb.SearchRequest,
 ) *SearchTask {
+	ctx, span := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, "schedule")
 	return &SearchTask{
 		ctx:              ctx,
 		collection:       collection,
@@ -68,6 +73,7 @@ func NewSearchTask(ctx context.Context,
 		originNqs:        []int64{req.GetReq().GetNq()},
 		notifier:         make(chan error, 1),
 		tr:               timerecord.NewTimeRecorderWithTrace(ctx, "searchTask"),
+		scheduleSpan:     span,
 	}
 }
 
@@ -118,6 +124,9 @@ func (t *SearchTask) Execute() error {
 		zap.String("shard", t.req.GetDmlChannels()[0]),
 	)
 
+	if t.scheduleSpan != nil {
+		t.scheduleSpan.End()
+	}
 	tr := timerecord.NewTimeRecorderWithTrace(t.ctx, "SearchTask")
 
 	req := t.req