enhance: Pass rpc stats via gin.Context

Related milvus-io#37223 RPC stats worked in middleware but faild to get method & collection info Signed-off-by: Congqi Xia <[email protected]>
congqixia · Nov 5, 2024 · 294258f · 294258f
1 parent 6b69170
commit 294258f
Showing 1 changed file with 14 additions and 1 deletion.
diff --git 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)
 	}