From 9a9de3df5cf6f306e1876c2b19a597244222afb4 Mon Sep 17 00:00:00 2001 From: congqixia Date: Tue, 5 Nov 2024 18:26:23 +0800 Subject: [PATCH] enhance: Pass rpc stats via gin.Context (#37439) Related #37223 RPC stats worked in middleware but faild to get method & collection info Signed-off-by: Congqi Xia --- .../distributed/proxy/httpserver/handler_v2.go | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) 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) }