cache: Per-client fair-share bandwidth on table uplink

steamcache/bandwidth.go

@@ -0,0 +1,158 @@
+// steamcache/bandwidth.go
+// Per-client fair-share / absolute bandwidth shaping for table-tier uplink.
+// Distinct from max_requests_per_client concurrency (semaphores in ratelimit.go).
+package steamcache
+
+import (
+	"context"
+	"net/http"
+	"sync"
+
+	"golang.org/x/time/rate"
+)
+
+const bandwidthWriteChunk = 32 * 1024
+
+// clientBandwidthLimiter fair-shares uplinkBytesPerSec among active clients and/or
+// applies an absolute per-client bytes/sec cap. Both 0 disables shaping.
+type clientBandwidthLimiter struct {
+	uplinkBytesPerSec int64
+	absoluteCap       int64
+
+	mu       sync.Mutex
+	active   map[string]int // refcount of in-flight shaped responses per client IP
+	limiters map[string]*rate.Limiter
+}
+
+func newClientBandwidthLimiter(uplinkBytesPerSec, absoluteCap int64) *clientBandwidthLimiter {
+	if uplinkBytesPerSec < 0 {
+		uplinkBytesPerSec = 0
+	}
+	if absoluteCap < 0 {
+		absoluteCap = 0
+	}
+	return &clientBandwidthLimiter{
+		uplinkBytesPerSec: uplinkBytesPerSec,
+		absoluteCap:       absoluteCap,
+		active:            make(map[string]int),
+		limiters:          make(map[string]*rate.Limiter),
+	}
+}
+
+func (b *clientBandwidthLimiter) enabled() bool {
+	return b != nil && (b.uplinkBytesPerSec > 0 || b.absoluteCap > 0)
+}
+
+// acquire registers clientIP as actively downloading and returns its limiter
+// (nil if shaping disabled) plus a release func that must be deferred.
+func (b *clientBandwidthLimiter) acquire(clientIP string) (*rate.Limiter, func()) {
+	if !b.enabled() {
+		return nil, func() {}
+	}
+	b.mu.Lock()
+	b.active[clientIP]++
+	lim := b.ensureLimiterLocked(clientIP)
+	b.recomputeRatesLocked()
+	b.mu.Unlock()
+
+	var once sync.Once
+	release := func() {
+		once.Do(func() {
+			b.mu.Lock()
+			defer b.mu.Unlock()
+			if n := b.active[clientIP]; n <= 1 {
+				delete(b.active, clientIP)
+			} else {
+				b.active[clientIP] = n - 1
+			}
+			b.recomputeRatesLocked()
+		})
+	}
+	return lim, release
+}
+
+func (b *clientBandwidthLimiter) ensureLimiterLocked(clientIP string) *rate.Limiter {
+	if lim, ok := b.limiters[clientIP]; ok {
+		return lim
+	}
+	// Start with a placeholder; recomputeRatesLocked sets the real rate.
+	lim := rate.NewLimiter(rate.Limit(1), 1)
+	b.limiters[clientIP] = lim
+	return lim
+}
+
+func (b *clientBandwidthLimiter) recomputeRatesLocked() {
+	n := len(b.active)
+	if n == 0 {
+		return
+	}
+	var fair int64
+	if b.uplinkBytesPerSec > 0 {
+		fair = b.uplinkBytesPerSec / int64(n)
+		if fair < 1 {
+			fair = 1
+		}
+	}
+	for ip := range b.active {
+		r := fair
+		if b.absoluteCap > 0 {
+			if r == 0 || b.absoluteCap < r {
+				r = b.absoluteCap
+			}
+		}
+		if r < 1 {
+			r = 1
+		}
+		lim := b.ensureLimiterLocked(ip)
+		burst := int(r)
+		if burst < bandwidthWriteChunk {
+			burst = bandwidthWriteChunk
+		}
+		// Cap burst to avoid huge memory spikes on huge uplinks.
+		if burst > 4*bandwidthWriteChunk {
+			burst = 4 * bandwidthWriteChunk
+		}
+		lim.SetLimit(rate.Limit(r))
+		lim.SetBurst(burst)
+	}
+}
+
+// limitedResponseWriter rate-limits response body Write calls. Headers/WriteHeader
+// are unlimited. Implements http.ResponseWriter (+ optional Flusher/Hijacker passthrough
+// is intentionally omitted — SteamCache body path only needs Write).
+type limitedResponseWriter struct {
+	http.ResponseWriter
+	lim *rate.Limiter
+	ctx context.Context
+}
+
+func (w *limitedResponseWriter) Write(p []byte) (int, error) {
+	if w.lim == nil || len(p) == 0 {
+		return w.ResponseWriter.Write(p)
+	}
+	ctx := w.ctx
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	total := 0
+	for total < len(p) {
+		chunk := p[total:]
+		if len(chunk) > bandwidthWriteChunk {
+			chunk = chunk[:bandwidthWriteChunk]
+		}
+		if err := w.lim.WaitN(ctx, len(chunk)); err != nil {
+			return total, err
+		}
+		n, err := w.ResponseWriter.Write(chunk)
+		total += n
+		if err != nil {
+			return total, err
+		}
+	}
+	return total, nil
+}
+
+// Unwrap exposes the underlying ResponseWriter for http.ResponseController etc.
+func (w *limitedResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}