grant: replace the rate check with a token bucket

The live throughput test found it: a 3-second run delivered 104 packets and
stopped after 50 milliseconds.

The rate check exempted the first 50 ms entirely, meaning to be lenient at
startup. The effect was the opposite. A sender could dump an unbounded burst
into that free window, and the instant the check switched on it compared those
bytes against 50 ms worth of allowance and refused everything until real time
caught up. Every short test passed — downtrain sends 50 packets, big_send seven
— and every sustained send died about fifty milliseconds in.

A token bucket (allowance = burst + rate x elapsed) has no such cliff; it is
smooth from t=0. The burst is 100 ms of the allowed rate, floored at one
ordinary datagram so a single packet is never refused outright. The floor is
deliberately one datagram: at 8 kbps a 64 KB floor would be sixty-four seconds'
worth, which is precisely the instant dump the ceiling exists to prevent. The
existing rate test caught that when I first tried it, and it was right.

Second half of the same bug: callers treated any refusal as terminal. TryAllow
now says why, so a sender can pace through a transient "too fast just now" and
still stop dead on a spent budget or an expired grant.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

server/internal/dataplane/throughput.go

@@ -138,19 +138,24 @@ func (s *Server) DownThroughput(
 	start := time.Now()
 	next := start
 
-	res.LimitedBy = "duration"
 	var seq uint32
 	for time.Now().Before(deadline) {
-		if !g.Allow(sizeBytes) {
-			// Distinguishing these two matters: a run stopped by the byte budget has not been
-			// running long enough for its rate to mean anything.
-			if g.Sent() >= g.MaxBytes {
-				res.LimitedBy = "budget"
-			} else {
+		ok, why := g.TryAllow(sizeBytes)
+		if !ok {
+			if why == session.RefusalRate {
+				// Transient: the bucket is momentarily empty. Wait for the next slot and carry
+				// on. Ending the run here would report a rate measured over a fraction of a
+				// second, which is worse than reporting no rate at all.
 				res.LimitedBy = "rate"
+				time.Sleep(time.Duration(perPacketNs))
+				continue
 			}
+			// Terminal: the budget is spent, or the grant expired.
+			res.LimitedBy = why
 			break
 		}
+		// Reaching here means the run is progressing normally; the clock will end it.
+		res.LimitedBy = "duration"
 		binary.BigEndian.PutUint32(payload[0:4], seq)
 		binary.BigEndian.PutUint64(payload[4:12], uint64(time.Since(s.start).Nanoseconds()))
 		if err := s.sendErr(conn, target, sess, TypeThroughputData, seq, payload); err != nil {

server/internal/session/grant.go

@@ -74,22 +74,64 @@ func (s *Session) NewGrant(actionID string, wantBytes int64, wantKbps int, lim G
 // enforces the byte ceiling, the expiry, and the average rate (by refusing early sends rather
 // than sleeping, so callers stay in control of pacing).
 func (g *Grant) Allow(n int) bool {
+	ok, _ := g.TryAllow(n)
+	return ok
+}
+
+// Refusal reasons from TryAllow. The distinction is not cosmetic: "too fast just now" is
+// transient and a caller should pace and carry on, while "budget" and "expired" are terminal and
+// a caller that keeps trying is only wasting its own run.
+const (
+	RefusalNone    = ""
+	RefusalBudget  = "budget"
+	RefusalExpired = "expired"
+	RefusalRate    = "rate"
+)
+
+// TryAllow reports whether n more bytes may be sent now, consuming the budget when they may, and
+// says why not when they may not.
+//
+// The rate limit is a token bucket: allowance = burst + rate x elapsed. An earlier version
+// exempted the first 50 ms from the check entirely, meaning to be lenient at startup. The effect
+// was the opposite - a sender could dump an unbounded burst into that window, and the moment the
+// check switched on it compared those bytes against 50 ms worth of allowance and refused
+// everything until real time caught up. A sustained send died about fifty milliseconds in, having
+// looked perfectly fine in every short test. A bucket has no such cliff: it is smooth from t=0.
+func (g *Grant) TryAllow(n int) (bool, string) {
 	g.mu.Lock()
 	defer g.mu.Unlock()
 	if time.Now().After(g.ExpiresAt) {
-		return false
+		return false, RefusalExpired
 	}
 	if g.sentBytes+int64(n) > g.MaxBytes {
-		return false
+		return false, RefusalBudget
 	}
-	// Average-rate check: bytes allowed so far = kbps/8 * elapsed_seconds.
+	// kbps -> bytes/s is kbps*1000/8 = kbps*125.
+	bytesPerSec := float64(g.MaxKbps) * 125
 	elapsed := time.Since(g.started).Seconds()
-	allowed := float64(g.MaxKbps) * 125 * elapsed // kbps -> bytes/s is kbps*1000/8 = kbps*125
-	if elapsed > 0.05 && float64(g.sentBytes+int64(n)) > allowed {
-		return false
+	allowed := burstBytes(bytesPerSec) + bytesPerSec*elapsed
+	if float64(g.sentBytes+int64(n)) > allowed {
+		return false, RefusalRate
 	}
 	g.sentBytes += int64(n)
-	return true
+	return true, RefusalNone
+}
+
+// burstBytes is the bucket's depth: 100 ms of the allowed rate, floored at a single ordinary
+// datagram.
+//
+// The floor exists only so that one packet is never refused outright by a very slow grant — it is
+// deliberately one datagram and not more. A generous floor would undo the rate ceiling at low
+// rates: at 8 kbps a 64 KB burst is sixty-four seconds' worth, which is exactly the instant dump
+// the ceiling is there to prevent. One datagram is 1.5 seconds' worth at that rate and nothing at
+// any realistic one.
+func burstBytes(bytesPerSec float64) float64 {
+	const oneDatagram = 1500
+	b := bytesPerSec * 0.1
+	if b < oneDatagram {
+		b = oneDatagram
+	}
+	return b
 }
 
 // Sent returns how many bytes this grant has consumed.

server/internal/session/grant_test.go

@@ -92,3 +92,67 @@ func TestGrantEnforcesRate(t *testing.T) {
 		t.Fatalf("rate limit let %d bytes through in ~60ms at 8kbps", sent)
 	}
 }
+
+// The bug this pins: the rate check used to exempt the first 50 ms entirely, so a sender could
+// dump an unbounded burst into that window and then be refused for as long as it took real time
+// to catch up. Every short test passed; a sustained send died about fifty milliseconds in. A
+// token bucket has no such cliff, and the property that matters is that a sender pacing *at* the
+// allowed rate is never refused for long.
+func TestSustainedSendAtTheAllowedRateIsNotCutOff(t *testing.T) {
+	s := sessionWithSource(t)
+	const kbps = 8000   // 1 MB/s
+	const packet = 1200 // bytes
+	g := s.NewGrant("a1", 8<<20, kbps, DefaultGrantLimits)
+
+	// Pace at the allowed rate for a short run and count how much got through. A correct
+	// limiter passes essentially all of it; the old one stopped almost immediately.
+	perPacket := time.Duration(float64(packet) / (float64(kbps) * 125) * float64(time.Second))
+	deadline := time.Now().Add(300 * time.Millisecond)
+	sent, refusals := 0, 0
+	for time.Now().Before(deadline) {
+		if ok, why := g.TryAllow(packet); ok {
+			sent += packet
+		} else if why == RefusalRate {
+			refusals++
+		} else {
+			t.Fatalf("unexpected terminal refusal %q after %d bytes", why, sent)
+		}
+		time.Sleep(perPacket)
+	}
+
+	// 300 ms at 1 MB/s is ~300 KB. Allow generous slack for scheduler granularity, but a run
+	// that delivered only a few packets means the limiter cut it off.
+	if sent < 100_000 {
+		t.Fatalf("a sender pacing at the allowed rate got only %d bytes through in 300ms "+
+			"(%d rate refusals) — the limiter is cutting off sustained sends", sent, refusals)
+	}
+}
+
+// The other half: a rate refusal must be distinguishable from a spent budget, because one is
+// transient and one is terminal, and a caller that cannot tell them apart either gives up early
+// or spins forever.
+func TestRefusalReasonsAreDistinguishable(t *testing.T) {
+	s := sessionWithSource(t)
+
+	// Budget: tiny ceiling, plenty of rate.
+	g := s.NewGrant("a1", 1000, 100_000, DefaultGrantLimits)
+	for i := 0; i < 20; i++ {
+		g.TryAllow(100)
+	}
+	if ok, why := g.TryAllow(100); ok || why != RefusalBudget {
+		t.Errorf("spent budget reported as ok=%v why=%q, want %q", ok, why, RefusalBudget)
+	}
+
+	// Rate: huge ceiling, minimal rate, so only the bucket can refuse.
+	g2 := s.NewGrant("a2", 1<<20, 8, DefaultGrantLimits)
+	sawRate := false
+	for i := 0; i < 100; i++ {
+		if ok, why := g2.TryAllow(1000); !ok && why == RefusalRate {
+			sawRate = true
+			break
+		}
+	}
+	if !sawRate {
+		t.Error("a sender far above the rate ceiling never got a rate refusal")
+	}
+}