Algorithm Design & Real-Time System Optimization Questions

package adaptive

import (
	"sync"
	"time"
	"math"
)

// TenantLimiter controls QPS for one tenant.
type TenantLimiter struct {
	mu sync.Mutex

	// token bucket
	tokens        float64
	lastRefill    time.Time
	capacity      float64 // max tokens == allowed QPS
	refillPerSec  float64 // equals allowed QPS

	// latency measurement (EMA)
	latencyEMA    float64
	emaAlpha      float64

	// AIMD params
	additiveInc   float64 // QPS add when healthy
	multShrink    float64 // factor to multiply when high latency
	targetLatency time.Duration
	minQPS, maxQPS float64

	// stabilization
	lastAdjust time.Time
	adjustInterval time.Duration
}

// NewTenantLimiter initializes limiter with initialQPS.
func NewTenantLimiter(initialQPS float64, targetLatency time.Duration) *TenantLimiter {
	now := time.Now()
	return &TenantLimiter{
		tokens:        initialQPS,
		lastRefill:    now,
		capacity:      math.Max(1, initialQPS),
		refillPerSec:  initialQPS,
		latencyEMA:    0,
		emaAlpha:      0.2,
		additiveInc:   1.0,
		multShrink:    0.7,
		targetLatency: targetLatency,
		minQPS:        1.0,
		maxQPS:        10000,
		lastAdjust:    now,
		adjustInterval: 1 * time.Second,
	}
}

// Allow attempts to take one token. Non-blocking.
func (t *TenantLimiter) Allow() bool {
	t.mu.Lock()
	defer t.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(t.lastRefill).Seconds()
	t.tokens += elapsed * t.refillPerSec
	if t.tokens > t.capacity {
		t.tokens = t.capacity
	}
	t.lastRefill = now

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

// Report records observed latency and possibly adjusts allowed QPS.
// Call Report after each request completes.
func (t *TenantLimiter) Report(latency time.Duration) {
	t.mu.Lock()
	defer t.mu.Unlock()

	// update EMA latency
	l := float64(latency.Milliseconds())
	if t.latencyEMA == 0 {
		t.latencyEMA = l
	} else {
		t.latencyEMA = t.emaAlpha*l + (1-t.emaAlpha)*t.latencyEMA
	}

	// only adjust at most once per adjustInterval
	if time.Since(t.lastAdjust) < t.adjustInterval {
		return
	}
	t.lastAdjust = time.Now()

	curQPS := t.refillPerSec
	target := float64(t.targetLatency.Milliseconds())

	if t.latencyEMA <= target {
		// additive increase
		curQPS += t.additiveInc
	} else {
		// multiplicative decrease
		curQPS = math.Max(t.minQPS, curQPS * t.multShrink)
	}

	// clamp
	if curQPS > t.maxQPS {
		curQPS = t.maxQPS
	}
	t.refillPerSec = curQPS
	t.capacity = math.Max(1.0, curQPS) // bucket capacity ~ QPS
}