impl connection recovery

2026-08-15 17:31:35 +08:00
parent e63a34d53a
commit 7bd84af48d
33 changed files with 3858 additions and 200 deletions
@@ -0,0 +1,268 @@
+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:
+	}
+}