trace web functions

Author: wxiaoguang

modules/gtprof/trace.go

@@ -83,7 +83,7 @@ func (t *Tracer) Start(ctx context.Context, spanName string) (context.Context, *
 		starters = globalTraceStarters
 	}
 	ts := &TraceSpan{name: spanName, startTime: time.Now()}
-	existingCtxSpan, _ := ctx.Value(ContextKeySpan).(*TraceSpan)
+	existingCtxSpan := GetContextSpan(ctx)
 	if existingCtxSpan != nil {
 		existingCtxSpan.mu.Lock()
 		existingCtxSpan.children = append(existingCtxSpan.children, ts)

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

@@ -6,22 +6,39 @@ package routing
 import (
 	"context"
 	"net/http"
+
+	"code.gitea.io/gitea/modules/gtprof"
+	"code.gitea.io/gitea/modules/reqctx"
 )
 
 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
+func StartContextSpan(ctx reqctx.RequestContext, spanName string) (*gtprof.TraceSpan, func()) {
+	curTraceSpan := gtprof.GetContextSpan(ctx)
+	_, newTraceSpan := gtprof.GetTracer().Start(ctx, spanName)
+	ctx.SetContextValue(gtprof.ContextKeySpan, newTraceSpan)
+	return newTraceSpan, func() {
+		newTraceSpan.End()
+		ctx.SetContextValue(gtprof.ContextKeySpan, curTraceSpan)
 	}
+}
 
-	record.lock.Lock()
-	record.funcInfo = funcInfo
-	record.lock.Unlock()
+// RecordFuncInfo records a func info into context
+func RecordFuncInfo(ctx context.Context, funcInfo *FuncInfo) (end func()) {
+	end = func() {}
+	if reqCtx := reqctx.FromContext(ctx); reqCtx != nil {
+		var traceSpan *gtprof.TraceSpan
+		traceSpan, end = StartContextSpan(reqCtx, "http.func")
+		traceSpan.SetAttributeString("func", funcInfo.shortName)
+	}
+	if record, ok := ctx.Value(contextKey).(*requestRecord); ok {
+		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

@@ -1623,7 +1623,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
 }