Add basic runtime tracing (#2996)

This allows us in almost all places to use regions to further trace down
long running tasks.
Also removes an unused function.

4 files changed, 98 insertions, 63 deletions

diff --git a/internal/httputil/httpapi.go b/internal/httputil/httpapi.go
index 37d144f4..289d1d2c 100644
--- a/internal/httputil/httpapi.go
+++ b/internal/httputil/httpapi.go
@@ -25,8 +25,6 @@ import (
 
 	"github.com/getsentry/sentry-go"
 	"github.com/matrix-org/util"
-	"github.com/opentracing/opentracing-go"
-	"github.com/opentracing/opentracing-go/ext"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -34,6 +32,7 @@ import (
 
 	"github.com/matrix-org/dendrite/clientapi/auth"
 	"github.com/matrix-org/dendrite/clientapi/jsonerror"
+	"github.com/matrix-org/dendrite/internal"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
 )
 
@@ -186,9 +185,9 @@ func MakeExternalAPI(metricsName string, f func(*http.Request) util.JSONResponse
 			}
 		}
 
-		span := opentracing.StartSpan(metricsName)
-		defer span.Finish()
-		req = req.WithContext(opentracing.ContextWithSpan(req.Context(), span))
+		trace, ctx := internal.StartTask(req.Context(), metricsName)
+		defer trace.EndTask()
+		req = req.WithContext(ctx)
 		h.ServeHTTP(nextWriter, req)
 
 	}
@@ -200,9 +199,9 @@ func MakeExternalAPI(metricsName string, f func(*http.Request) util.JSONResponse
 // This is used to serve HTML alongside JSON error messages
 func MakeHTMLAPI(metricsName string, enableMetrics bool, f func(http.ResponseWriter, *http.Request)) http.Handler {
 	withSpan := func(w http.ResponseWriter, req *http.Request) {
-		span := opentracing.StartSpan(metricsName)
-		defer span.Finish()
-		req = req.WithContext(opentracing.ContextWithSpan(req.Context(), span))
+		trace, ctx := internal.StartTask(req.Context(), metricsName)
+		defer trace.EndTask()
+		req = req.WithContext(ctx)
 		f(w, req)
 	}
 
@@ -223,57 +222,6 @@ func MakeHTMLAPI(metricsName string, enableMetrics bool, f func(http.ResponseWri
 	)
 }
 
-// MakeInternalAPI turns a util.JSONRequestHandler function into an http.Handler.
-// This is used for APIs that are internal to dendrite.
-// If we are passed a tracing context in the request headers then we use that
-// as the parent of any tracing spans we create.
-func MakeInternalAPI(metricsName string, enableMetrics bool, f func(*http.Request) util.JSONResponse) http.Handler {
-	h := util.MakeJSONAPI(util.NewJSONRequestHandler(f))
-	withSpan := func(w http.ResponseWriter, req *http.Request) {
-		carrier := opentracing.HTTPHeadersCarrier(req.Header)
-		tracer := opentracing.GlobalTracer()
-		clientContext, err := tracer.Extract(opentracing.HTTPHeaders, carrier)
-		var span opentracing.Span
-		if err == nil {
-			// Default to a span without RPC context.
-			span = tracer.StartSpan(metricsName)
-		} else {
-			// Set the RPC context.
-			span = tracer.StartSpan(metricsName, ext.RPCServerOption(clientContext))
-		}
-		defer span.Finish()
-		req = req.WithContext(opentracing.ContextWithSpan(req.Context(), span))
-		h.ServeHTTP(w, req)
-	}
-
-	if !enableMetrics {
-		return http.HandlerFunc(withSpan)
-	}
-
-	return promhttp.InstrumentHandlerCounter(
-		promauto.NewCounterVec(
-			prometheus.CounterOpts{
-				Name:      metricsName + "_requests_total",
-				Help:      "Total number of internal API calls",
-				Namespace: "dendrite",
-			},
-			[]string{"code"},
-		),
-		promhttp.InstrumentHandlerResponseSize(
-			promauto.NewHistogramVec(
-				prometheus.HistogramOpts{
-					Namespace: "dendrite",
-					Name:      metricsName + "_response_size_bytes",
-					Help:      "A histogram of response sizes for requests.",
-					Buckets:   []float64{200, 500, 900, 1500, 5000, 15000, 50000, 100000},
-				},
-				[]string{},
-			),
-			http.HandlerFunc(withSpan),
-		),
-	)
-}
-
 // WrapHandlerInBasicAuth adds basic auth to a handler. Only used for /metrics
 func WrapHandlerInBasicAuth(h http.Handler, b BasicAuth) http.HandlerFunc {
 	if b.Username == "" || b.Password == "" {
diff --git a/internal/pushgateway/client.go b/internal/pushgateway/client.go
index 259239b8..d5671be3 100644
--- a/internal/pushgateway/client.go
+++ b/internal/pushgateway/client.go
@@ -10,8 +10,6 @@ import (
 	"time"
 
 	"github.com/matrix-org/dendrite/internal"
-
-	"github.com/opentracing/opentracing-go"
 )
 
 type httpClient struct {
@@ -34,8 +32,8 @@ func NewHTTPClient(disableTLSValidation bool) Client {
 }
 
 func (h *httpClient) Notify(ctx context.Context, url string, req *NotifyRequest, resp *NotifyResponse) error {
-	span, ctx := opentracing.StartSpanFromContext(ctx, "Notify")
-	defer span.Finish()
+	trace, ctx := internal.StartRegion(ctx, "Notify")
+	defer trace.EndRegion()
 
 	body, err := json.Marshal(req)
 	if err != nil {
diff --git a/internal/tracing.go b/internal/tracing.go
new file mode 100644
index 00000000..4e062aed
--- /dev/null
+++ b/internal/tracing.go
@@ -0,0 +1,64 @@
+// Copyright 2023 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+	"context"
+	"runtime/trace"
+
+	"github.com/opentracing/opentracing-go"
+)
+
+type Trace struct {
+	span   opentracing.Span
+	region *trace.Region
+	task   *trace.Task
+}
+
+func StartTask(inCtx context.Context, name string) (Trace, context.Context) {
+	ctx, task := trace.NewTask(inCtx, name)
+	span, ctx := opentracing.StartSpanFromContext(ctx, name)
+	return Trace{
+		span: span,
+		task: task,
+	}, ctx
+}
+
+func StartRegion(inCtx context.Context, name string) (Trace, context.Context) {
+	region := trace.StartRegion(inCtx, name)
+	span, ctx := opentracing.StartSpanFromContext(inCtx, name)
+	return Trace{
+		span:   span,
+		region: region,
+	}, ctx
+}
+
+func (t Trace) EndRegion() {
+	t.span.Finish()
+	if t.region != nil {
+		t.region.End()
+	}
+}
+
+func (t Trace) EndTask() {
+	t.span.Finish()
+	if t.task != nil {
+		t.task.End()
+	}
+}
+
+func (t Trace) SetTag(key string, value any) {
+	t.span.SetTag(key, value)
+}
diff --git a/internal/tracing_test.go b/internal/tracing_test.go
new file mode 100644
index 00000000..582f50c3
--- /dev/null
+++ b/internal/tracing_test.go
@@ -0,0 +1,25 @@
+package internal
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTracing(t *testing.T) {
+	inCtx := context.Background()
+
+	task, ctx := StartTask(inCtx, "testing")
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, task)
+	assert.NotEqual(t, inCtx, ctx)
+	task.SetTag("key", "value")
+
+	region, ctx2 := StartRegion(ctx, "testing")
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, region)
+	assert.NotEqual(t, ctx, ctx2)
+	defer task.EndTask()
+	defer region.EndRegion()
+}