chore: remove NodeID middleware and replace with static function (#2664)

Author: miparnisari

internal/dispatch/caching/caching.go

@@ -16,7 +16,6 @@ import (
 
 	"github.com/authzed/spicedb/internal/dispatch"
 	"github.com/authzed/spicedb/internal/dispatch/keys"
-	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/telemetry/otelconv"
 	"github.com/authzed/spicedb/pkg/cache"
 	"github.com/authzed/spicedb/pkg/middleware/nodeid"
@@ -162,11 +161,7 @@ func (cd *Dispatcher) DispatchCheck(ctx context.Context, req *v1.DispatchCheckRe
 			cd.checkFromCacheCounter.Inc()
 			// If debugging is requested, add the req and the response to the trace.
 			if req.Debug == v1.DispatchCheckRequest_ENABLE_BASIC_DEBUGGING {
-				nodeID, err := nodeid.FromContext(ctx)
-				if err != nil {
-					log.Err(err).Msg("failed to get nodeID from context")
-				}
-
+				nodeID := nodeid.Get()
 				response.Metadata.DebugInfo = &v1.DebugInformation{
 					Check: &v1.CheckDebugTrace{
 						Request:        req,

internal/dispatch/graph/graph.go

@@ -241,10 +241,7 @@ func (ld *localDispatcher) DispatchCheck(ctx context.Context, req *v1.DispatchCh
 	resourceType := tuple.StringCoreRR(req.ResourceRelation)
 	spanName := "DispatchCheck → " + resourceType + "@" + req.Subject.Namespace + "#" + req.Subject.Relation
 
-	nodeID, err := nodeid.FromContext(ctx)
-	if err != nil {
-		log.Err(err).Msg("failed to get node ID")
-	}
+	nodeID := nodeid.Get()
 
 	ctx, span := tracer.Start(ctx, spanName, trace.WithAttributes(
 		attribute.String(otelconv.AttrDispatchResourceType, resourceType),
@@ -263,12 +260,6 @@ func (ld *localDispatcher) DispatchCheck(ctx context.Context, req *v1.DispatchCh
 			}, rewriteError(ctx, err)
 		}
 
-		// NOTE: we return debug information here to ensure tooling can see the cycle.
-		nodeID, nerr := nodeid.FromContext(ctx)
-		if nerr != nil {
-			log.Err(nerr).Msg("failed to get nodeID from context")
-		}
-
 		return &v1.DispatchCheckResponse{
 			Metadata: &v1.ResponseMeta{
 				DispatchCount: 0,
@@ -337,11 +328,7 @@ func (ld *localDispatcher) DispatchCheck(ctx context.Context, req *v1.DispatchCh
 
 // DispatchExpand implements dispatch.Expand interface
 func (ld *localDispatcher) DispatchExpand(ctx context.Context, req *v1.DispatchExpandRequest) (*v1.DispatchExpandResponse, error) {
-	nodeID, err := nodeid.FromContext(ctx)
-	if err != nil {
-		log.Err(err).Msg("failed to get node ID")
-	}
-
+	nodeID := nodeid.Get()
 	ctx, span := tracer.Start(ctx, "DispatchExpand", trace.WithAttributes(
 		attribute.String(otelconv.AttrDispatchStart, tuple.StringCoreONR(req.ResourceAndRelation)),
 		attribute.String(otelconv.AttrDispatchNodeID, nodeID),
@@ -377,10 +364,7 @@ func (ld *localDispatcher) DispatchLookupResources2(
 	req *v1.DispatchLookupResources2Request,
 	stream dispatch.LookupResources2Stream,
 ) error {
-	nodeID, err := nodeid.FromContext(stream.Context())
-	if err != nil {
-		log.Err(err).Msg("failed to get node ID")
-	}
+	nodeID := nodeid.Get()
 
 	ctx, span := tracer.Start(stream.Context(), "DispatchLookupResources2", trace.WithAttributes(
 		attribute.String(otelconv.AttrDispatchResourceType, tuple.StringCoreRR(req.ResourceRelation)),
@@ -413,10 +397,7 @@ func (ld *localDispatcher) DispatchLookupResources3(
 	req *v1.DispatchLookupResources3Request,
 	stream dispatch.LookupResources3Stream,
 ) error {
-	nodeID, err := nodeid.FromContext(stream.Context())
-	if err != nil {
-		log.Err(err).Msg("failed to get node ID")
-	}
+	nodeID := nodeid.Get()
 
 	ctx, span := tracer.Start(stream.Context(), "DispatchLookupResources3", trace.WithAttributes(
 		attribute.String(otelconv.AttrDispatchResourceType, tuple.StringCoreRR(req.ResourceRelation)),
@@ -450,10 +431,7 @@ func (ld *localDispatcher) DispatchLookupSubjects(
 	req *v1.DispatchLookupSubjectsRequest,
 	stream dispatch.LookupSubjectsStream,
 ) error {
-	nodeID, err := nodeid.FromContext(stream.Context())
-	if err != nil {
-		log.Err(err).Msg("failed to get node ID")
-	}
+	nodeID := nodeid.Get()
 
 	resourceType := tuple.StringCoreRR(req.ResourceRelation)
 	subjectRelation := tuple.StringCoreRR(req.SubjectRelation)

internal/graph/check.go

@@ -116,12 +116,7 @@ func (cc *ConcurrentChecker) Check(ctx context.Context, req ValidatedCheckReques
 		return resolved.Resp, resolved.Err
 	}
 
-	nodeID, err := nodeid.FromContext(ctx)
-	if err != nil {
-		// NOTE: we ignore this error here as if the node ID is missing, the debug
-		// trace is still valid.
-		log.Err(err).Msg("failed to get node ID")
-	}
+	nodeID := nodeid.Get()
 
 	// Add debug information if requested.
 	debugInfo := resolved.Resp.Metadata.DebugInfo
@@ -1326,15 +1321,10 @@ func combineResponseMetadata(ctx context.Context, existing *v1.ResponseMeta, res
 		return combined
 	}
 
-	nodeID, err := nodeid.FromContext(ctx)
-	if err != nil {
-		log.Err(err).Msg("failed to get nodeID from context")
-	}
-
 	debugInfo := &v1.DebugInformation{
 		Check: &v1.CheckDebugTrace{
 			TraceId:  NewTraceID(),
-			SourceId: nodeID,
+			SourceId: nodeid.Get(),
 		},
 	}
 

pkg/cmd/server/defaults.go

@@ -39,7 +39,6 @@ import (
 	"github.com/authzed/spicedb/pkg/datastore"
 	consistencymw "github.com/authzed/spicedb/pkg/middleware/consistency"
 	logmw "github.com/authzed/spicedb/pkg/middleware/logging"
-	"github.com/authzed/spicedb/pkg/middleware/nodeid"
 	"github.com/authzed/spicedb/pkg/middleware/requestid"
 	"github.com/authzed/spicedb/pkg/middleware/serverversion"
 	"github.com/authzed/spicedb/pkg/releases"
@@ -167,7 +166,6 @@ var alwaysDebugOption = grpclog.WithLevels(func(code codes.Code) grpclog.Level {
 
 const (
 	DefaultMiddlewareRequestID     = "requestid"
-	DefaultMiddlewareNodeID        = "nodeid"
 	DefaultMiddlewareLog           = "log"
 	DefaultMiddlewareGRPCLog       = "grpclog"
 	DefaultMiddlewareGRPCAuth      = "grpcauth"
@@ -302,25 +300,18 @@ func DefaultUnaryMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.UnaryS
 			WithInterceptor(logmw.UnaryServerInterceptor(logmw.ExtractMetadataField(string(requestmeta.RequestIDKey), "requestID"))).
 			Done(),
 
-		NewUnaryMiddleware().
-			WithName(DefaultMiddlewareNodeID).
-			WithInterceptor(nodeid.UnaryServerInterceptor("")).
-			Done(),
-
 		NewUnaryMiddleware().
 			WithName(DefaultMiddlewareGRPCLog + "-debug").
 			WithInterceptor(selector.UnaryServerInterceptor(
 				grpclog.UnaryServerInterceptor(InterceptorLogger(opts.Logger), determineEventsToLog(opts), alwaysDebugOption, durationFieldOption, traceIDFieldOption),
 				selector.MatchFunc(matchesRoute(healthCheckRoute)))).
-			EnsureAlreadyExecuted(DefaultMiddlewareNodeID).
 			Done(),
 
 		NewUnaryMiddleware().
 			WithName(DefaultMiddlewareGRPCLog).
 			WithInterceptor(selector.UnaryServerInterceptor(
 				grpclog.UnaryServerInterceptor(InterceptorLogger(opts.Logger), determineEventsToLog(opts), defaultCodeToLevel, durationFieldOption, traceIDFieldOption),
 				selector.MatchFunc(doesNotMatchRoute(healthCheckRoute)))).
-			EnsureAlreadyExecuted(DefaultMiddlewareNodeID).
 			Done(),
 
 		NewUnaryMiddleware().
@@ -375,25 +366,18 @@ func DefaultStreamingMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.St
 			WithInterceptor(logmw.StreamServerInterceptor(logmw.ExtractMetadataField(string(requestmeta.RequestIDKey), "requestID"))).
 			Done(),
 
-		NewStreamMiddleware().
-			WithName(DefaultMiddlewareNodeID).
-			WithInterceptor(nodeid.StreamServerInterceptor("")).
-			Done(),
-
 		NewStreamMiddleware().
 			WithName(DefaultMiddlewareGRPCLog + "-debug").
 			WithInterceptor(selector.StreamServerInterceptor(
 				grpclog.StreamServerInterceptor(InterceptorLogger(opts.Logger), determineEventsToLog(opts), alwaysDebugOption, durationFieldOption, traceIDFieldOption),
 				selector.MatchFunc(matchesRoute(healthCheckRoute)))).
-			EnsureInterceptorAlreadyExecuted(DefaultMiddlewareNodeID).
 			Done(),
 
 		NewStreamMiddleware().
 			WithName(DefaultMiddlewareGRPCLog).
 			WithInterceptor(selector.StreamServerInterceptor(
 				grpclog.StreamServerInterceptor(InterceptorLogger(opts.Logger), determineEventsToLog(opts), defaultCodeToLevel, durationFieldOption, traceIDFieldOption),
 				selector.MatchFunc(doesNotMatchRoute(healthCheckRoute)))).
-			EnsureInterceptorAlreadyExecuted(DefaultMiddlewareNodeID).
 			Done(),
 
 		NewStreamMiddleware().
@@ -454,7 +438,6 @@ func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc
 	grpcMetricsUnaryInterceptor, grpcMetricsStreamingInterceptor := GRPCMetrics(disableGRPCLatencyHistogram)
 	return []grpc.UnaryServerInterceptor{
 			requestid.UnaryServerInterceptor(requestid.GenerateIfMissing(true)),
-			nodeid.UnaryServerInterceptor(""),
 			logmw.UnaryServerInterceptor(logmw.ExtractMetadataField(string(requestmeta.RequestIDKey), "requestID")),
 			grpclog.UnaryServerInterceptor(InterceptorLogger(logger), dispatchDefaultCodeToLevel, durationFieldOption, traceIDFieldOption),
 			grpcMetricsUnaryInterceptor,
@@ -465,7 +448,6 @@ func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc
 			// NOTE: the logging middlewares are not present here in streaming, to remove their significant overhead
 			// when returning streaming messages.
 			requestid.StreamServerInterceptor(requestid.GenerateIfMissing(true)),
-			nodeid.StreamServerInterceptor(""),
 			grpcMetricsStreamingInterceptor,
 			grpcauth.StreamServerInterceptor(authFunc),
 			datastoremw.StreamServerInterceptor(ds),

pkg/middleware/nodeid/nodeid.go

@@ -1,27 +1,16 @@
 package nodeid
 
 import (
-	"context"
 	"encoding/hex"
 	"os"
 
 	"github.com/cespare/xxhash/v2"
-	middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
 	"github.com/rs/zerolog/log"
-	"google.golang.org/grpc"
 )
 
 const spiceDBPrefix = "spicedb:"
 
-type ctxKeyType struct{}
-
-var nodeIDKey ctxKeyType = struct{}{}
-
-type nodeIDHandle struct {
-	nodeID string
-}
-
-var defaultNodeID string
+var computedNodeID string
 
 func init() {
 	hostname, err := os.Hostname()
@@ -37,66 +26,10 @@ func init() {
 		return
 	}
 
-	defaultNodeID = spiceDBPrefix + hex.EncodeToString(hasher.Sum(nil))
-}
-
-// ContextWithHandle adds a placeholder to a context that will later be
-// filled by the Node ID.
-func ContextWithHandle(ctx context.Context) context.Context {
-	return context.WithValue(ctx, nodeIDKey, &nodeIDHandle{})
+	computedNodeID = spiceDBPrefix + hex.EncodeToString(hasher.Sum(nil))
 }
 
-// FromContext reads the node's ID out of a context.Context.
-func FromContext(ctx context.Context) (string, error) {
-	if c := ctx.Value(nodeIDKey); c != nil {
-		handle := c.(*nodeIDHandle)
-		if handle.nodeID != "" {
-			return handle.nodeID, nil
-		}
-	}
-
-	if err := setInContext(ctx, defaultNodeID); err != nil {
-		return "", err
-	}
-
-	return defaultNodeID, nil
-}
-
-// setInContext adds a node ID to the given context
-func setInContext(ctx context.Context, nodeID string) error {
-	handle := ctx.Value(nodeIDKey)
-	if handle == nil {
-		return nil
-	}
-	handle.(*nodeIDHandle).nodeID = nodeID
-	return nil
-}
-
-// UnaryServerInterceptor returns a new unary server interceptor that adds the
-// node ID to the context. If empty, spicedb:$hostname is used.
-func UnaryServerInterceptor(nodeID string) grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
-		newCtx := ContextWithHandle(ctx)
-		if nodeID != "" {
-			if err := setInContext(newCtx, nodeID); err != nil {
-				return nil, err
-			}
-		}
-		return handler(newCtx, req)
-	}
-}
-
-// StreamServerInterceptor returns a new stream server interceptor that adds the
-// node ID to the context. If empty, spicedb:$hostname is used.
-func StreamServerInterceptor(nodeID string) grpc.StreamServerInterceptor {
-	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		wrapped := middleware.WrapServerStream(stream)
-		wrapped.WrappedContext = ContextWithHandle(wrapped.WrappedContext)
-		if nodeID != "" {
-			if err := setInContext(wrapped.WrappedContext, nodeID); err != nil {
-				return err
-			}
-		}
-		return handler(srv, wrapped)
-	}
+// Get returns the hostname.
+func Get() string {
+	return computedNodeID
 }