feat: add tracing to proto validation (#2950)

Author: miparnisari

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## [Unreleased]
 ### Changed
 - Updated CI so that Postgres tests run against v18 which is GA and not against v13 which is EOL (https://github.com/authzed/spicedb/pull/2926)
+- Added tracing to request validation (https://github.com/authzed/spicedb/pull/2950)
 
 ### Fixed
 - Regression introduced in 1.49.2: missing spans in ReadSchema calls (https://github.com/authzed/spicedb/pull/2947)

internal/middleware/handwrittenvalidation/doc.go

@@ -1,2 +1,2 @@
-// Package handwrittenvalidation defines middleware that runs custom-made validations on incoming requests.
+// Package handwrittenvalidation defines middleware that runs validations on incoming requests.
 package handwrittenvalidation

internal/middleware/interceptorwrapper/interceptorwrapper.go

@@ -0,0 +1,33 @@
+package interceptorwrapper
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel"
+	"google.golang.org/grpc"
+)
+
+var tracer = otel.Tracer("spicedb/internal/middleware")
+
+// WrapUnaryServerInterceptorWithSpans returns a new interceptor that wraps the given interceptor
+// with a span, measuring the duration of the interceptor's pre-handler logic.
+func WrapUnaryServerInterceptorWithSpans(
+	inner grpc.UnaryServerInterceptor,
+	spanName string,
+) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		_, span := tracer.Start(ctx, spanName)
+		// NOTE: this shim is what lets us measure how long the interceptor is doing work.
+		// It's the handler that we pass to the wrapped interceptor, so `span.End()` will be
+		// called when the handler itself is called.
+		shimHandler := func(ctx context.Context, req any) (any, error) {
+			span.End()
+			return handler(ctx, req)
+		}
+		resp, err := inner(ctx, req, info, shimHandler)
+		if span.IsRecording() {
+			span.End()
+		}
+		return resp, err
+	}
+}

internal/services/v1/experimental.go

@@ -24,6 +24,7 @@ import (
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/middleware"
 	"github.com/authzed/spicedb/internal/middleware/handwrittenvalidation"
+	"github.com/authzed/spicedb/internal/middleware/interceptorwrapper"
 	"github.com/authzed/spicedb/internal/middleware/perfinsights"
 	"github.com/authzed/spicedb/internal/middleware/streamtimeout"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
@@ -99,7 +100,7 @@ func NewExperimentalServer(dispatch dispatch.Dispatcher, permServerConfig Permis
 	return &experimentalServer{
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
+				interceptorwrapper.WrapUnaryServerInterceptorWithSpans(grpcvalidate.UnaryServerInterceptor(validator), "protovalidate"),
 				handwrittenvalidation.UnaryServerInterceptor,
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(permServerConfig.PerformanceInsightMetricsEnabled),

internal/services/v1/relationships.go

@@ -19,6 +19,7 @@ import (
 	"github.com/authzed/spicedb/internal/dispatch"
 	"github.com/authzed/spicedb/internal/middleware"
 	"github.com/authzed/spicedb/internal/middleware/handwrittenvalidation"
+	"github.com/authzed/spicedb/internal/middleware/interceptorwrapper"
 	"github.com/authzed/spicedb/internal/middleware/perfinsights"
 	"github.com/authzed/spicedb/internal/middleware/streamtimeout"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
@@ -153,7 +154,7 @@ func NewPermissionsServer(
 		config:   configWithDefaults,
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
+				interceptorwrapper.WrapUnaryServerInterceptorWithSpans(grpcvalidate.UnaryServerInterceptor(validator), "protovalidate"),
 				handwrittenvalidation.UnaryServerInterceptor,
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(configWithDefaults.PerformanceInsightMetricsEnabled),

internal/services/v1/schema.go

@@ -11,6 +11,7 @@ import (
 
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/middleware"
+	"github.com/authzed/spicedb/internal/middleware/interceptorwrapper"
 	"github.com/authzed/spicedb/internal/middleware/perfinsights"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
@@ -50,7 +51,7 @@ func NewSchemaServer(config SchemaServerConfig) v1.SchemaServiceServer {
 	return &schemaServer{
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
+				interceptorwrapper.WrapUnaryServerInterceptorWithSpans(grpcvalidate.UnaryServerInterceptor(validator), "protovalidate"),
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(config.PerformanceInsightMetricsEnabled),
 			),

internal/services/v1/watch.go

@@ -52,7 +52,7 @@ func (ws *watchServer) Watch(req *v1.WatchRequest, stream v1.WatchService_WatchS
 		}
 	}
 
-	objectTypes := mapz.NewSet[string](req.GetOptionalObjectTypes()...)
+	objectTypes := mapz.NewSet(req.GetOptionalObjectTypes()...)
 
 	ctx := stream.Context()
 	dl := datalayer.MustFromContext(ctx)