redapricot/client/shaper.go

package client

import (
	"sync"
	"time"
)

// minShaperWait floors the dispatcher's sleep so floating-point dust in the
// token arithmetic cannot spin it.
const minShaperWait = time.Millisecond

// Shaper caps the aggregate rate at which the client writes DATA to the hub and
// divides that budget across streams.
//
// The credit windows of PROTOCOL.md §7.3 bound how many bytes may be *in flight*
// per stream; they say nothing about bytes per *second*. That is the gap this
// fills. On a residential uplink one player loading chunks will otherwise
// saturate the line and push every other player's keepalive past its timeout.
//
// Two mechanisms are layered:
//
//   - A token bucket sets the long-run rate and the size of the burst that may
//     be spent after an idle period.
//   - Start-time fair queueing decides who spends those tokens. A global virtual
//     clock advances with each grant; every stream remembers the virtual time at
//     which its last request finished. A request is stamped
//     max(share.vfinish, vclock) and the lowest stamp is served first, so a
//     stream that keeps sending pushes its own stamp further out and yields to
//     quieter streams. The clamp to vclock is what keeps bursts cheap: a stream
//     returning from idle is pulled back to the head of the clock, so it cannot
//     hoard credit while it was idle, but it is not punished for the idleness
//     either. One stream alone gets the whole rate.
//
// A nil *Shaper means "no limit"; every method short-circuits, so call sites do
// not branch.
type Shaper struct {
	rate  float64 // bytes per second
	burst float64 // token bucket capacity, bytes
	chunk int     // how much a caller should request at a time

	mu      sync.Mutex
	tokens  float64
	last    time.Time
	vclock  float64      // virtual time, in bytes of service granted
	waiting []*shaperReq // unordered; the dispatcher scans for the lowest vstart

	wake chan struct{} // cap 1, non-blocking: nudges the dispatcher
	done chan struct{}
	once sync.Once
}

// shaperShare is one stream's position in the fair queue. It lives on the
// Stream and dies with it; a fresh share starts at zero and is clamped up to
// the current virtual clock on its first request.
type shaperShare struct{ vfinish float64 }

// shaperReq is one pending Acquire. granted and membership in Shaper.waiting
// are both guarded by Shaper.mu.
type shaperReq struct {
	n       int
	vstart  float64
	grant   chan struct{}
	granted bool
}

// NewShaper builds a shaper for the given rate. A non-positive rate returns nil,
// which every method treats as "unlimited".
func NewShaper(bytesPerSec int64) *Shaper {
	if bytesPerSec <= 0 {
		return nil
	}
	rate := float64(bytesPerSec)

	burst := rate * ShaperBurstSeconds
	// The floor is a correctness constraint, not a preference: a request larger
	// than the bucket could never be afforded and would park forever.
	if burst < MinShaperBurst {
		burst = MinShaperBurst
	}
	if burst > MaxShaperBurst {
		burst = MaxShaperBurst
	}

	chunk := int(rate * ShaperSliceSeconds)
	if chunk < MinShaperChunk {
		chunk = MinShaperChunk
	}
	if chunk > DataChunkSize {
		chunk = DataChunkSize
	}

	sh := &Shaper{
		rate:   rate,
		burst:  burst,
		chunk:  chunk,
		tokens: burst,
		last:   time.Now(),
		wake:   make(chan struct{}, 1),
		done:   make(chan struct{}),
	}
	go sh.dispatch()
	return sh
}

// chunkSize is how many bytes a sender should offer per request. It is sized to
// ShaperSliceSeconds of transmission so no stream holds the link for long before
// the scheduler can switch: at 1 Mbps a full 32 KiB chunk takes ~256 ms, which is
// enough dead air to drag other players towards a keepalive timeout.
func (sh *Shaper) chunkSize() int {
	if sh == nil {
		return DataChunkSize
	}
	return sh.chunk
}

// Acquire blocks until n bytes of bandwidth budget are available for the stream
// owning share. It returns false only when cancel fires first, in which case
// nothing was charged.
//
// cancel is the stream's done channel: a stream torn down while parked here must
// not keep a goroutine (and its Stream) alive waiting for tokens it will never
// use.
func (sh *Shaper) Acquire(share *shaperShare, n int, cancel <-chan struct{}) bool {
	if sh == nil || n <= 0 {
		return true
	}
	req := &shaperReq{n: n, grant: make(chan struct{})}

	sh.mu.Lock()
	// Stamp the request and reserve this stream's slot in virtual time up front,
	// so a stream cannot queue many requests at the same cheap stamp.
	req.vstart = share.vfinish
	if req.vstart < sh.vclock {
		req.vstart = sh.vclock
	}
	share.vfinish = req.vstart + float64(n)
	sh.waiting = append(sh.waiting, req)
	sh.mu.Unlock()
	sh.nudge()

	select {
	case <-req.grant:
		return true
	case <-sh.done:
		// Shaping stopped: let live traffic through rather than stalling it.
		sh.mu.Lock()
		sh.removeLocked(req)
		sh.mu.Unlock()
		return true
	case <-cancel:
		sh.mu.Lock()
		granted := req.granted
		if !granted {
			sh.removeLocked(req)
		}
		sh.mu.Unlock()
		return granted
	}
}

// Stop shuts the dispatcher down and releases everyone parked in Acquire.
func (sh *Shaper) Stop() {
	if sh == nil {
		return
	}
	sh.once.Do(func() { close(sh.done) })
}

// dispatch is the single goroutine that hands out tokens. It sleeps exactly as
// long as the next waiter needs rather than polling on a fixed tick, so an idle
// shaper costs nothing.
func (sh *Shaper) dispatch() {
	for {
		wait := sh.grantReady()
		var tick <-chan time.Time
		var timer *time.Timer
		if wait > 0 {
			timer = time.NewTimer(wait)
			tick = timer.C
		}
		select {
		case <-tick:
		case <-sh.wake:
		case <-sh.done:
			if timer != nil {
				timer.Stop()
			}
			return
		}
		if timer != nil {
			timer.Stop()
		}
	}
}

// grantReady refills the bucket and grants every waiter it can afford, lowest
// virtual start time first. It returns how long until the next waiter becomes
// affordable, or 0 when nothing is pending.
func (sh *Shaper) grantReady() time.Duration {
	sh.mu.Lock()
	defer sh.mu.Unlock()

	now := time.Now()
	if elapsed := now.Sub(sh.last); elapsed > 0 {
		sh.tokens += sh.rate * elapsed.Seconds()
		if sh.tokens > sh.burst {
			sh.tokens = sh.burst
		}
		sh.last = now
	}

	for {
		req := sh.headLocked()
		if req == nil {
			return 0
		}
		// Callers stay under chunkSize, which NewShaper keeps below the bucket.
		// Should a future caller not, wait for a full bucket rather than for a
		// token count that can never be reached, and let the balance go negative:
		// the debt is repaid by the next refill, so the long-run rate still holds.
		need := min(float64(req.n), sh.burst)
		if need > sh.tokens {
			wait := time.Duration((need - sh.tokens) / sh.rate * float64(time.Second))
			if wait < minShaperWait {
				wait = minShaperWait
			}
			return wait
		}
		sh.tokens -= float64(req.n)
		// The clock follows the request being served, never runs ahead of it.
		if req.vstart > sh.vclock {
			sh.vclock = req.vstart
		}
		req.granted = true
		sh.removeLocked(req)
		close(req.grant)
	}
}

// headLocked returns the pending request with the lowest virtual start time.
// A linear scan is deliberate: the queue holds at most one entry per live
// stream (tens, not thousands), so a heap would cost more in complexity than it
// saves in comparisons.
func (sh *Shaper) headLocked() *shaperReq {
	var best *shaperReq
	for _, w := range sh.waiting {
		if best == nil || w.vstart < best.vstart {
			best = w
		}
	}
	return best
}

func (sh *Shaper) removeLocked(req *shaperReq) {
	for i, w := range sh.waiting {
		if w == req {
			sh.waiting = append(sh.waiting[:i], sh.waiting[i+1:]...)
			return
		}
	}
}

func (sh *Shaper) nudge() {
	select {
	case sh.wake <- struct{}{}:
	default:
	}
}