chore: add more spans in middlewares

Author: miparnisari

internal/middleware/handwrittenvalidation/handwrittenvalidation.go

@@ -3,6 +3,7 @@ package handwrittenvalidation
 import (
 	"context"
 
+	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -12,12 +13,16 @@ type handwrittenValidator interface {
 	HandwrittenValidate() error
 }
 
+var tracer = otel.Tracer("spicedb/internal/middleware")
+
 // UnaryServerInterceptor returns a new unary server interceptor that runs the handwritten validation
 // on the incoming request, if any.
 func UnaryServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 	validator, ok := req.(handwrittenValidator)
 	if ok {
+		_, span := tracer.Start(ctx, "handwritten.validation")
 		err := validator.HandwrittenValidate()
+		span.End()
 		if err != nil {
 			return nil, status.Errorf(codes.InvalidArgument, "%s", err)
 		}

internal/middleware/memoryprotection/memory_protection.go

@@ -9,13 +9,16 @@ import (
 	middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
 	log "github.com/authzed/spicedb/internal/logging"
 )
 
+var tracer = otel.Tracer("spicedb/internal/middleware/memory_protection")
+
 // RequestsProcessed tracks requests that were processed by this middleware.
 var RequestsProcessed = promauto.NewCounterVec(prometheus.CounterOpts{
 	Namespace: "spicedb",
@@ -44,7 +47,7 @@ func New(usageProvider MemoryUsageProvider, name string) *MemoryProtectionMiddle
 // UnaryServerInterceptor returns a unary server interceptor that rejects incoming requests is memory usage is too high
 func (am *MemoryProtectionMiddleware) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
 	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
-		if err := am.checkAdmission(info.FullMethod); err != nil {
+		if err := am.checkAdmission(ctx, info.FullMethod); err != nil {
 			return nil, err
 		}
 
@@ -55,7 +58,7 @@ func (am *MemoryProtectionMiddleware) UnaryServerInterceptor() grpc.UnaryServerI
 // StreamServerInterceptor returns a stream server interceptor that rejects incoming requests is memory usage is too high
 func (am *MemoryProtectionMiddleware) StreamServerInterceptor() grpc.StreamServerInterceptor {
 	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		if err := am.checkAdmission(info.FullMethod); err != nil {
+		if err := am.checkAdmission(stream.Context(), info.FullMethod); err != nil {
 			return err
 		}
 
@@ -65,7 +68,10 @@ func (am *MemoryProtectionMiddleware) StreamServerInterceptor() grpc.StreamServe
 }
 
 // checkAdmission returns an error if the request should be denied because memory usage is too high.
-func (am *MemoryProtectionMiddleware) checkAdmission(method string) error {
+func (am *MemoryProtectionMiddleware) checkAdmission(ctx context.Context, method string) error {
+	_, span := tracer.Start(ctx, "checkMemoryUsage")
+	span.End()
+
 	accept := true
 	defer func() {
 		am.recordMetric(method, accept)

internal/middleware/memoryprotection/memory_protection_test.go

@@ -48,7 +48,7 @@ func TestNew(t *testing.T) {
 			am := New(tt.inputProvider, "name")
 			require.NotNil(t, am)
 
-			err := am.checkAdmission("some_method")
+			err := am.checkAdmission(t.Context(), "some_method")
 			if tt.expectReqLetThrough {
 				require.NoError(t, err) // if the middleware is off, every request is let through
 			} else {

internal/middleware/perfinsights/perfinsights.go

@@ -11,6 +11,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/rs/zerolog/log"
+	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
 
@@ -69,6 +70,8 @@ var APIShapeLatency = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	NativeHistogramBucketFactor: 1.1,
 }, append([]string{"api_kind"}, allLabels...))
 
+var tracer = otel.Tracer("spicedb/internal/middleware")
+
 // ShapeBuilder is a function that returns a slice of strings representing the shape of the API call.
 // This is used to report the shape of the API call to Prometheus.
 type ShapeBuilder func() APIShapeLabels
@@ -79,6 +82,8 @@ func ObserveShapeLatency(ctx context.Context, methodName string, shape APIShapeL
 }
 
 func observeShapeLatency(ctx context.Context, metric *prometheus.HistogramVec, methodName string, shape APIShapeLabels, duration time.Duration) {
+	ctx, span := tracer.Start(ctx, "perfInsights.observeShapeLatency")
+	defer span.End()
 	labels := buildLabels(methodName, shape)
 	if len(labels) == 0 {
 		log.Warn().Str("method", methodName).

internal/middleware/usagemetrics/usagemetrics.go

@@ -9,6 +9,7 @@ import (
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"go.opentelemetry.io/otel"
 	"google.golang.org/grpc"
 
 	"github.com/authzed/authzed-go/pkg/responsemeta"
@@ -34,6 +35,8 @@ var (
 		Help:      "Histogram of cluster dispatches performed by the instance.",
 		Buckets:   []float64{1, 5, 10, 25, 50, 100, 250},
 	}, DispatchedCountLabels)
+
+	tracer = otel.Tracer("spicedb/internal/middleware")
 )
 
 type reporter struct{}
@@ -52,6 +55,9 @@ type serverReporter struct {
 
 // PostCall is invoked after all PostMsgSend operations.
 func (r *serverReporter) PostCall(_ error, _ time.Duration) {
+	_, span := tracer.Start(r.ctx, "usagemetrics.PostCall")
+	defer span.End()
+
 	responseMeta := FromContext(r.ctx)
 	if responseMeta == nil {
 		responseMeta = &dispatch.ResponseMeta{}