chore: improvements and fix building of dispatch middleware

Author: miparnisari

.gitattributes

@@ -0,0 +1,4 @@
+internal/mocks/*.go linguist-generated=true
+*.pb.go linguist-generated=true
+*.pb.*.go linguist-generated=true
+proto/internal/buf.lock linguist-generated=true

development/prometheus.yaml

@@ -6,6 +6,9 @@ global:
 scrape_configs:
   - job_name: "spicedb"
     static_configs:
-      - targets: ["spicedb:9090"]
+      - targets: ["spicedb-1:9090"]
         labels:
-          service: "spicedb"
+          service: "spicedb-1"
+      - targets: ["spicedb-2:9090"]
+        labels:
+          service: "spicedb-2"

go.mod

@@ -112,6 +112,7 @@ require (
 	go.opentelemetry.io/otel/trace v1.38.0
 	go.uber.org/atomic v1.11.0
 	go.uber.org/goleak v1.3.0
+	go.uber.org/mock v0.6.0
 	golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b
 	golang.org/x/mod v0.28.0
 	golang.org/x/sync v0.17.0
@@ -137,6 +138,8 @@ tool (
 	github.com/golangci/golangci-lint/v2/cmd/golangci-lint
 	// support running mage with go run mage.go
 	github.com/magefile/mage/mage
+	// mocks are generated with go:generate directives.
+	go.uber.org/mock/mockgen
 	// vulncheck always uses the current directory's go.mod.
 	golang.org/x/vuln/cmd/govulncheck
 )

go.sum

@@ -2595,6 +2595,8 @@ go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwE
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
+go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y=
+go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU=
 go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=

internal/dispatch/dispatch.go

@@ -1,3 +1,5 @@
+//go:generate go run go.uber.org/mock/mockgen -source dispatch.go -destination ../mocks/mock_dispatcher.go -package mocks Dispatcher
+
 package dispatch
 
 import (

internal/middleware/memoryprotection/memoryprotection.go

@@ -18,27 +18,29 @@ import (
 	log "github.com/authzed/spicedb/internal/logging"
 )
 
+const DefaultSampleIntervalSeconds = 1
+
 var (
 	// RejectedRequestsCounter tracks requests rejected due to memory pressure
 	RejectedRequestsCounter = promauto.NewCounterVec(prometheus.CounterOpts{
 		Namespace: "spicedb",
-		Subsystem: "admission",
-		Name:      "memory_overload_rejected_requests_total",
+		Subsystem: "memory_admission",
+		Name:      "rejected_requests_total",
 		Help:      "Total requests rejected due to memory pressure",
 	}, []string{"endpoint"})
 
 	// MemoryUsageGauge tracks current memory usage percentage
 	MemoryUsageGauge = promauto.NewGauge(prometheus.GaugeOpts{
 		Namespace: "spicedb",
-		Subsystem: "admission",
+		Subsystem: "memory_admission",
 		Name:      "memory_usage_percent",
 		Help:      "Current memory usage as percentage of GOMEMLIMIT",
 	})
 )
 
 // Config holds configuration for the memory protection middleware
 type Config struct {
-	// ThresholdPercent is the memory usage threshold for requests (0-100)
+	// ThresholdPercent is the memory usage threshold for requests. If zero or negative, this middleware has no effect
 	ThresholdPercent int
 	// SampleIntervalSeconds controls how often memory usage is sampled
 	SampleIntervalSeconds int
@@ -60,56 +62,90 @@ func DefaultDispatchConfig() Config {
 	}
 }
 
-// AdmissionMiddleware implements memory-based admission control
-type AdmissionMiddleware struct {
-	config          Config
-	memoryLimit     int64
-	lastMemoryUsage atomic.Int64
-	metricsSamples  []metrics.Sample
-	ctx             context.Context
+// MemoryLimitProvider gets and sets the limit of memory usage.
+// In production, use DefaultMemoryLimitProvider.
+// For testing, use HardCodedMemoryLimitProvider.
+type MemoryLimitProvider interface {
+	Get() int64
+	Set(int64)
+}
+
+var (
+	_ MemoryLimitProvider = (*DefaultMemoryLimitProvider)(nil)
+	_ MemoryLimitProvider = (*HardCodedMemoryLimitProvider)(nil)
+)
+
+type DefaultMemoryLimitProvider struct{}
+
+func (p *DefaultMemoryLimitProvider) Get() int64 {
+	// SetMemoryLimit returns the previously set memory limit.
+	// A negative input does not adjust the limit, and allows for retrieval of the currently set memory limit
+	return debug.SetMemoryLimit(-1)
+}
+
+func (p *DefaultMemoryLimitProvider) Set(limit int64) {
+	debug.SetMemoryLimit(limit)
+}
+
+type HardCodedMemoryLimitProvider struct {
+	Hardcodedlimit int64
 }
 
-// New creates a new memory protection middleware with the given context
-func New(ctx context.Context, config Config) *AdmissionMiddleware {
-	// Use the provided context directly
-	mwCtx := ctx
+func (p *HardCodedMemoryLimitProvider) Get() int64 {
+	return p.Hardcodedlimit
+}
+
+func (p *HardCodedMemoryLimitProvider) Set(limit int64) {
+	p.Hardcodedlimit = limit
+}
+
+type MemoryAdmissionMiddleware struct {
+	config         Config
+	memoryLimit    int64 // -1 means no limit
+	metricsSamples []metrics.Sample
+	ctx            context.Context // to stop the background process
+
+	lastMemorySampleInBytes   *atomic.Uint64 // atomic because it's written inside a goroutine but can be read from anywhere
+	timestampLastMemorySample *atomic.Pointer[time.Time]
+}
+
+// New creates a new memory admission middleware with the given context.
+// Whe the context is cancelled, this middleware stops its background processing.
+func New(ctx context.Context, config Config, limitProvider MemoryLimitProvider) MemoryAdmissionMiddleware {
+	am := MemoryAdmissionMiddleware{
+		config:                    config,
+		lastMemorySampleInBytes:   &atomic.Uint64{},
+		timestampLastMemorySample: &atomic.Pointer[time.Time]{},
+		memoryLimit:               -1, // disabled initially
+		ctx:                       ctx,
+	}
 
 	// Get the current GOMEMLIMIT
-	memoryLimit := debug.SetMemoryLimit(-1)
+	memoryLimit := limitProvider.Get()
 	if memoryLimit < 0 {
 		// If no limit is set, we can't provide memory protection
 		log.Info().Msg("GOMEMLIMIT not set, memory protection disabled")
-		return &AdmissionMiddleware{
-			config:      config,
-			memoryLimit: -1, // Disabled
-			ctx:         mwCtx,
-		}
+		return am
 	}
 
-	// Check if memory protection is disabled via config
 	if config.ThresholdPercent <= 0 {
 		log.Info().Msg("memory protection disabled via configuration")
-		return &AdmissionMiddleware{
-			config:      config,
-			memoryLimit: -1, // Disabled
-			ctx:         mwCtx,
-		}
+		return am
 	}
 
-	am := &AdmissionMiddleware{
-		config:      config,
-		memoryLimit: memoryLimit,
-		metricsSamples: []metrics.Sample{
-			{Name: "/memory/classes/heap/objects:bytes"},
-		},
-		ctx: mwCtx,
+	if config.SampleIntervalSeconds <= 0 {
+		log.Info().Msgf("memory protection sample interval cannot be zero or negative; using default value of %q seconds", DefaultSampleIntervalSeconds)
+		am.config.SampleIntervalSeconds = DefaultSampleIntervalSeconds
 	}
 
-	// Initialize with current memory usage
-	if err := am.sampleMemory(); err != nil {
-		log.Warn().Err(err).Msg("failed to get initial memory sample")
+	am.memoryLimit = memoryLimit
+	am.metricsSamples = []metrics.Sample{
+		{Name: "/memory/classes/heap/objects:bytes"},
 	}
 
+	// Initialize with current memory usage
+	am.sampleMemory()
+
 	// Start background sampling with context
 	am.startBackgroundSampling()
 
@@ -122,9 +158,9 @@ func New(ctx context.Context, config Config) *AdmissionMiddleware {
 	return am
 }
 
-// UnaryServerInterceptor returns a unary server interceptor that implements admission control
-func (am *AdmissionMiddleware) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+// UnaryServerInterceptor returns a unary server interceptor that rejects incoming requests is memory usage is too high
+func (am *MemoryAdmissionMiddleware) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 		if am.memoryLimit < 0 {
 			// Memory protection is disabled
 			return handler(ctx, req)
@@ -139,9 +175,9 @@ func (am *AdmissionMiddleware) UnaryServerInterceptor() grpc.UnaryServerIntercep
 	}
 }
 
-// StreamServerInterceptor returns a stream server interceptor that implements admission control
-func (am *AdmissionMiddleware) StreamServerInterceptor() grpc.StreamServerInterceptor {
-	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+// StreamServerInterceptor returns a stream server interceptor that rejects incoming requests is memory usage is too high
+func (am *MemoryAdmissionMiddleware) StreamServerInterceptor() grpc.StreamServerInterceptor {
+	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		if am.memoryLimit < 0 {
 			// Memory protection is disabled
 			return handler(srv, stream)
@@ -158,8 +194,8 @@ func (am *AdmissionMiddleware) StreamServerInterceptor() grpc.StreamServerInterc
 }
 
 // checkAdmission determines if a request should be admitted based on current memory usage
-func (am *AdmissionMiddleware) checkAdmission(fullMethod string) error {
-	memoryUsage := am.getCurrentMemoryUsage()
+func (am *MemoryAdmissionMiddleware) checkAdmission(fullMethod string) error {
+	memoryUsage := am.getLastMemorySampleInBytes()
 
 	usagePercent := float64(memoryUsage) / float64(am.memoryLimit) * 100
 
@@ -181,24 +217,23 @@ func (am *AdmissionMiddleware) checkAdmission(fullMethod string) error {
 }
 
 // startBackgroundSampling starts a background goroutine that samples memory usage periodically
-func (am *AdmissionMiddleware) startBackgroundSampling() {
+func (am *MemoryAdmissionMiddleware) startBackgroundSampling() {
 	interval := time.Duration(am.config.SampleIntervalSeconds) * time.Second
 	ticker := time.NewTicker(interval)
 
 	go func() {
 		defer ticker.Stop()
-		defer log.Debug().Msg("memory protection background sampling stopped")
-
-		log.Debug().
-			Dur("interval", interval).
-			Msg("memory protection background sampling started")
+		// TODO this code might start running before the logger is setup, therefore we have a data race
+		//defer log.Debug().Msg("memory protection background sampling stopped")
+		//
+		//log.Debug().
+		//	Dur("interval", interval).
+		//	Msg("memory protection background sampling started")
 
 		for {
 			select {
 			case <-ticker.C:
-				if err := am.sampleMemory(); err != nil {
-					log.Warn().Err(err).Msg("background memory sampling failed")
-				}
+				am.sampleMemory()
 			case <-am.ctx.Done():
 				return
 			}
@@ -207,33 +242,36 @@ func (am *AdmissionMiddleware) startBackgroundSampling() {
 }
 
 // sampleMemory samples the current memory usage and updates the cached value
-func (am *AdmissionMiddleware) sampleMemory() error {
+func (am *MemoryAdmissionMiddleware) sampleMemory() {
 	defer func() {
 		if r := recover(); r != nil {
 			log.Warn().Interface("panic", r).Msg("memory sampling panicked")
 		}
 	}()
 
+	now := time.Now()
 	metrics.Read(am.metricsSamples)
-	newUsage := int64(am.metricsSamples[0].Value.Uint64())
-	am.lastMemoryUsage.Store(newUsage)
+	newUsage := am.metricsSamples[0].Value.Uint64()
+	am.lastMemorySampleInBytes.Store(newUsage)
+	am.timestampLastMemorySample.Store(&now)
 
 	// Update metrics gauge
 	if am.memoryLimit > 0 {
 		usagePercent := float64(newUsage) / float64(am.memoryLimit) * 100
 		MemoryUsageGauge.Set(usagePercent)
 	}
+}
 
-	return nil
+func (am *MemoryAdmissionMiddleware) getLastMemorySampleInBytes() uint64 {
+	return am.lastMemorySampleInBytes.Load()
 }
 
-// getCurrentMemoryUsage returns the cached memory usage in bytes
-func (am *AdmissionMiddleware) getCurrentMemoryUsage() int64 {
-	return am.lastMemoryUsage.Load()
+func (am *MemoryAdmissionMiddleware) getTimestampLastMemorySample() *time.Time {
+	return am.timestampLastMemorySample.Load()
 }
 
 // recordRejection records metrics for rejected requests
-func (am *AdmissionMiddleware) recordRejection(fullMethod string) {
+func (am *MemoryAdmissionMiddleware) recordRejection(fullMethod string) {
 	endpointType := "api"
 	if strings.HasPrefix(fullMethod, "/dispatch.v1.DispatchService") {
 		endpointType = "dispatch"