diff --git a/internal/distributed/proxy/httpserver/handler_v2.go b/internal/distributed/proxy/httpserver/handler_v2.go index fced1d4877ed1..88a533637e984 100644 --- a/internal/distributed/proxy/httpserver/handler_v2.go +++ b/internal/distributed/proxy/httpserver/handler_v2.go @@ -192,7 +192,7 @@ func wrapperPost(newReq newReqFunc, v2 handlerFuncV2) gin.HandlerFunc { } } username, _ := c.Get(ContextUsername) - ctx, span := otel.Tracer(typeutil.ProxyRole).Start(c, c.Request.URL.Path) + ctx, span := otel.Tracer(typeutil.ProxyRole).Start(getCtx(c), c.Request.URL.Path) defer span.End() ctx = proxy.NewContextWithMetadata(ctx, username.(string), dbName) traceID := span.SpanContext().TraceID().String() @@ -204,10 +204,23 @@ func wrapperPost(newReq newReqFunc, v2 handlerFuncV2) gin.HandlerFunc { } } +const ( + v2CtxKey = `milvus_restful_v2_ctxkey` +) + +func getCtx(ctx *gin.Context) context.Context { + v, ok := ctx.Get(v2CtxKey) + if !ok { + return ctx + } + return v.(context.Context) +} + // restfulSizeMiddleware is the middleware fetchs metrics stats from gin struct. func restfulSizeMiddleware(handler gin.HandlerFunc, observeOutbound bool) gin.HandlerFunc { return func(ctx *gin.Context) { h := metrics.WrapRestfulContext(ctx, ctx.Request.ContentLength) + ctx.Set(v2CtxKey, h) handler(ctx) metrics.RecordRestfulMetrics(h, int64(ctx.Writer.Size()), observeOutbound) }