chore: improvements and fix building of dispatch middleware

Author: miparnisari

internal/middleware/memoryprotection/memoryprotection.go

@@ -22,15 +22,15 @@ 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",
 	})
@@ -60,55 +60,84 @@ 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)
 }
 
-// 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
+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
+}
+
+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,
+	am.memoryLimit = memoryLimit
+	am.metricsSamples = []metrics.Sample{
+		{Name: "/memory/classes/heap/objects:bytes"},
 	}
 
 	// Initialize with current memory usage
-	if err := am.sampleMemory(); err != nil {
-		log.Warn().Err(err).Msg("failed to get initial memory sample")
-	}
+	am.sampleMemory()
 
 	// Start background sampling with context
 	am.startBackgroundSampling()
@@ -122,9 +151,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 +168,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 +187,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 +210,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 +235,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"