broad-announce/internal/concurrency/perip.go

// Package concurrency is the in-memory per-IP concurrency cap
// for ingestd (SPEC §22 layer 2). It is shared by:
//
//   - cmd/ingestd's HTTP server (via the httpserver.ConnState
//     callback that the http handler installs)
//   - cmd/ingestd's WebSocket ingest path (incremented on upgrade,
//     decremented on close)
//
// Why in-memory and not Redis: this is the DoS gate. A Redis
// round-trip on every connection would self-DoS. The cap is
// approximate (a few extra conns during a multi-replica failover
// is fine). A janitor goroutine prunes idle entries every minute
// so the map doesn't grow unbounded.
package concurrency

import (
	"sync"
	"sync/atomic"
	"time"
)

// PerIP is a per-source-IP concurrency cap. Each unique IP gets
// an atomic counter; Acquire returns false if the cap is hit.
// The cap is a hard cap; one process's limit is independent of
// other ingestd replicas (which is what we want — each replica
// guards its own TCP listener).
type PerIP struct {
	cap    int64
	mu     sync.Mutex
	conns  map[string]*atomic.Int64
	idle  map[string]time.Time // last released-at
	stop  chan struct{}
}

// NewPerIP creates a PerIP with the given cap. The cap is the
// maximum concurrent connections per IP. A non-positive cap
// disables the gate (Acquire always returns true).
//
// The janitor runs every minute and prunes IPs whose counter is
// zero and whose last-release is older than 5 min. This bounds
// the map size to "IPs that hit ingestd in the last 5 min",
// which is small.
func NewPerIP(cap int) *PerIP {
	p := &PerIP{
		cap:   int64(cap),
		conns: make(map[string]*atomic.Int64),
		idle:  make(map[string]time.Time),
		stop:  make(chan struct{}),
	}
	go p.janitor()
	return p
}

// Close stops the janitor. Safe to call multiple times.
func (p *PerIP) Close() {
	select {
	case <-p.stop:
	default:
		close(p.stop)
	}
}

// Acquire tries to increment the counter for ip. Returns true on
// success, false if the cap is hit. Always increments on success;
// the caller MUST call Release exactly once per successful
// Acquire when the connection closes.
func (p *PerIP) Acquire(ip string) bool {
	if p.cap <= 0 {
		return true // gate disabled
	}
	c := p.counter(ip)
	// Atomic add; if the result exceeds the cap, decrement back
	// and reject. This is the standard compare-and-swap pattern
	// for admission control.
	n := c.Add(1)
	if n > p.cap {
		c.Add(-1)
		return false
	}
	p.touchActive(ip)
	return true
}

// Release decrements the counter for ip. Safe to call without a
// matching Acquire (the counter clamps at 0; we use Add(-1) but
// the case where the result goes negative is treated as a no-op
// because the caller is misbehaving — log via Warn, not panic).
func (p *PerIP) Release(ip string) {
	if p.cap <= 0 {
		return
	}
	p.mu.Lock()
	c, ok := p.conns[ip]
	p.mu.Unlock()
	if !ok {
		return
	}
	n := c.Add(-1)
	if n < 0 {
		c.Add(1) // clamp
		return
	}
	p.touchIdle(ip)
}

// InUse returns the current count for ip (0 if not tracked).
// Useful for tests and the operator's /v1/debug/perip endpoint.
func (p *PerIP) InUse(ip string) int64 {
	p.mu.Lock()
	c, ok := p.conns[ip]
	p.mu.Unlock()
	if !ok {
		return 0
	}
	v := c.Load()
	if v < 0 {
		return 0
	}
	return v
}

func (p *PerIP) counter(ip string) *atomic.Int64 {
	p.mu.Lock()
	defer p.mu.Unlock()
	c, ok := p.conns[ip]
	if !ok {
		c = &atomic.Int64{}
		p.conns[ip] = c
	}
	return c
}

func (p *PerIP) touchActive(ip string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	delete(p.idle, ip)
}

func (p *PerIP) touchIdle(ip string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.idle[ip] = time.Now()
}

// janitor prunes idle entries whose counter is 0. The 5-min
// window is wide enough that a reconnecting source doesn't lose
// its slot but narrow enough that the map doesn't grow
// unbounded over weeks of traffic.
func (p *PerIP) janitor() {
	t := time.NewTicker(time.Minute)
	defer t.Stop()
	for {
		select {
		case <-p.stop:
			return
		case now := <-t.C:
			p.mu.Lock()
			for ip, last := range p.idle {
				if now.Sub(last) > 5*time.Minute {
					c := p.conns[ip]
					if c != nil && c.Load() == 0 {
						delete(p.conns, ip)
						delete(p.idle, ip)
					}
				}
			}
			p.mu.Unlock()
		}
	}
}