test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-01 19:39:21 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

enhance: Pass rpc stats via gin.Context (#37439)

Browse Source 
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 18:26:23 +08:00 committed by GitHub
parent bd04cac4b3
commit 9a9de3df5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
internal/distributed/proxy/httpserver/handler_v2.go
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)
}