enhance: [2.4] Pass rpc stats via gin.Context (#37440)

Cherry pick from master
pr: #37439
Related #37223

RPC stats worked in middleware but faild to get method & collection info

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-11-05 14:24:24 +08:00 committed by GitHub
parent 6b69170a64
commit c195f9f76a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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)
}