amadan
edge: an upgrade holds its slot until the handshake is answered (#358)
A Connection: Upgrade header is the client's to send, and the upgrade branch admitted it without a slot whether or not the instance upgraded. A slow ordinary request carrying the header therefore never counted toward the host's or the client's in-flight slots. The upgrade branch now holds a slot like any request, behind the same slotWriter: the slot goes back at the first final answer or at the reverse proxy's hijack on a 101, so an open WebSocket still does not count. slotWriter already forwards Hijack and Unwrap, so the proxy's ResponseController hijack works through it. Tests: a Connection: Upgrade request that never upgrades is counted and the same client is shed behind it; a real handshake through the stdlib proxy holds the slot while the instance considers it, upgrades, and leaves InFlight at 0 while bytes flow. TestRateUpgradeJudgedNotHeld is now TestRateUpgradeHeldUntilAnswered, matching spec §4.2.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 files changed,
+176
−38
internal/edge/edge.go+24 −15internal/edge/ratelimit.go+13 −15internal/edge/ratelimit_test.go+137 −6internal/ratelimit/limiter.go+2 −2
diff --git a/internal/edge/edge.go b/internal/edge/edge.go| index ed155629d..e2ccd7cbd 100644 |
| --- a/internal/edge/edge.go |
| +++ b/internal/edge/edge.go |
| @@ -1729,23 +1729,32 @@ func (e *Edge) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
| } |
| // A Connection: Upgrade request (a WebSocket handshake, typically GET) |
| - // must reach the activator/proxy with the real ResponseWriter, completely |
| - // untouched by caching: the stdlib reverse proxy's upgrade path type- |
| - // asserts the ResponseWriter to http.Hijacker to take over the raw |
| - // connection, and once upgraded, the open WebSocket is also what keeps a |
| - // hibernating instance's idle clock from firing (the activator's inflight |
| - // counting — see internal/activator/activator.go's actEntry.inflight doc |
| - // comment). So this path skips caching entirely and stays byte-for-byte |
| - // identical to pre-cache behavior, rather than relying on |
| - // recordingResponseWriter's Hijack forwarding below. |
| + // must reach the activator/proxy with a hijackable ResponseWriter, |
| + // completely untouched by caching: the stdlib reverse proxy's upgrade |
| + // path hijacks the ResponseWriter (through http.ResponseController) to |
| + // take over the raw connection, and once upgraded, the open WebSocket is |
| + // also what keeps a hibernating instance's idle clock from firing (the |
| + // activator's inflight counting — see internal/activator/activator.go's |
| + // actEntry.inflight doc comment). So this path skips caching and |
| + // compression entirely, rather than relying on recordingResponseWriter's |
| + // Hijack forwarding below; the only wrapper it gets is the rate limit's |
| + // slotWriter, which forwards Hijack and Unwrap. |
| if isUpgradeRequest(r) { |
| - // The rate limit judges the handshake and holds no slot: an open |
| - // WebSocket is not work the instance is doing for the busy count. A |
| - // refusal is written to the raw writer; an admission leaves it |
| - // untouched. |
| - if _, ok := e.rateAdmit(w, r, up, false); !ok { |
| + // The rate limit holds a slot until the instance answers the |
| + // handshake, as for any request (§4.2): the header is the client's |
| + // to send, so it must not buy a slow request out of the busy count. |
| + // slotWriter gives the slot back at the first final answer or at the |
| + // proxy's hijack on a 101, so an open WebSocket is not counted, and |
| + // it is itself an http.Hijacker (with Unwrap) for the proxy. A |
| + // refusal is written to the raw writer. |
| + release, ok := e.rateAdmit(w, r, up) |
| + if !ok { |
| return |
| } |
| + if release != nil { |
| + defer release() |
| + w = &slotWriter{ResponseWriter: w, release: release} |
| + } |
| if up.hibernating { |
| if e.act.ServeHTTP(w, r, up.route, up.proxy) { |
| return |
| @@ -1816,7 +1825,7 @@ func (e *Edge) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
| // the defer covers an answer that never wrote anything. A request that |
| // holds no slot (off, exempt, not an instance) keeps its writer as it |
| // was. |
| - release, admitted := e.rateAdmit(w, r, up, true) |
| + release, admitted := e.rateAdmit(w, r, up) |
| if !admitted { |
| return |
| } |
diff --git a/internal/edge/ratelimit.go b/internal/edge/ratelimit.go| index 72551febe..124db8f74 100644 |
| --- a/internal/edge/ratelimit.go |
| +++ b/internal/edge/ratelimit.go |
| @@ -108,12 +108,12 @@ func (e *Edge) parseRateSettings(rt registry.Route) ratelimit.Settings { |
| return s |
| } |
| -// rateAdmit judges r for up. holdSlot is false for an upgrade, which is |
| -// judged when its handshake arrives and never holds a slot. When ok is |
| -// false the 429 has been written and the caller returns. Otherwise release |
| -// is nil when the request holds no slot (not judged, or judged without one), |
| -// and when it is not, the caller must call it once the request is done; it |
| -// gives the slot back if the slot writer has not already. |
| +// rateAdmit judges r for up. When ok is false the 429 has been written and |
| +// the caller returns. Otherwise release is nil when the request holds no |
| +// slot (not judged, or admitted untracked), and when it is not, the caller |
| +// must call it once the request is done; it gives the slot back if the slot |
| +// writer has not already. An upgrade holds one like any request, until its |
| +// handshake is answered. |
| // |
| // Not judged: anything but an instance, a door caller (the box itself, on |
| // this box or on the box that backhauled it, hopDoorFrom), and a request |
| @@ -121,7 +121,7 @@ func (e *Edge) parseRateSettings(rt registry.Route) ratelimit.Settings { |
| // request comes from the hub, so keying on it would put every client in one |
| // bucket; the hub has the route row and the real client, and limits there |
| // (its route is viaTunnel). |
| -func (e *Edge) rateAdmit(w http.ResponseWriter, r *http.Request, up upstream, holdSlot bool) (release func(), ok bool) { |
| +func (e *Edge) rateAdmit(w http.ResponseWriter, r *http.Request, up upstream) (release func(), ok bool) { |
| if e.rateMode == RateOff || up.route.Kind != "instance" { |
| return nil, true |
| } |
| @@ -141,7 +141,7 @@ func (e *Edge) rateAdmit(w http.ResponseWriter, r *http.Request, up upstream, ho |
| observe := e.rateMode == RateObserve |
| // Keyed on the route's host, never the request's: an alias, a wildcard |
| // name and the primary are one instance with one set of buckets. |
| - t, d, retryAfter := e.rateLimiter.Admit(up.route.Host, eff, client, r.URL.Path, holdSlot, observe) |
| + t, d, retryAfter := e.rateLimiter.Admit(up.route.Host, eff, client, r.URL.Path, true, observe) |
| if d != ratelimit.Allow { |
| e.noteRateLimited(up.route, d, client.String(), observe) |
| if !observe { |
| @@ -149,9 +149,6 @@ func (e *Edge) rateAdmit(w http.ResponseWriter, r *http.Request, up upstream, ho |
| return nil, false |
| } |
| } |
| - if !holdSlot { |
| - return nil, true |
| - } |
| return t.Release, true |
| } |
| @@ -189,7 +186,7 @@ func writeRateLimited(w http.ResponseWriter, retryAfter time.Duration) { |
| } |
| // slotWriter gives a request's slot back when the instance starts its |
| -// answer: the first final WriteHeader, Write or Flush. The busy count |
| +// answer: the first final WriteHeader, Write or Flush, or a Hijack. The busy count |
| // measures work the instance is doing, not how fast a client downloads, so a |
| // long body or an event stream stops counting once its headers are out. An |
| // informational (1xx) header is not the answer and keeps the slot. |
| @@ -229,9 +226,10 @@ func (s *slotWriter) Flush() { |
| } |
| } |
| -// Hijack forwards like countingWriter's, so the writer stack below a |
| -// recordingResponseWriter keeps the shape it had. Nothing on the path this |
| -// writer serves upgrades; upgrades never get one. |
| +// Hijack gives the slot back and forwards like countingWriter's. On an |
| +// upgrade the stdlib reverse proxy hijacks once the instance has answered |
| +// 101, so this is where an upgrade's slot goes back; the open connection is |
| +// not work the instance is doing for the busy count. |
| func (s *slotWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { |
| hj, ok := s.ResponseWriter.(http.Hijacker) |
| if !ok { |
diff --git a/internal/edge/ratelimit_test.go b/internal/edge/ratelimit_test.go| index 3866307e0..da99aebea 100644 |
| --- a/internal/edge/ratelimit_test.go |
| +++ b/internal/edge/ratelimit_test.go |
| @@ -5,6 +5,7 @@ package edge |
| // ServeHTTP) against a real upstream; nothing calls rateAdmit directly. |
| import ( |
| + "bufio" |
| "context" |
| "fmt" |
| "io" |
| @@ -385,10 +386,12 @@ func TestRateSlotReleasedAtHeaders(t *testing.T) { |
| io.Copy(io.Discard, resp.Body) |
| } |
| -// An upgrade is judged when the handshake arrives and never holds a slot: |
| -// it keeps the raw, hijackable writer, and the open connection does not count |
| -// toward busy. Over the backstop it is refused before the activator. |
| -func TestRateUpgradeJudgedNotHeld(t *testing.T) { |
| +// An upgrade holds a slot like any request until the instance answers its |
| +// handshake (§4.2), through a writer that is still an http.Hijacker. The |
| +// activator answers each one at once, so the slot goes back and ten |
| +// upgrades leave nothing held. Over the backstop the next is refused before |
| +// the activator. |
| +func TestRateUpgradeHeldUntilAnswered(t *testing.T) { |
| var calls atomic.Int64 |
| var e *Edge |
| var heldDuring atomic.Int64 |
| @@ -419,8 +422,8 @@ func TestRateUpgradeJudgedNotHeld(t *testing.T) { |
| t.Fatalf("upgrade %d: status %d, want 101", i, rec.Code) |
| } |
| } |
| - if n := heldDuring.Load(); n != 0 { |
| - t.Fatalf("slots held during admitted upgrades summed to %d, want 0", n) |
| + if n := heldDuring.Load(); n != 10 { |
| + t.Fatalf("slots held during admitted upgrades summed to %d, want 10: each holds one until it is answered", n) |
| } |
| if n := e.rateLimiter.InFlight(rlHost); n != 0 { |
| t.Fatalf("InFlight = %d after ten upgrades, want 0", n) |
| @@ -433,6 +436,134 @@ func TestRateUpgradeJudgedNotHeld(t *testing.T) { |
| } |
| } |
| +// A Connection: Upgrade header is the client's to send, so it buys nothing: |
| +// a request carrying one that the instance never upgrades is counted while |
| +// it waits like any other, and the same client is shed behind it while the |
| +// instance is busy. |
| +func TestRateUpgradeHeaderWithoutUpgradeHoldsSlot(t *testing.T) { |
| + up := newBlockingUpstream(t) |
| + s := newRateScene(t, "enforce", nil, rlRoute(up.addr(), `{"busy_at":1,"client_share":1}`)) |
| + t.Cleanup(up.open) // before the front server closes on a failure |
| + |
| + upgradeHeaderGet := func(path string) <-chan rlResult { |
| + ch := make(chan rlResult, 1) |
| + go func() { |
| + req, err := http.NewRequest(http.MethodGet, s.front.URL+path, nil) |
| + if err != nil { |
| + ch <- rlResult{err: err} |
| + return |
| + } |
| + req.Host = rlHost |
| + req.Header.Set(rlPeerHeader, "203.0.113.1") |
| + req.Header.Set("Connection", "Upgrade") |
| + req.Header.Set("Upgrade", "websocket") |
| + resp, err := rlClient.Do(req) |
| + if err != nil { |
| + ch <- rlResult{err: err} |
| + return |
| + } |
| + defer resp.Body.Close() |
| + b, err := io.ReadAll(resp.Body) |
| + ch <- rlResult{status: resp.StatusCode, header: resp.Header, body: string(b), err: err} |
| + }() |
| + return ch |
| + } |
| + |
| + slow := upgradeHeaderGet("/slow") |
| + up.waitArrived(t, "the slow request") |
| + if n := s.e.rateLimiter.InFlight(rlHost); n != 1 { |
| + t.Fatalf("InFlight = %d while a Connection: Upgrade request that never upgrades waits on the instance, want 1", n) |
| + } |
| + want429(t, waitResult(t, upgradeHeaderGet("/again"), "the same client's second request"), "1", |
| + "the same client's second Connection: Upgrade request while busy") |
| + if n := up.hits.Load(); n != 1 { |
| + t.Fatalf("upstream hits = %d, want 1: a shed request must never reach the upstream", n) |
| + } |
| + |
| + up.open() |
| + wantOK(t, waitResult(t, slow, "the slow request"), "the slow request") |
| + if n := s.e.rateLimiter.InFlight(rlHost); n != 0 { |
| + t.Fatalf("InFlight = %d after the slow request was answered, want 0", n) |
| + } |
| +} |
| + |
| +// A real WebSocket handshake through ServeHTTP and the stdlib reverse proxy: |
| +// the slot is held while the instance considers the handshake, the proxy can |
| +// still hijack the connection, and once the 101 is out the slot is back |
| +// while bytes keep flowing both ways. With busy_at 1, the open socket does |
| +// not make the instance busy for anyone, its own client included. |
| +func TestRateWebSocketUpgradeReleasesOnceUpgraded(t *testing.T) { |
| + var e atomic.Pointer[Edge] |
| + var heldAtHandshake atomic.Int64 |
| + heldAtHandshake.Store(-1) |
| + ws := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| + if r.URL.Path != "/socket" { |
| + io.WriteString(w, "ok") |
| + return |
| + } |
| + heldAtHandshake.Store(int64(e.Load().rateLimiter.InFlight(rlHost))) |
| + conn, buf, err := w.(http.Hijacker).Hijack() |
| + if err != nil { |
| + return |
| + } |
| + defer conn.Close() |
| + buf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") |
| + if buf.Flush() != nil { |
| + return |
| + } |
| + for { |
| + line, err := buf.ReadString('\n') |
| + if err != nil { |
| + return |
| + } |
| + buf.WriteString("echo:" + line) |
| + if buf.Flush() != nil { |
| + return |
| + } |
| + } |
| + })) |
| + t.Cleanup(ws.Close) |
| + s := newRateScene(t, "enforce", nil, rlRoute(ws.Listener.Addr().String(), `{"busy_at":1,"client_share":1}`)) |
| + e.Store(s.e) |
| + |
| + c, err := net.Dial("tcp", s.front.Listener.Addr().String()) |
| + if err != nil { |
| + t.Fatal(err) |
| + } |
| + defer c.Close() |
| + if err := c.SetDeadline(time.Now().Add(10 * time.Second)); err != nil { |
| + t.Fatal(err) |
| + } |
| + fmt.Fprintf(c, "GET /socket HTTP/1.1\r\nHost: %s\r\n%s: 203.0.113.1\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\n", rlHost, rlPeerHeader) |
| + br := bufio.NewReader(c) |
| + resp, err := http.ReadResponse(br, nil) |
| + if err != nil { |
| + t.Fatalf("read handshake response: %v", err) |
| + } |
| + if resp.StatusCode != http.StatusSwitchingProtocols { |
| + t.Fatalf("handshake status = %d, want 101", resp.StatusCode) |
| + } |
| + if n := heldAtHandshake.Load(); n != 1 { |
| + t.Fatalf("InFlight = %d while the instance considered the handshake, want 1: an upgrade holds a slot until it is answered", n) |
| + } |
| + |
| + for _, msg := range []string{"one\n", "two\n"} { |
| + fmt.Fprint(c, msg) |
| + line, err := br.ReadString('\n') |
| + if err != nil { |
| + t.Fatalf("read echo of %q: %v", msg, err) |
| + } |
| + if line != "echo:"+msg { |
| + t.Fatalf("echo = %q, want %q", line, "echo:"+msg) |
| + } |
| + if n := s.e.rateLimiter.InFlight(rlHost); n != 0 { |
| + t.Fatalf("InFlight = %d while the socket is open, want 0: the slot goes back once upgraded", n) |
| + } |
| + } |
| + wantOK(t, s.get(t, "203.0.113.1", rlHost, "/other"), "the socket's own client during the open socket") |
| + wantOK(t, s.get(t, "203.0.113.2", rlHost, "/other"), "a second client during the open socket") |
| +} |
| + |
| // Review Focus 1: an alias is the primary's instance, so it spends the |
| // primary's bucket. |
| func TestRateAliasUsesPrimaryHost(t *testing.T) { |
diff --git a/internal/ratelimit/limiter.go b/internal/ratelimit/limiter.go| index 9510d9358..dacca5d42 100644 |
| --- a/internal/ratelimit/limiter.go |
| +++ b/internal/ratelimit/limiter.go |
| @@ -97,7 +97,7 @@ type Ticket struct { |
| done bool // guarded by l.mu |
| } |
| -// noSlot is the Ticket for a request let through without a slot: an upgrade |
| +// noSlot is the Ticket for a request let through without a slot: one |
| // judged with holdSlot false, or a client admitted untracked. |
| var noSlot = &Ticket{} |
| @@ -125,7 +125,7 @@ func (t *Ticket) Release() { |
| l.mu.Unlock() |
| } |
| -// Admit decides one request. holdSlot false (upgrades) judges without taking a slot. |
| +// Admit decides one request. holdSlot false judges without taking a slot. |
| // On Shed/OverCeiling the returned Ticket is nil unless observe is true, in which |
| // case the request is let through and the Ticket holds a slot like Allow. |
| // |