Added a flag to disable tracing of health check requests

Author: davide-baldo

internal/gateway/gateway.go

@@ -18,6 +18,7 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
 
 	"github.com/authzed/authzed-go/proto"
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
@@ -26,6 +27,9 @@ import (
 	"github.com/authzed/spicedb/internal/grpchelpers"
 )
 
+const healthCheckRoute = "/grpc.health.v1.Health/Check"
+const healthCheckHttpPath = "/healthz"
+
 var histogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	Namespace: "spicedb",
 	Subsystem: "rest_gateway",
@@ -35,14 +39,23 @@ var histogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
 
 // NewHandler creates an REST gateway HTTP CloserHandler with the provided upstream
 // configuration.
-func NewHandler(ctx context.Context, upstreamAddr, upstreamTLSCertPath string) (*CloserHandler, error) {
+func NewHandler(ctx context.Context, upstreamAddr, upstreamTLSCertPath string, disableHealthCheckTracing bool) (*CloserHandler, error) {
 	if upstreamAddr == "" {
 		return nil, fmt.Errorf("upstreamAddr must not be empty")
 	}
 
-	opts := []grpc.DialOption{
-		grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
+	var opts []grpc.DialOption
+
+	if disableHealthCheckTracing {
+		opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler(
+			otelgrpc.WithFilter(func(info *stats.RPCTagInfo) bool {
+				return info.FullMethodName != healthCheckRoute
+			}),
+		)))
+	} else {
+		opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler()))
 	}
+
 	if upstreamTLSCertPath == "" {
 		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	} else {
@@ -85,7 +98,15 @@ func NewHandler(ctx context.Context, upstreamAddr, upstreamTLSCertPath string) (
 	}))
 	mux.Handle("/", gwMux)
 
-	finalHandler := promhttp.InstrumentHandlerDuration(histogram, otelhttp.NewHandler(mux, "gateway"))
+	var finalHandler http.Handler
+	if disableHealthCheckTracing {
+		// Only apply tracing to non-health check endpoints
+		finalHandler = promhttp.InstrumentHandlerDuration(histogram,
+			skipHealthCheckTracing(mux, otelhttp.NewHandler(mux, "gateway")))
+	} else {
+		finalHandler = promhttp.InstrumentHandlerDuration(histogram, otelhttp.NewHandler(mux, "gateway"))
+	}
+
 	return newCloserHandler(finalHandler, schemaConn, permissionsConn, watchConn, healthConn, experimentalConn), nil
 }
 
@@ -143,6 +164,19 @@ func registerHandler(ctx context.Context, mux *runtime.ServeMux, endpoint string
 	return conn, nil
 }
 
+// skipHealthCheckTracing creates an HTTP handler that skips OpenTelemetry tracing for health check endpoints
+func skipHealthCheckTracing(directHandler http.Handler, tracedHandler http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == healthCheckHttpPath {
+			// Skip tracing for health check endpoints
+			directHandler.ServeHTTP(w, r)
+		} else {
+			// Apply OpenTelemetry tracing for all other endpoints
+			tracedHandler.ServeHTTP(w, r)
+		}
+	})
+}
+
 var defaultOtelOpts = []otelgrpc.Option{
 	otelgrpc.WithPropagators(otel.GetTextMapPropagator()),
 	otelgrpc.WithTracerProvider(otel.GetTracerProvider()),

internal/gateway/gateway_test.go

@@ -1,15 +1,24 @@
 package gateway
 
 import (
+	"context"
+	"net"
 	"net/http"
+	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/goleak"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 
 	"github.com/authzed/spicedb/pkg/testutil"
 )
@@ -48,11 +57,159 @@ func TestOtelForwarding(t *testing.T) {
 func TestCloseConnections(t *testing.T) {
 	defer goleak.VerifyNone(t, append(testutil.GoLeakIgnores(), goleak.IgnoreCurrent())...)
 
-	gatewayHandler, err := NewHandler(t.Context(), "192.0.2.0:4321", "")
+	gatewayHandler, err := NewHandler(t.Context(), "192.0.2.0:4321", "", false)
 	require.NoError(t, err)
 	// 4 conns for permission+schema+watch+experimental services, 1 for health check
 	require.Len(t, gatewayHandler.closers, 5)
 
 	// if connections are not closed, goleak would detect it
 	require.NoError(t, gatewayHandler.Close())
 }
+
+func TestGatewayHealthCheckTracing(t *testing.T) {
+	tests := []struct {
+		name                      string
+		disableHealthCheckTracing bool
+		description               string
+	}{
+		{
+			name:                      "health check tracing disabled",
+			disableHealthCheckTracing: true,
+			description:               "gateway should skip tracing for health checks when flag is true",
+		},
+		{
+			name:                      "health check tracing enabled",
+			disableHealthCheckTracing: false,
+			description:               "gateway should trace health checks when flag is false",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer goleak.VerifyNone(t, append(testutil.GoLeakIgnores(), goleak.IgnoreCurrent())...)
+
+			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+			defer cancel()
+
+			// Since we can't use a real gRPC server, we need to mock one
+			mockServer := setupMockGRPCServer(t, ctx)
+			defer mockServer.Stop()
+
+			defaultProvider := otel.GetTracerProvider()
+			defer otel.SetTracerProvider(defaultProvider)
+
+			provider := sdktrace.NewTracerProvider(
+				sdktrace.WithSampler(sdktrace.AlwaysSample()),
+			)
+			spanrecorder := tracetest.NewSpanRecorder()
+			provider.RegisterSpanProcessor(spanrecorder)
+			otel.SetTracerProvider(provider)
+
+			gatewayHandler, err := NewHandler(ctx, mockServer.Addr(), "", tt.disableHealthCheckTracing)
+			require.NoError(t, err)
+			defer gatewayHandler.Close()
+
+			// Test health check endpoint
+			req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
+			w := httptest.NewRecorder()
+			gatewayHandler.ServeHTTP(w, req)
+			require.Equal(t, http.StatusOK, w.Code)
+
+			healthSpans := spanrecorder.Ended()
+			spanrecorder.Reset()
+
+			// Test regular endpoint
+			req = httptest.NewRequest(http.MethodGet, "/openapi.json", nil)
+			w = httptest.NewRecorder()
+			gatewayHandler.ServeHTTP(w, req)
+			require.Equal(t, http.StatusOK, w.Code)
+
+			regularSpans := spanrecorder.Ended()
+
+			if tt.disableHealthCheckTracing {
+				require.Len(t, healthSpans, 0, "Health check should not create spans when tracing is disabled")
+
+				healthFound := false
+				for _, span := range healthSpans {
+					if span.Name() == "grpc.health.v1.Health/Check" {
+						healthFound = true
+						break
+					}
+				}
+				require.False(t, healthFound, "Unexpectedly found health check span when tracing is disabled")
+
+				gatewayFound := false
+				for _, span := range healthSpans {
+					if span.Name() == "gateway" {
+						gatewayFound = true
+						break
+					}
+				}
+				require.False(t, gatewayFound, "Unexpectedly found gateway span for health check when tracing is disabled")
+			} else {
+				require.True(t, len(healthSpans) > 0, "Health check should create spans when tracing is enabled")
+
+				healthFound := false
+				for _, span := range healthSpans {
+					if span.Name() == "grpc.health.v1.Health/Check" {
+						healthFound = true
+						break
+					}
+				}
+				require.True(t, healthFound, "Expected to find health check span when tracing is enabled")
+			}
+
+			require.True(t, len(regularSpans) > 0, "Regular endpoints should create spans")
+			regularFound := false
+			for _, span := range regularSpans {
+				if span.Name() == "gateway" {
+					regularFound = true
+					break
+				}
+			}
+			require.True(t, regularFound, "Expected to find gateway span for regular endpoint")
+
+		})
+	}
+}
+
+type mockGRPCServer struct {
+	server   *grpc.Server
+	listener net.Listener
+	addr     string
+}
+
+func (m *mockGRPCServer) Addr() string {
+	return m.addr
+}
+
+func (m *mockGRPCServer) Stop() {
+	if m.server != nil {
+		m.server.Stop()
+	}
+	if m.listener != nil {
+		m.listener.Close()
+	}
+}
+
+func setupMockGRPCServer(t *testing.T, ctx context.Context) *mockGRPCServer {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	server := grpc.NewServer()
+
+	healthServer := health.NewServer()
+	healthpb.RegisterHealthServer(server, healthServer)
+
+	go func() {
+		if err := server.Serve(listener); err != nil {
+			t.Logf("Mock gRPC server error: %v", err)
+		}
+	}()
+
+	return &mockGRPCServer{
+		server:   server,
+		listener: listener,
+		addr:     listener.Addr().String(),
+	}
+}

pkg/cmd/serve.go

@@ -184,6 +184,7 @@ func RegisterServeFlags(cmd *cobra.Command, config *server.Config) error {
 	otel := cobraotel.New("spicedb")
 	otel.RegisterFlags(observabilityFlags)
 	runtime.RegisterFlags(observabilityFlags)
+	observabilityFlags.BoolVar(&config.DisableHealthCheckOTelTracingMiddleware, "otel-disable-healthcheck-tracing", false, "disables telemetry tracing for health check API calls to reduce trace volume")
 
 	metricsFlags := nfs.FlagSet(BoldBlue("Metrics Server"))
 	// Flags for metrics