steamcache2/steamcache/bandwidth.go

// 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
}