s1d3sw1ped/steamcache2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 19 Wiki Activity

Add concurrency limits and configuration options for SteamCache

Browse Source 
- Introduced maxConcurrentRequests and maxRequestsPerClient fields in the Config struct to manage request limits.
- Updated the SteamCache implementation to utilize these new configuration options for controlling concurrent requests.
- Enhanced the ServeHTTP method to enforce global and per-client rate limiting using semaphores.
- Modified the root command to accept new flags for configuring concurrency limits via command-line arguments.
- Updated tests to reflect changes in the SteamCache initialization and request handling logic.
This commit is contained in:
Justin Harms
2025-09-02 06:50:42 -05:00
parent 7fb1fcf21f
commit 9ca8fa4a5e
6 changed files with 203 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

184
steamcache/steamcache.go
View File

@@ -22,6 +22,7 @@ import (
"time"
"github.com/docker/go-units"
"golang.org/x/sync/semaphore"
)
// generateURLHash creates a SHA256 hash of the entire URL path for cache key
@@ -56,12 +57,17 @@ var hopByHopHeaders = map[string]struct{}{
"Server": {},
}
var (
// Request coalescing structures
coalescedRequests = make(map[string]*coalescedRequest)
coalescedRequestsMu sync.RWMutex
// Constants for limits
const (
defaultMaxConcurrentRequests = int64(200) // Max total concurrent requests
defaultMaxRequestsPerClient = int64(5) // Max concurrent requests per IP
)
type clientLimiter struct {
semaphore *semaphore.Weighted
lastSeen time.Time
}
type coalescedRequest struct {
responseChan chan *http.Response
errorChan chan error
@@ -107,25 +113,84 @@ func (cr *coalescedRequest) complete(resp *http.Response, err error) {
}
// getOrCreateCoalescedRequest gets an existing coalesced request or creates a new one
func getOrCreateCoalescedRequest(cacheKey string) (*coalescedRequest, bool) {
coalescedRequestsMu.Lock()
defer coalescedRequestsMu.Unlock()
func (sc *SteamCache) getOrCreateCoalescedRequest(cacheKey string) (*coalescedRequest, bool) {
sc.coalescedRequestsMu.Lock()
defer sc.coalescedRequestsMu.Unlock()
if cr, exists := coalescedRequests[cacheKey]; exists {
if cr, exists := sc.coalescedRequests[cacheKey]; exists {
cr.addWaiter()
return cr, false
}
cr := newCoalescedRequest()
coalescedRequests[cacheKey] = cr
sc.coalescedRequests[cacheKey] = cr
return cr, true
}
// removeCoalescedRequest removes a completed coalesced request
func removeCoalescedRequest(cacheKey string) {
coalescedRequestsMu.Lock()
defer coalescedRequestsMu.Unlock()
delete(coalescedRequests, cacheKey)
func (sc *SteamCache) removeCoalescedRequest(cacheKey string) {
sc.coalescedRequestsMu.Lock()
defer sc.coalescedRequestsMu.Unlock()
delete(sc.coalescedRequests, cacheKey)
}
// getClientIP extracts the client IP address from the request
func getClientIP(r *http.Request) string {
// Check for forwarded headers first (common in proxy setups)
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
// X-Forwarded-For can contain multiple IPs, take the first one
if idx := strings.Index(xff, ","); idx > 0 {
return strings.TrimSpace(xff[:idx])
}
return strings.TrimSpace(xff)
}
if xri := r.Header.Get("X-Real-IP"); xri != "" {
return strings.TrimSpace(xri)
}
// Fall back to RemoteAddr
if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
return host
}
return r.RemoteAddr
}
// getOrCreateClientLimiter gets or creates a rate limiter for a client IP
func (sc *SteamCache) getOrCreateClientLimiter(clientIP string) *clientLimiter {
sc.clientRequestsMu.Lock()
defer sc.clientRequestsMu.Unlock()
limiter, exists := sc.clientRequests[clientIP]
if !exists || time.Since(limiter.lastSeen) > 5*time.Minute {
// Create new limiter or refresh existing one
limiter = &clientLimiter{
semaphore: semaphore.NewWeighted(sc.maxRequestsPerClient),
lastSeen: time.Now(),
}
sc.clientRequests[clientIP] = limiter
} else {
limiter.lastSeen = time.Now()
}
return limiter
}
// cleanupOldClientLimiters removes old client limiters to prevent memory leaks
func (sc *SteamCache) cleanupOldClientLimiters() {
for {
time.Sleep(10 * time.Minute) // Clean up every 10 minutes
sc.clientRequestsMu.Lock()
now := time.Now()
for ip, limiter := range sc.clientRequests {
if now.Sub(limiter.lastSeen) > 30*time.Minute {
delete(sc.clientRequests, ip)
}
}
sc.clientRequestsMu.Unlock()
}
}
type SteamCache struct {
@@ -144,9 +209,22 @@ type SteamCache struct {
client *http.Client
cancel context.CancelFunc
wg sync.WaitGroup
// Request coalescing structures
coalescedRequests map[string]*coalescedRequest
coalescedRequestsMu sync.RWMutex
// Concurrency control
maxConcurrentRequests int64
requestSemaphore *semaphore.Weighted
// Per-client rate limiting
clientRequests map[string]*clientLimiter
clientRequestsMu sync.RWMutex
maxRequestsPerClient int64
}
func New(address string, memorySize string, diskSize string, diskPath, upstream, memoryGC, diskGC string) *SteamCache {
func New(address string, memorySize string, diskSize string, diskPath, upstream, memoryGC, diskGC string, maxConcurrentRequests int64, maxRequestsPerClient int64) *SteamCache {
memorysize, err := units.FromHumanSize(memorySize)
if err != nil {
panic(err)
@@ -238,6 +316,13 @@ func New(address string, memorySize string, diskSize string, diskPath, upstream,
ReadHeaderTimeout: 10 * time.Second, // New, for header attacks
MaxHeaderBytes: 1 << 20, // 1MB, optional
},
// Initialize concurrency control fields
coalescedRequests: make(map[string]*coalescedRequest),
maxConcurrentRequests: maxConcurrentRequests,
requestSemaphore: semaphore.NewWeighted(maxConcurrentRequests),
clientRequests: make(map[string]*clientLimiter),
maxRequestsPerClient: maxRequestsPerClient,
}
// Log GC algorithm configuration
@@ -272,6 +357,13 @@ func (sc *SteamCache) Run() {
ctx, cancel := context.WithCancel(context.Background())
sc.cancel = cancel
// Start cleanup goroutine for old client limiters
sc.wg.Add(1)
go func() {
defer sc.wg.Done()
sc.cleanupOldClientLimiters()
}()
sc.wg.Add(1)
go func() {
defer sc.wg.Done()
@@ -295,18 +387,49 @@ func (sc *SteamCache) Shutdown() {
}
func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Apply global concurrency limit first
if err := sc.requestSemaphore.Acquire(context.Background(), 1); err != nil {
logger.Logger.Warn().Str("client_ip", getClientIP(r)).Msg("Server at capacity, rejecting request")
http.Error(w, "Server busy, please try again later", http.StatusServiceUnavailable)
return
}
defer sc.requestSemaphore.Release(1)
// Apply per-client rate limiting
clientIP := getClientIP(r)
clientLimiter := sc.getOrCreateClientLimiter(clientIP)
if err := clientLimiter.semaphore.Acquire(context.Background(), 1); err != nil {
logger.Logger.Warn().
Str("client_ip", clientIP).
Int("max_per_client", int(sc.maxRequestsPerClient)).
Msg("Client exceeded concurrent request limit")
http.Error(w, "Too many concurrent requests from this client", http.StatusTooManyRequests)
return
}
defer clientLimiter.semaphore.Release(1)
if r.Method != http.MethodGet {
logger.Logger.Warn().Str("method", r.Method).Msg("Only GET method is supported")
logger.Logger.Warn().
Str("method", r.Method).
Str("client_ip", clientIP).
Msg("Only GET method is supported")
http.Error(w, "Only GET method is supported", http.StatusMethodNotAllowed)
return
}
if r.URL.Path == "/" {
logger.Logger.Debug().
Str("client_ip", clientIP).
Msg("Health check request")
w.WriteHeader(http.StatusOK) // this is used by steamcache2's upstream verification at startup
return
}
if r.URL.String() == "/lancache-heartbeat" {
logger.Logger.Debug().
Str("client_ip", clientIP).
Msg("LanCache heartbeat request")
w.Header().Add("X-LanCache-Processed-By", "SteamCache2")
w.WriteHeader(http.StatusNoContent)
w.Write(nil)
@@ -358,19 +481,21 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
logger.Logger.Info().
Str("key", cacheKey).
Str("host", r.Host).
Str("client_ip", clientIP).
Str("status", "HIT").
Dur("duration", time.Since(tstart)).
Msg("request")
Msg("cache request")
return
}
}
// Check for coalesced request (another client already downloading this)
coalescedReq, isNew := getOrCreateCoalescedRequest(cacheKey)
coalescedReq, isNew := sc.getOrCreateCoalescedRequest(cacheKey)
if !isNew {
// Wait for the existing download to complete
logger.Logger.Debug().
Str("key", cacheKey).
Str("client_ip", clientIP).
Int("waiting_clients", coalescedReq.waitingCount).
Msg("Joining coalesced request")
@@ -402,21 +527,26 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
logger.Logger.Info().
Str("key", cacheKey).
Str("host", r.Host).
Str("client_ip", clientIP).
Str("status", "HIT-COALESCED").
Dur("duration", time.Since(tstart)).
Msg("request")
Msg("cache request")
return
case err := <-coalescedReq.errorChan:
logger.Logger.Error().Err(err).Str("key", cacheKey).Msg("Coalesced request failed")
logger.Logger.Error().
Err(err).
Str("key", cacheKey).
Str("client_ip", clientIP).
Msg("Coalesced request failed")
http.Error(w, "Upstream request failed", http.StatusInternalServerError)
return
}
}
// Remove coalesced request when done
defer removeCoalescedRequest(cacheKey)
defer sc.removeCoalescedRequest(cacheKey)
var req *http.Request
if sc.upstream != "" { // if an upstream server is configured, proxy the request to the upstream server
@@ -590,25 +720,35 @@ func (sc *SteamCache) ServeHTTP(w http.ResponseWriter, r *http.Request) {
logger.Logger.Info().
Str("key", cacheKey).
Str("host", r.Host).
Str("client_ip", clientIP).
Str("status", "MISS").
Dur("duration", time.Since(tstart)).
Msg("request")
Msg("cache request")
return
}
if r.URL.Path == "/favicon.ico" {
logger.Logger.Debug().
Str("client_ip", clientIP).
Msg("Favicon request")
w.WriteHeader(http.StatusNoContent)
return
}
if r.URL.Path == "/robots.txt" {
logger.Logger.Debug().
Str("client_ip", clientIP).
Msg("Robots.txt request")
w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(http.StatusOK)
w.Write([]byte("User-agent: *\nDisallow: /\n"))
return
}
logger.Logger.Warn().Str("url", r.URL.String()).Msg("Not found")
logger.Logger.Warn().
Str("url", r.URL.String()).
Str("client_ip", clientIP).
Msg("Request not found")
http.Error(w, "Not found", http.StatusNotFound)
}

6
steamcache/steamcache_test.go
View File

@@ -14,7 +14,7 @@ func TestCaching(t *testing.T) {
os.WriteFile(filepath.Join(td, "key2"), []byte("value2"), 0644)
sc := New("localhost:8080", "1G", "1G", td, "", "lru", "lru")
sc := New("localhost:8080", "1G", "1G", td, "", "lru", "lru", 200, 5)
w, err := sc.vfs.Create("key", 5)
if err != nil {
@@ -85,7 +85,7 @@ func TestCaching(t *testing.T) {
}
func TestCacheMissAndHit(t *testing.T) {
sc := New("localhost:8080", "0", "1G", t.TempDir(), "", "lru", "lru")
sc := New("localhost:8080", "0", "1G", t.TempDir(), "", "lru", "lru", 200, 5)
key := "testkey"
value := []byte("testvalue")
@@ -166,7 +166,7 @@ func TestURLHashing(t *testing.T) {
// Removed hash calculation tests since we switched to lightweight validation
func TestSteamKeySharding(t *testing.T) {
sc := New("localhost:8080", "0", "1G", t.TempDir(), "", "lru", "lru")
sc := New("localhost:8080", "0", "1G", t.TempDir(), "", "lru", "lru", 200, 5)
// Test with a Steam-style key that should trigger sharding
steamKey := "steam/0016cfc5019b8baa6026aa1cce93e685d6e06c6e"