Prepare for support performance trace (#33286)

For #32973

Author: wxiaoguang

modules/web/handler.go

@@ -121,7 +121,7 @@ func wrapHandlerProvider[T http.Handler](hp func(next http.Handler) T, funcInfo
 	return func(next http.Handler) http.Handler {
 		h := hp(next) // this handle could be dynamically generated, so we can't use it for debug info
 		return http.HandlerFunc(func(resp http.ResponseWriter, req *http.Request) {
-			routing.UpdateFuncInfo(req.Context(), funcInfo)
+			defer routing.RecordFuncInfo(req.Context(), funcInfo)()
 			h.ServeHTTP(resp, req)
 		})
 	}
@@ -157,7 +157,7 @@ func toHandlerProvider(handler any) func(next http.Handler) http.Handler {
 				return // it's doing pre-check, just return
 			}
 
-			routing.UpdateFuncInfo(req.Context(), funcInfo)
+			defer routing.RecordFuncInfo(req.Context(), funcInfo)()
 			ret := fn.Call(argsIn)
 
 			// handle the return value (no-op at the moment)

modules/web/routing/context.go

@@ -12,16 +12,18 @@ type contextKeyType struct{}
 
 var contextKey contextKeyType
 
-// UpdateFuncInfo updates a context's func info
-func UpdateFuncInfo(ctx context.Context, funcInfo *FuncInfo) {
-	record, ok := ctx.Value(contextKey).(*requestRecord)
-	if !ok {
-		return
+// RecordFuncInfo records a func info into context
+func RecordFuncInfo(ctx context.Context, funcInfo *FuncInfo) (end func()) {
+	// TODO: reqCtx := reqctx.FromContext(ctx), add trace support
+	end = func() {}
+
+	// save the func info into the context record
+	if record, ok := ctx.Value(contextKey).(*requestRecord); ok {
+		record.lock.Lock()
+		record.funcInfo = funcInfo
+		record.lock.Unlock()
 	}
-
-	record.lock.Lock()
-	record.funcInfo = funcInfo
-	record.lock.Unlock()
+	return end
 }
 
 // MarkLongPolling marks the request is a long-polling request, and the logger may output different message for it

routers/init.go

@@ -213,7 +213,7 @@ func NormalRoutes() *web.Router {
 	}
 
 	r.NotFound(func(w http.ResponseWriter, req *http.Request) {
-		routing.UpdateFuncInfo(req.Context(), routing.GetFuncInfo(http.NotFound, "GlobalNotFound"))
+		defer routing.RecordFuncInfo(req.Context(), routing.GetFuncInfo(http.NotFound, "GlobalNotFound"))()
 		http.NotFound(w, req)
 	})
 	return r

routers/web/base.go

@@ -34,7 +34,7 @@ func storageHandler(storageSetting *setting.Storage, prefix string, objStore sto
 				http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
 				return
 			}
-			routing.UpdateFuncInfo(req.Context(), funcInfo)
+			defer routing.RecordFuncInfo(req.Context(), funcInfo)()
 
 			rPath := strings.TrimPrefix(req.URL.Path, "/"+prefix+"/")
 			rPath = util.PathJoinRelX(rPath)
@@ -65,7 +65,7 @@ func storageHandler(storageSetting *setting.Storage, prefix string, objStore sto
 			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
 			return
 		}
-		routing.UpdateFuncInfo(req.Context(), funcInfo)
+		defer routing.RecordFuncInfo(req.Context(), funcInfo)()
 
 		rPath := strings.TrimPrefix(req.URL.Path, "/"+prefix+"/")
 		rPath = util.PathJoinRelX(rPath)

routers/web/web.go

@@ -1622,7 +1622,7 @@ func registerRoutes(m *web.Router) {
 
 	m.NotFound(func(w http.ResponseWriter, req *http.Request) {
 		ctx := context.GetWebContext(req)
-		routing.UpdateFuncInfo(ctx, routing.GetFuncInfo(ctx.NotFound, "WebNotFound"))
+		defer routing.RecordFuncInfo(ctx, routing.GetFuncInfo(ctx.NotFound, "WebNotFound"))()
 		ctx.NotFound("", nil)
 	})
 }

services/context/base.go

@@ -4,7 +4,6 @@
 package context
 
 import (
-	"context"
 	"fmt"
 	"html/template"
 	"io"
@@ -25,8 +24,7 @@ type BaseContextKeyType struct{}
 var BaseContextKey BaseContextKeyType
 
 type Base struct {
-	context.Context
-	reqctx.RequestDataStore
+	reqctx.RequestContext
 
 	Resp ResponseWriter
 	Req  *http.Request
@@ -172,19 +170,19 @@ func (b *Base) TrN(cnt any, key1, keyN string, args ...any) template.HTML {
 }
 
 func NewBaseContext(resp http.ResponseWriter, req *http.Request) *Base {
-	ds := reqctx.GetRequestDataStore(req.Context())
+	reqCtx := reqctx.FromContext(req.Context())
 	b := &Base{
-		Context:          req.Context(),
-		RequestDataStore: ds,
-		Req:              req,
-		Resp:             WrapResponseWriter(resp),
-		Locale:           middleware.Locale(resp, req),
-		Data:             ds.GetData(),
+		RequestContext: reqCtx,
+
+		Req:    req,
+		Resp:   WrapResponseWriter(resp),
+		Locale: middleware.Locale(resp, req),
+		Data:   reqCtx.GetData(),
 	}
 	b.Req = b.Req.WithContext(b)
-	ds.SetContextValue(BaseContextKey, b)
-	ds.SetContextValue(translation.ContextKey, b.Locale)
-	ds.SetContextValue(httplib.RequestContextKey, b.Req)
+	reqCtx.SetContextValue(BaseContextKey, b)
+	reqCtx.SetContextValue(translation.ContextKey, b.Locale)
+	reqCtx.SetContextValue(httplib.RequestContextKey, b.Req)
 	return b
 }