Day 2: Implement a Token-Bucket Interceptor on the gRPC Server — and Shed Excess Load Under CPU Saturation

Day 2: Implement a Token-Bucket Interceptor on the gRPC Server — and Shed Excess Load Under CPU Saturation

In yesterday's lesson, you defined the Protobuf contract for our core inventory allocation service and verified that changes to our messages would not break wire-format compatibility for existing clients. Today, we move from the schema layer to the runtime protection layer.

We are going to build a high-performance gRPC interceptor that implements a thread-safe token-bucket rate limiter. But we are not stopping at static limits. In production, static limits are a guessing game. If your server is running on a noisy neighbor VM or undergoing heavy garbage collection, its actual processing capacity drops. To survive, the server must dynamically adjust its rate limits and aggressively shed excess load before CPU saturation degrades service latency for everyone.

The Production Stakes: The Danger of Cooperative Degradation

When a distributed system is overloaded, it faces a choice: degrade cooperatively by processing some requests successfully while quickly rejecting others, or fail catastrophically by trying to process everything, resulting in timeout cascades, thread pool exhaustion, and eventual crash-loops.

This is exactly what happened during the infamous 2021 Roblox 73-hour outage. As backend systems struggled to recover, a thundering herd of client retries saturated the gateway layers. Because the internal services lacked aggressive, early load-shedding mechanisms, the CPU spent almost all its cycles on context switching, connection management, and half-completed requests that eventually timed out on the client side. The servers were doing massive amounts of work, but their useful throughput was zero.

To prevent this, we use Load Shedding. Instead of queueing requests until we run out of memory, we reject excess requests at the very edge of our service using a gRPC interceptor. By returning a gRPC RESOURCE_EXHAUSTED status code immediately, we consume virtually no CPU or memory for rejected requests, preserving our precious compute cycles for the requests we can successfully complete.

Component Architecture and Control Flow

Our load-shedding architecture consists of three main components running inside our Go-based gRPC server:

1. 1.The Dynamic Token Bucket: A thread-safe rate limiter that controls request flow. It has a maximum capacity (burst size) and a refill rate (sustained requests per second).

2. 2.The CPU Monitor: A background routine that periodically samples the system's CPU utilization.

3. 3.The gRPC Server Interceptor: A middleware layer that intercepts every incoming unary RPC, checks the Token Bucket, and either allows the request to proceed or terminates it instantly with a Codes.ResourceExhausted error.

When a request arrives, the interceptor must make a decision in sub-microsecond time. It queries the Token Bucket using a non-blocking Allow() check. If a token is available, it is consumed, and the request is dispatched to the handler. If no tokens are available, or if the CPU Monitor has flagged the system as saturated and throttled the bucket capacity, the request is rejected immediately.

Designing the Mechanism: Token Bucket vs. Alternatives

Why use a Token Bucket instead of other rate-limiting algorithms? Let's look at the trade-offs:

For high-performance gRPC services, the Token Bucket is the industry standard because we can implement it with zero-allocation, lock-free, or low-contention mutex operations, keeping our interceptor overhead under 50 nanoseconds per request.

Core Code Implementation Deep-Dive

Let's look at how the token-bucket state transitions are calculated. Instead of running a background timer to constantly add tokens to the bucket (which would waste CPU cycles), we calculate the token count lazily on every request based on the elapsed time since the last request.

// TokenBucket represents a thread-safe lazy-refill rate limiter.
type TokenBucket struct {
	mu         sync.Mutex
	capacity   float64
	tokens     float64
	refillRate float64 // Tokens per second
	lastRefill time.Time
}

// Allow checks if a single token can be consumed from the bucket.
func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastRefill).Seconds()
	tb.lastRefill = now

	// Refill tokens based on elapsed time
	tb.tokens = tb.tokens + (elapsed * tb.refillRate)
	if tb.tokens > tb.capacity {
		tb.tokens = tb.capacity
	}

	if tb.tokens >= 1.0 {
		tb.tokens -= 1.0
		return true
	}
	return false
}

The interceptor wraps this logic and integrates with our CPU monitor. If the CPU utilization exceeds our threshold (e.g., 80%), we dynamically scale down the refillRate and capacity to force incoming clients to back off.

// UnaryLoadShedder returns a gRPC interceptor that sheds load when the bucket is empty.
func UnaryLoadShedder(tb *TokenBucket, monitor *CPUMonitor) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// Dynamically adjust rate limits based on CPU stress
		if monitor.IsSaturated() {
			tb.ScaleLimits(0.2) // Reduce capacity to 20% under heavy CPU load
		} else {
			tb.ResetLimits()
		}

		if !tb.Allow() {
			return nil, status.Error(codes.ResourceExhausted, "server is saturated; load shedding active")
		}
		return handler(ctx, req)
	}
}

Capacity Math: The Overhead of Protection

When designing an interceptor, you must ensure that the protection mechanism does not cost more than the work it prevents.

Let's calculate our resource budget:

• Memory Footprint: A single TokenBucket struct in Go occupies 48 bytes. Even if we run a separate bucket per-client IP or per-API route (e.g., 10,000 active clients), the total memory overhead is less than 500 KB.

• CPU Overhead: A mutex lock/unlock cycle in Go takes approximately 15–25 nanoseconds under low contention. At 100,000 requests per second (RPS), the interceptor consumes roughly 2% of a single CPU core.

• The Load-Shedding Payback: Releasing a request via Codes.ResourceExhausted bypasses database queries, JSON parsing, and business logic serialization. If a normal inventory allocation request takes 5ms of CPU time, and a rejected request takes 1 microsecond, shedding load is 5,000 times cheaper than processing a failure under saturation.

The Before/After Delta

Before today, if our inventory service was hit with a traffic spike that exceeded its CPU capacity, requests would pile up in the OS TCP backlog and Go's goroutine queue. Latency would degrade exponentially, eventually causing clients to time out and retry, worsening the collapse.

After today, when traffic spikes beyond our configured threshold, the server will process exactly its maximum sustainable capacity (e.g., 100 RPS) with flat, healthy latencies (under 5ms), while immediately rejecting excess requests with a clean ResourceExhausted error in under 1 millisecond.

Assignment: Build a Dynamic Cooldown Controller

Your assignment is to extend the load-shedding interceptor to prevent rate-limit flapping. Currently, when CPU drops below 80%, we immediately restore full limits. If a massive queue of requests is waiting, the CPU will instantly spike back to 100%, causing the server to flap rapidly between normal and degraded states.

Tasks:

1. 1.Implement a Cooldown Period (e.g., 5 seconds) in the CPUMonitor. Once the server enters a saturated state, it must remain in the degraded rate-limiting state for at least 5 seconds after the CPU drops below the threshold.

2. 2.Write a unit test verifying that if CPU spikes to 90% and then immediately drops to 30%, the rate limits remain scaled down for the duration of the cooldown period.

Hint: Use a time.Time timestamp to track when the CPU last exceeded the threshold, and check if time.Since(lastSpike) < cooldownDuration before resetting limits.

Solution Walkthrough

To solve this, modify your CPUMonitor struct to track the last saturation timestamp:

type CPUMonitor struct {
	mu            sync.RWMutex
	threshold     float64
	currentCPU    float64
	lastSaturated time.Time
	cooldown      time.Duration
}

func (cm *CPUMonitor) IsSaturated() bool {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	isCurrentlyOver := cm.currentCPU >= cm.threshold
	if isCurrentlyOver {
		return true
	}

	// Check if we are still within the cooldown window
	return time.Since(cm.lastSaturated) < cm.cooldown
}

func (cm *CPUMonitor) UpdateCPU(val float64) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.currentCPU = val
	if val >= cm.threshold {
		cm.lastSaturated = time.Now()
	}
}

This simple hysteresis loop prevents system resonance and guarantees that the backend has fully stabilized before we open the floodgates to normal traffic levels again.

In tomorrow's lesson, we will build on this resilient foundation by writing a complete Protobuf schema for the CoreCart Inventory Service, generating production-ready Go stubs, and writing an advanced compatibility test suite to prove wire-format compatibility across schema versions.