From 9c0b23b21aba4c83eb5f35700b358d3942dcf490 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Tue, 30 Jun 2026 12:29:44 +0530 Subject: [PATCH 01/11] drpcwire: add KindWindowUpdate flow-control frame Define the KindWindowUpdate control frame that carries a flow-control credit grant: a varint byte delta in the body, for the stream named by the frame's stream id. Add WindowUpdateFrame/ParseWindowUpdate helpers to build and read it. Stream id 0 is reserved for a possible future connection-level window and is unused in v1. The control bit marks the frame as out-of-band signaling, so it is emitted without blocking on data backpressure. It is not relied on for backward compatibility: an unknown control frame is only dropped after packet assembly, so it is not safely ignorable on an active stream. Instead, grants are only ever sent once flow control is active cluster-wide, so a peer that predates flow control never receives one. ParseWindowUpdate enforces the whole v1 wire contract in one place, since these frames are intercepted before packet assembly and its message-id checks: it accepts only a self-contained control frame (Control and Done set) naming a real stream (id != 0) with a positive delta and no trailing bytes. A non-conforming frame yields ok=false and is dropped by the caller, so a reserved stream-0 update from a future peer is ignored rather than mishandled. This is dormant scaffolding for the flow-control work: nothing emits or acts on the frame yet. Co-Authored-By: roachdev-claude --- drpcwire/packet.go | 9 ++++ drpcwire/packet_string.go | 5 ++- drpcwire/window_update.go | 45 ++++++++++++++++++++ drpcwire/window_update_test.go | 75 ++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 drpcwire/window_update.go create mode 100644 drpcwire/window_update_test.go diff --git a/drpcwire/packet.go b/drpcwire/packet.go index 447eb40e..5b490947 100644 --- a/drpcwire/packet.go +++ b/drpcwire/packet.go @@ -36,6 +36,15 @@ const ( // KindInvokeMetadata includes metadata about the next Invoke packet. KindInvokeMetadata Kind = 7 + + // KindWindowUpdate carries a flow-control credit grant: a varint byte delta + // in the body for the stream named by the frame's stream id. Stream id 0 is + // reserved for a possible future connection-level window and is unused in v1. + // The control bit marks it as out-of-band signaling (so it is emitted without + // blocking on data backpressure); it is not relied on for backward + // compatibility — grants are only ever sent once flow control is active + // cluster-wide, so a peer that predates flow control never receives one. + KindWindowUpdate Kind = 8 ) // diff --git a/drpcwire/packet_string.go b/drpcwire/packet_string.go index ceb317a6..98ee6ae1 100644 --- a/drpcwire/packet_string.go +++ b/drpcwire/packet_string.go @@ -15,11 +15,12 @@ func _() { _ = x[KindClose-5] _ = x[KindCloseSend-6] _ = x[KindInvokeMetadata-7] + _ = x[KindWindowUpdate-8] } -const _Kind_name = "InvokeMessageErrorCancelCloseCloseSendInvokeMetadata" +const _Kind_name = "InvokeMessageErrorCancelCloseCloseSendInvokeMetadataWindowUpdate" -var _Kind_index = [...]uint8{0, 6, 13, 18, 24, 29, 38, 52} +var _Kind_index = [...]uint8{0, 6, 13, 18, 24, 29, 38, 52, 64} func (i Kind) String() string { i -= 1 diff --git a/drpcwire/window_update.go b/drpcwire/window_update.go new file mode 100644 index 00000000..43ce6249 --- /dev/null +++ b/drpcwire/window_update.go @@ -0,0 +1,45 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcwire + +// WindowUpdateFrame builds a flow-control credit grant for the given stream id +// and byte delta. Callers must pass a real stream id (streamID > 0; stream id 0 +// is reserved for a possible future connection-level window and is unused in +// v1) and a positive delta. These are exactly the conditions ParseWindowUpdate +// enforces: a frame built with streamID == 0 or delta == 0 is well-formed but +// will be rejected on parse, so this helper does not validate them itself (its +// callers, the grant-emit path, are controlled and always satisfy them). The +// frame is marked Control (out-of-band signaling, emitted without blocking on +// data backpressure) and Done (a single self-contained frame). +func WindowUpdateFrame(streamID, delta uint64) Frame { + return Frame{ + Data: AppendVarint(nil, delta), + ID: ID{Stream: streamID}, + Kind: KindWindowUpdate, + Done: true, + Control: true, + } +} + +// ParseWindowUpdate extracts the stream id and credit delta from a window update +// frame, enforcing the whole v1 wire contract in one place: the frame must be a +// self-contained control frame (Control and Done set), name a real stream +// (id != 0; stream 0 is reserved for a future connection-level window), and +// carry a positive delta with no trailing bytes. ok is false for any frame that +// does not conform (a non-KindWindowUpdate frame or a malformed/empty payload +// included). A non-conforming frame is meant to be dropped by the caller, not +// acted on or treated as fatal; this also means a reserved stream-0 update from +// a future peer is ignored rather than mishandled. The message id is +// deliberately unconstrained: window updates are intercepted before packet +// assembly, so message-id monotonicity does not apply to them. +func ParseWindowUpdate(fr Frame) (streamID, delta uint64, ok bool) { + if fr.Kind != KindWindowUpdate || !fr.Control || !fr.Done || fr.ID.Stream == 0 { + return 0, 0, false + } + rem, d, parsed, err := ReadVarint(fr.Data) + if !parsed || err != nil || len(rem) != 0 || d == 0 { + return 0, 0, false + } + return fr.ID.Stream, d, true +} diff --git a/drpcwire/window_update_test.go b/drpcwire/window_update_test.go new file mode 100644 index 00000000..8ff45586 --- /dev/null +++ b/drpcwire/window_update_test.go @@ -0,0 +1,75 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcwire + +import ( + "testing" + + "github.com/zeebo/assert" +) + +func TestKindWindowUpdateString(t *testing.T) { + assert.Equal(t, KindWindowUpdate.String(), "WindowUpdate") +} + +func TestWindowUpdateFrameRoundTrip(t *testing.T) { + for _, tc := range []struct { + stream uint64 + delta uint64 + }{ + {stream: 7, delta: 128 << 10}, // per-stream + {stream: 123456, delta: 1 << 40}, // large delta + } { + fr := WindowUpdateFrame(tc.stream, tc.delta) + + // It is out-of-band signaling: a control frame (control bit set, so it is + // emitted without blocking on data backpressure) and a single + // self-contained frame (Done set). + assert.That(t, fr.Control) + assert.That(t, fr.Done) + assert.Equal(t, fr.Kind, KindWindowUpdate) + + rem, got, ok, err := ParseFrame(AppendFrame(nil, fr)) + assert.NoError(t, err) + assert.That(t, ok) + assert.Equal(t, len(rem), 0) + + sid, delta, ok := ParseWindowUpdate(got) + assert.That(t, ok) + assert.Equal(t, sid, tc.stream) + assert.Equal(t, delta, tc.delta) + } +} + +func TestParseWindowUpdateRejectsNonconforming(t *testing.T) { + // A conforming frame parses. + if _, _, ok := ParseWindowUpdate(WindowUpdateFrame(7, 128)); !ok { + t.Fatal("conforming window update did not parse") + } + + notControl := WindowUpdateFrame(7, 128) + notControl.Control = false + notDone := WindowUpdateFrame(7, 128) + notDone.Done = false + trailing := WindowUpdateFrame(7, 128) + trailing.Data = append(append([]byte(nil), trailing.Data...), 0xff) + empty := WindowUpdateFrame(7, 128) + empty.Data = nil + + // The v1 wire contract: self-contained control frame (Control+Done), real + // stream (id != 0), positive delta, no trailing bytes. + for name, fr := range map[string]Frame{ + "wrong kind": {Kind: KindMessage, Control: true, Done: true, ID: ID{Stream: 7}, Data: AppendVarint(nil, 1)}, + "not control": notControl, + "not done": notDone, + "stream zero": WindowUpdateFrame(0, 1), + "zero delta": WindowUpdateFrame(7, 0), + "trailing bytes": trailing, + "empty payload": empty, + } { + if _, _, ok := ParseWindowUpdate(fr); ok { + t.Fatalf("expected %q to be rejected", name) + } + } +} From d88576b98292eb53ab5c7eb6d62339a0e5b9bfdb Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 19:58:05 +0530 Subject: [PATCH 02/11] drpcstream: add per-stream send-window credit primitive Add sendWindow, the per-stream flow-control credit balance on the sender. acquire spends credit, blocking until enough is available; grant adds credit; close terminates the window. acquire returns early on close (with the close error) or context cancellation (with ctx.Err()), consuming no credit in either case; both are checked before any debit, so a canceled context never consumes credit or lets a frame proceed even when credit is available. Grants are no-revoke: grant takes an unsigned delta (the wire type), so a grant can only ever raise the balance, and the addition saturates at math.MaxInt64 so a large or malicious delta can never wrap it negative. The balance is a signed int64; acquire never drives it below zero, but the enablement layer may debit credit for bytes sent before flow control begins enforcing, leaving it negative, in which case acquire blocks until grants restore it. This is the primitive only -- it is not yet wired into the stream write path (that gate comes in a later commit). Co-Authored-By: roachdev-claude --- drpcstream/send_window.go | 115 +++++++++++++++++++++++++ drpcstream/send_window_test.go | 149 +++++++++++++++++++++++++++++++++ 2 files changed, 264 insertions(+) create mode 100644 drpcstream/send_window.go create mode 100644 drpcstream/send_window_test.go diff --git a/drpcstream/send_window.go b/drpcstream/send_window.go new file mode 100644 index 00000000..2ce9dca9 --- /dev/null +++ b/drpcstream/send_window.go @@ -0,0 +1,115 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "math" + "sync" +) + +// sendWindow is a per-stream flow-control credit balance on the sender. It +// tracks how many more bytes the stream is allowed to put on the wire right +// now. acquire spends credit (blocking until enough is available), grant adds +// credit, and close terminates the window. +// +// Grants are no-revoke: grant is strictly additive, and the receiver never +// takes back credit it has already issued. The balance is a signed int64; +// acquire never drives it below zero, but the enablement layer may debit credit +// for bytes sent before flow control begins enforcing, leaving it negative, in +// which case acquire blocks until grants restore it. +type sendWindow struct { + mu sync.Mutex + avail int64 // available credit; signed (may be negative — see doc) + closed bool // set once by close; no further acquires succeed + err error // terminal error returned by acquire after close + notify chan struct{} // closed+replaced to wake parked acquirers +} + +// newSendWindow returns a sendWindow seeded with initial credit. +func newSendWindow(initial int64) *sendWindow { + return &sendWindow{avail: initial, notify: make(chan struct{})} +} + +// available returns the current credit balance. +func (w *sendWindow) available() int64 { + w.mu.Lock() + defer w.mu.Unlock() + return w.avail +} + +// acquire blocks until n bytes of credit are available, then debits them and +// returns nil. It returns early if the window is closed (with the close error) +// or ctx is canceled (with ctx.Err()), consuming no credit in either case. Both +// are checked before any debit, so a canceled context never consumes credit or +// lets a frame proceed even when credit is available. +func (w *sendWindow) acquire(ctx context.Context, n int64) error { + for { + w.mu.Lock() + switch { + case w.closed: + err := w.err + w.mu.Unlock() + return err + case ctx.Err() != nil: + w.mu.Unlock() + return ctx.Err() + case w.avail >= n: + w.avail -= n + w.mu.Unlock() + return nil + } + // Snapshot the notify channel under the lock before parking, so a grant + // or close that fires the instant we unlock is not missed. + ch := w.notify + w.mu.Unlock() + + select { + case <-ch: + // Credit was granted or the window closed; loop and re-check. + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// grant adds n bytes of credit and wakes any parked acquirer. n is unsigned +// (the wire delta type), so a grant can only ever raise the balance — no-revoke. +// The addition saturates at math.MaxInt64 so a large or malicious delta can +// never wrap the balance negative. Grants after close are ignored (the window +// is dead). +func (w *sendWindow) grant(n uint64) { + w.mu.Lock() + if !w.closed { + if n >= uint64(math.MaxInt64) { + w.avail = math.MaxInt64 + } else if sum := w.avail + int64(n); sum < w.avail { + w.avail = math.MaxInt64 // overflowed past MaxInt64 + } else { + w.avail = sum + } + w.wakeLocked() + } + w.mu.Unlock() +} + +// close terminates the window with err, waking every parked acquirer, which +// then returns err. Subsequent acquires also return err. It is a no-op if the +// window is already closed. +func (w *sendWindow) close(err error) { + w.mu.Lock() + if !w.closed { + w.closed = true + w.err = err + w.wakeLocked() + } + w.mu.Unlock() +} + +// wakeLocked broadcasts to all parked acquirers by closing the current notify +// channel and installing a fresh one. It must be called with w.mu held. +func (w *sendWindow) wakeLocked() { + close(w.notify) + w.notify = make(chan struct{}) +} diff --git a/drpcstream/send_window_test.go b/drpcstream/send_window_test.go new file mode 100644 index 00000000..0e1a6b74 --- /dev/null +++ b/drpcstream/send_window_test.go @@ -0,0 +1,149 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "errors" + "math" + "testing" + "time" + + "github.com/zeebo/assert" + "github.com/zeebo/errs" +) + +// blockShort is how long we wait to conclude that an acquire is (correctly) +// blocked before we release it. +const blockShort = 20 * time.Millisecond + +func TestSendWindowAcquireImmediate(t *testing.T) { + w := newSendWindow(1000) + assert.Equal(t, w.available(), int64(1000)) + + assert.NoError(t, w.acquire(context.Background(), 400)) + assert.Equal(t, w.available(), int64(600)) + + assert.NoError(t, w.acquire(context.Background(), 600)) + assert.Equal(t, w.available(), int64(0)) +} + +func TestSendWindowGrantsAccumulate(t *testing.T) { + w := newSendWindow(0) + w.grant(100) + w.grant(50) + assert.Equal(t, w.available(), int64(150)) + + assert.NoError(t, w.acquire(context.Background(), 150)) + assert.Equal(t, w.available(), int64(0)) +} + +func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) { + w := newSendWindow(100) + done := make(chan error, 1) + go func() { done <- w.acquire(context.Background(), 300) }() + + // Not enough credit yet: acquire must block. + select { + case <-done: + t.Fatal("acquire returned before sufficient credit") + case <-time.After(blockShort): + } + + w.grant(250) // 100 + 250 = 350 >= 300 + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("acquire did not return after grant") + } + assert.Equal(t, w.available(), int64(50)) +} + +func TestSendWindowAcquireContextCancel(t *testing.T) { + w := newSendWindow(0) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- w.acquire(ctx, 100) }() + + select { + case <-done: + t.Fatal("acquire returned before cancellation") + case <-time.After(blockShort): + } + + cancel() + + select { + case err := <-done: + assert.That(t, errors.Is(err, context.Canceled)) + case <-time.After(time.Second): + t.Fatal("acquire did not wake on context cancellation") + } + // Credit was not consumed by a failed acquire. + assert.Equal(t, w.available(), int64(0)) +} + +func TestSendWindowGrantSaturates(t *testing.T) { + // Adding to a near-max balance saturates at MaxInt64 instead of wrapping. + w := newSendWindow(math.MaxInt64 - 10) + w.grant(100) + assert.Equal(t, w.available(), int64(math.MaxInt64)) + + // A single delta larger than MaxInt64 (the wire delta is uint64) saturates + // rather than becoming a negative balance. + w2 := newSendWindow(0) + w2.grant(math.MaxUint64) + assert.Equal(t, w2.available(), int64(math.MaxInt64)) + assert.That(t, w2.available() > 0) // never wrapped negative + + // A grant that raises a negative balance is applied exactly (no saturation). + w3 := newSendWindow(0) + w3.avail = -300 + w3.grant(500) + assert.Equal(t, w3.available(), int64(200)) +} + +func TestSendWindowAcquireCanceledCtxWithCredit(t *testing.T) { + w := newSendWindow(1000) + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already canceled, with credit available + + err := w.acquire(ctx, 100) + assert.That(t, errors.Is(err, context.Canceled)) + // A canceled context must not consume credit even though it was available. + assert.Equal(t, w.available(), int64(1000)) +} + +func TestSendWindowCloseWakesAcquire(t *testing.T) { + w := newSendWindow(0) + closeErr := errs.New("terminated") + done := make(chan error, 1) + go func() { done <- w.acquire(context.Background(), 100) }() + + select { + case <-done: + t.Fatal("acquire returned before close") + case <-time.After(blockShort): + } + + w.close(closeErr) + + select { + case err := <-done: + assert.That(t, errors.Is(err, closeErr)) + case <-time.After(time.Second): + t.Fatal("acquire did not wake on close") + } +} + +func TestSendWindowAcquireAfterClose(t *testing.T) { + w := newSendWindow(1000) + closeErr := errs.New("closed") + w.close(closeErr) + + // Even though credit is available, a closed window returns the close error. + assert.That(t, errors.Is(w.acquire(context.Background(), 1), closeErr)) +} From 6594c1418b567ca4df29db5b9248f415e6d1e9f1 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 20:24:09 +0530 Subject: [PATCH 03/11] drpcstream: gate data-frame sends on the per-stream send window Wire the sendWindow credit gate into rawWriteLocked: a KindMessage frame now acquires len(frame) bytes of per-stream send credit before it is handed to the writer. Control frames (invoke/metadata) bypass the gate. The window is opt-in: a stream has no send window by default, so data writes stay ungated (unlimited) and behavior is unchanged until one is installed. terminate closes the window, so a send parked on credit wakes with the termination error -- matching the write path's existing preemption semantics. Per-stream only; the connection-level window and the enablement path that installs the window come in later commits. Co-Authored-By: roachdev-claude --- drpcstream/send_window_gate_test.go | 88 +++++++++++++++++++++++++++++ drpcstream/stream.go | 18 ++++++ 2 files changed, 106 insertions(+) create mode 100644 drpcstream/send_window_gate_test.go diff --git a/drpcstream/send_window_gate_test.go b/drpcstream/send_window_gate_test.go new file mode 100644 index 00000000..268094cf --- /dev/null +++ b/drpcstream/send_window_gate_test.go @@ -0,0 +1,88 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "testing" + "time" + + "github.com/zeebo/assert" + "github.com/zeebo/errs" + + "storj.io/drpc/drpcwire" +) + +// newGateStream builds a stream writing to io.Discard with an explicit +// SplitSize so small payloads are a single frame. +func newGateStream(t *testing.T) *Stream { + mw := testMuxWriter(t) + return NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) +} + +// By default no send window is installed, so data writes are ungated +// (unlimited) and behavior is unchanged. +func TestStream_SendWindowDefaultUngated(t *testing.T) { + st := newGateStream(t) + assert.That(t, st.sendw == nil) + assert.NoError(t, st.RawWrite(drpcwire.KindMessage, []byte("hello"))) +} + +// With a finite send window, a data write blocks until enough credit is +// granted. +func TestStream_SendWindowGatesDataWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(4) // 4 bytes of credit + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("hello")) }() // 5 bytes > 4 + + select { + case <-done: + t.Fatal("data write returned before sufficient credit") + case <-time.After(blockShort): + } + + st.sendw.grant(1) // 4 + 1 = 5 >= 5 + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("data write did not complete after grant") + } +} + +// Control kinds (here, invoke) are not flow-controlled: they proceed even with +// zero send credit. +func TestStream_SendWindowControlKindsBypassGate(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // no credit at all + + assert.NoError(t, st.WriteInvoke("service.Method", nil)) +} + +// Terminating the stream wakes a send parked on credit. +func TestStream_SendWindowTerminateWakesParkedWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // send will park immediately + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }() + + select { + case <-done: + t.Fatal("data write returned before termination") + case <-time.After(blockShort): + } + + st.Cancel(errs.New("boom")) + + select { + case err := <-done: + assert.Error(t, err) + case <-time.After(time.Second): + t.Fatal("parked data write did not wake on termination") + } +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index fe7b0b48..0d6e9ba1 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -55,6 +55,10 @@ type Stream struct { recvQueue ringBuffer wbuf []byte + // sendw is the per-stream send-side flow-control window. It is nil when + // flow control is not enabled, in which case data writes are ungated. + sendw *sendWindow + mu sync.Mutex // protects state transitions sigs struct { send drpcsignal.Signal // set when done sending messages @@ -352,6 +356,9 @@ func (s *Stream) terminate(err error) { s.sigs.recv.Set(err) s.sigs.term.Set(err) s.recvQueue.Close(err) + if s.sendw != nil { + s.sendw.close(err) + } s.checkFinished() } @@ -405,6 +412,17 @@ func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) { drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(fr.Data))) s.log("SEND", fr.String) + // Flow control: a data frame must acquire send credit before it goes on + // the wire. Only KindMessage is flow-controlled; control frames (e.g. + // invoke/metadata) bypass. When no window is installed (the default), + // sends are ungated. acquire blocks until credit arrives and is + // interruptible by stream termination, which closes the window. + if kind == drpcwire.KindMessage && s.sendw != nil { + if err := s.sendw.acquire(s.Context(), int64(len(fr.Data))); err != nil { + return err + } + } + // Pass the send signal, the write side's own stop signal, so a parked // WriteFrame is woken when sending ends (cancel, error, close) and // returns that signal's error directly. That is io.EOF in the cancel and From 2f61b3df1d8e780d2ce45b39044021a77c85028c Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 20:35:47 +0530 Subject: [PATCH 04/11] drpcstream: add per-stream receive-window grant policy Add recvWindow, the per-stream receive side of flow control. It tracks buffered bytes (in-progress reassembly plus completed, not-yet-consumed data) and decides when to return credit to the sender: grants are withheld while buffered is at or above the high-water mark, and the accrued returnable credit is otherwise released once it reaches the threshold, coalescing many frames into a single grant. Consuming reopens a gate that the high-water mark had closed. dispatched/consumed return the credit delta the caller should emit as a KindWindowUpdate. Emitting the frame, wiring into the read path (HandleFrame/MsgRecv), and the connection-level receive budget come in later commits. Co-Authored-By: roachdev-claude --- drpcstream/recv_window.go | 75 ++++++++++++++++++++++++++++++++++ drpcstream/recv_window_test.go | 62 ++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 drpcstream/recv_window.go create mode 100644 drpcstream/recv_window_test.go diff --git a/drpcstream/recv_window.go b/drpcstream/recv_window.go new file mode 100644 index 00000000..dd8fc188 --- /dev/null +++ b/drpcstream/recv_window.go @@ -0,0 +1,75 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import "sync" + +// recvWindow is the per-stream receive side of flow control. It tracks how many +// bytes are currently buffered for the stream (in-progress reassembly plus +// completed, not-yet-consumed data) and decides when to return credit to the +// sender. +// +// The receiver holds no credit balance of its own; it accrues "returnable" +// credit as bytes are dispatched off the wire and releases it as grants: +// - while buffered is at or above the high-water mark, grants are withheld +// (backpressure), though the returnable credit keeps accruing; +// - otherwise, accrued credit is released once it reaches the threshold, +// coalescing many frames into a single grant. +// +// dispatched and consumed each return the credit delta the caller should send +// as a KindWindowUpdate (0 when nothing should be sent). Emitting the frame is +// the caller's job. dispatched (reader goroutine) and consumed (application +// goroutine) can run concurrently, so the state is mutex-guarded. +type recvWindow struct { + mu sync.Mutex + buffered int64 // in-progress reassembly + completed, not-yet-consumed bytes + pending int64 // returnable credit accrued but not yet granted + highWater int64 // withhold grants while buffered is at or above this + threshold int64 // release accrued credit once it reaches this +} + +// newRecvWindow returns a recvWindow with the given high-water mark and grant +// threshold. +func newRecvWindow(highWater, threshold int64) *recvWindow { + return &recvWindow{highWater: highWater, threshold: threshold} +} + +// dispatched records that n bytes were dispatched off the wire into the +// stream's buffers and returns the credit delta to grant (0 if none). +func (w *recvWindow) dispatched(n int64) int64 { + w.mu.Lock() + defer w.mu.Unlock() + w.buffered += n + w.pending += n + return w.maybeGrantLocked() +} + +// consumed records that n bytes were consumed by the application and returns +// the credit delta to grant (0 if none). Consuming can reopen a gate that was +// closed by the high-water mark, flushing credit accrued during the pause. +func (w *recvWindow) consumed(n int64) int64 { + w.mu.Lock() + defer w.mu.Unlock() + w.buffered -= n + return w.maybeGrantLocked() +} + +// bufferedBytes returns the current buffered byte count. +func (w *recvWindow) bufferedBytes() int64 { + w.mu.Lock() + defer w.mu.Unlock() + return w.buffered +} + +// maybeGrantLocked releases the accrued returnable credit when the high-water +// gate is open and the threshold is met, returning the delta (0 otherwise). It +// must be called with w.mu held. +func (w *recvWindow) maybeGrantLocked() int64 { + if w.buffered < w.highWater && w.pending >= w.threshold { + g := w.pending + w.pending = 0 + return g + } + return 0 +} diff --git a/drpcstream/recv_window_test.go b/drpcstream/recv_window_test.go new file mode 100644 index 00000000..3331c1cf --- /dev/null +++ b/drpcstream/recv_window_test.go @@ -0,0 +1,62 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "testing" + + "github.com/zeebo/assert" +) + +// Below the high-water mark, accrued returnable credit is released as a grant +// once it reaches the threshold, and withheld (0) below it. +func TestRecvWindowGrantsAtThreshold(t *testing.T) { + w := newRecvWindow(1000, 100) + assert.Equal(t, w.dispatched(60), int64(0)) // pending 60 < 100 + assert.Equal(t, w.dispatched(60), int64(120)) // pending 120 >= 100, buffered 120 < 1000 + assert.Equal(t, w.bufferedBytes(), int64(120)) + assert.Equal(t, w.dispatched(50), int64(0)) // pending reset by the grant; 50 < 100 +} + +// At or above the high-water mark, grants are withheld even when the accrued +// credit is past the threshold. +func TestRecvWindowWithholdsAtOrAboveHighWater(t *testing.T) { + w := newRecvWindow(100, 50) + assert.Equal(t, w.dispatched(120), int64(0)) // buffered 120 >= 100 -> withhold + assert.Equal(t, w.bufferedBytes(), int64(120)) +} + +// Consuming buffered data drops below the high-water mark and flushes the +// credit accrued while the gate was closed (resume-on-consume). +func TestRecvWindowResumesOnConsume(t *testing.T) { + w := newRecvWindow(100, 50) + assert.Equal(t, w.dispatched(120), int64(0)) // withheld; pending 120 + assert.Equal(t, w.consumed(30), int64(120)) // buffered 90 < 100 -> flush pending + assert.Equal(t, w.bufferedBytes(), int64(90)) +} + +// Consuming while the accrued credit is below the threshold emits no grant. +func TestRecvWindowConsumeBelowThresholdNoGrant(t *testing.T) { + w := newRecvWindow(1000, 100) + assert.Equal(t, w.dispatched(50), int64(0)) // pending 50 < 100 + assert.Equal(t, w.consumed(20), int64(0)) // buffered 30, pending 50 < 100 + assert.Equal(t, w.bufferedBytes(), int64(30)) +} + +// Many small dispatches coalesce into few grants, each carrying the accrued +// amount. +func TestRecvWindowCoalescesGrants(t *testing.T) { + w := newRecvWindow(1_000_000, 100) + grants, granted := 0, int64(0) + for i := 0; i < 10; i++ { + if g := w.dispatched(30); g > 0 { + grants++ + granted += g + } + } + // 10 dispatches of 30 (300 bytes) with threshold 100 coalesce into 2 grants + // of 120 each; the remaining 60 stays accrued. + assert.Equal(t, grants, 2) + assert.Equal(t, granted, int64(240)) +} From 74ad6c38bbf5fda9c6691e976d1f680e3ca53a6b Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 20:45:38 +0530 Subject: [PATCH 05/11] drpcstream: wire the receive-window grant policy into the read path Install a per-stream recvWindow and drive it from the read/recv path: - HandleFrame intercepts an incoming KindWindowUpdate before the assembler and applies it to the send window. Intercepting early keeps a grant that interleaves with an in-progress message (grants are emitted off the write lock) from disturbing reassembly. - A dispatched KindMessage frame returns credit via recvWindow, emitting a KindWindowUpdate once the accrued amount crosses the threshold and the high-water gate is open. - RawRecv/MsgRecv call recvWindow after dequeuing; consuming can reopen a gate the high-water mark had closed and flush the accrued grant. Grants are control frames, appended without blocking, so emitting them from the reader goroutine is safe. The window is opt-in: with no recvw installed no grants are emitted, and an incoming grant with no send window is ignored, so default behavior is unchanged. Stream-level only; the connection-level receive budget and the enablement path that installs the windows come in later commits. Co-Authored-By: roachdev-claude --- drpcstream/recv_window_wiring_test.go | 142 ++++++++++++++++++++++++++ drpcstream/stream.go | 46 +++++++++ 2 files changed, 188 insertions(+) create mode 100644 drpcstream/recv_window_wiring_test.go diff --git a/drpcstream/recv_window_wiring_test.go b/drpcstream/recv_window_wiring_test.go new file mode 100644 index 00000000..85107199 --- /dev/null +++ b/drpcstream/recv_window_wiring_test.go @@ -0,0 +1,142 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "io" + "testing" + "time" + + "github.com/zeebo/assert" + + "storj.io/drpc/drpcwire" +) + +// captureWriter returns a MuxWriter whose output frames are decoded and +// delivered on the returned channel, so tests can observe emitted grants. +func captureWriter(t *testing.T) (*drpcwire.MuxWriter, <-chan drpcwire.Frame) { + t.Helper() + pr, pw := io.Pipe() + mw := drpcwire.NewMuxWriter(pw, func(error) {}) + frames := make(chan drpcwire.Frame, 64) + rd := drpcwire.NewReader(pr) + go func() { + for { + fr, err := rd.ReadFrame() + if err != nil { + return + } + frames <- fr + } + }() + t.Cleanup(func() { + mw.Stop(nil) + <-mw.Done() + _ = pw.Close() + _ = pr.Close() + }) + return mw, frames +} + +func waitFrame(t *testing.T, frames <-chan drpcwire.Frame) drpcwire.Frame { + t.Helper() + select { + case fr := <-frames: + return fr + case <-time.After(time.Second): + t.Fatal("expected a frame to be written, got none") + return drpcwire.Frame{} + } +} + +func assertNoFrame(t *testing.T, frames <-chan drpcwire.Frame) { + t.Helper() + select { + case fr := <-frames: + t.Fatalf("unexpected frame written: %v", fr) + case <-time.After(blockShort): + } +} + +func msgFrame(sid, mid uint64, data []byte, done bool) drpcwire.Frame { + return drpcwire.Frame{ + ID: drpcwire.ID{Stream: sid, Message: mid}, + Kind: drpcwire.KindMessage, + Data: data, + Done: done, + } +} + +// An incoming KindWindowUpdate is intercepted and applied to the send window. +func TestStream_IncomingGrantAppliesToSendWindow(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) + + assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 500))) + assert.Equal(t, st.sendw.available(), int64(500)) +} + +// A grant interleaved with an in-progress message is intercepted before the +// assembler, so reassembly is undisturbed. +func TestStream_GrantDoesNotDisturbReassembly(t *testing.T) { + st := newGateStream(t) + sid := st.ID() + + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, []byte("foo"), false))) + assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(sid, 100))) // interleaved + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, []byte("bar"), true))) + + got, err := st.RawRecv() + assert.NoError(t, err) + assert.Equal(t, string(got), "foobar") +} + +// Dispatching a data frame returns credit once the threshold is met. +func TestStream_DispatchEmitsGrant(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + st.recvw = newRecvWindow(1<<20, 100) + + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 150), true))) + + fr := waitFrame(t, frames) + sid, delta, ok := drpcwire.ParseWindowUpdate(fr) + assert.That(t, ok) + assert.Equal(t, sid, st.ID()) + assert.Equal(t, delta, uint64(150)) +} + +// Above the high-water mark, dispatch withholds; consuming resumes and flushes +// the accrued credit. +func TestStream_ConsumeEmitsGrantOnResume(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + st.recvw = newRecvWindow(100, 50) + + // 120 bytes buffered >= high-water 100 -> withheld. + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 120), true))) + assertNoFrame(t, frames) + + // Consume -> buffered drops below high-water -> flush accrued 120. + _, err := st.RawRecv() + assert.NoError(t, err) + + fr := waitFrame(t, frames) + _, delta, ok := drpcwire.ParseWindowUpdate(fr) + assert.That(t, ok) + assert.Equal(t, delta, uint64(120)) +} + +// With no receive window installed, no grants are emitted and an incoming grant +// with no send window is harmlessly ignored. +func TestStream_NoWindowsNoGrants(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 150), true))) + assertNoFrame(t, frames) + + assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 100))) +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 0d6e9ba1..bd1e0b42 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -59,6 +59,10 @@ type Stream struct { // flow control is not enabled, in which case data writes are ungated. sendw *sendWindow + // recvw is the per-stream receive-side flow-control policy. It is nil when + // flow control is not enabled, in which case no grants are emitted. + recvw *recvWindow + mu sync.Mutex // protects state transitions sigs struct { send drpcsignal.Signal // set when done sending messages @@ -211,10 +215,30 @@ func (s *Stream) HandleFrame(fr drpcwire.Frame) (err error) { return nil } + // Flow-control grants are out-of-band signaling, not part of the message + // stream, so they are intercepted before the assembler. Grants are emitted + // off the write lock, so one can arrive interleaved with an in-progress + // message; intercepting here keeps it from disturbing reassembly. Without a + // send window (flow control not enabled) the grant is simply ignored. + if fr.Kind == drpcwire.KindWindowUpdate { + if s.sendw != nil { + if _, delta, ok := drpcwire.ParseWindowUpdate(fr); ok { + s.sendw.grant(delta) + } + } + return nil + } + packet, packetReady, err := s.pa.AppendFrame(fr) if err != nil { return err } + + // A data frame dispatched off the wire may return credit to the sender. + if fr.Kind == drpcwire.KindMessage && s.recvw != nil { + s.emitGrant(s.recvw.dispatched(int64(len(fr.Data)))) + } + if !packetReady { return nil } @@ -280,6 +304,18 @@ func (s *Stream) handlePacket(pkt drpcwire.Packet) (err error) { } } +// emitGrant sends a per-stream flow-control credit grant of delta bytes to the +// peer, and is a no-op when delta is not positive. Grants are control frames, +// which WriteFrame appends immediately without blocking on backpressure, so +// this is safe to call from the reader goroutine. It is best-effort: a write +// error means the stream is going away. +func (s *Stream) emitGrant(delta int64) { + if delta <= 0 { + return + } + _ = s.wr.WriteFrame(drpcwire.WindowUpdateFrame(s.id.Stream, uint64(delta)), nil) +} + // // helpers // @@ -447,8 +483,13 @@ func (s *Stream) RawRecv() (data []byte, err error) { return nil, err } data = append([]byte(nil), b...) + n := len(b) s.recvQueue.Done() + // Consuming buffered data can return credit to the sender. + if s.recvw != nil { + s.emitGrant(s.recvw.consumed(int64(n))) + } return data, nil } @@ -488,9 +529,14 @@ func (s *Stream) MsgRecv(msg drpc.Message, enc drpc.Encoding) (err error) { if err != nil { return err } + n := len(b) err = enc.Unmarshal(b, msg) s.recvQueue.Done() + // Consuming buffered data can return credit to the sender. + if s.recvw != nil { + s.emitGrant(s.recvw.consumed(int64(n))) + } return err } From 21d9e46c7d098a82d4a89539c88a669fabe83f06 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 21:15:44 +0530 Subject: [PATCH 06/11] drpcstream: install per-stream windows from a flow-control option Add Options.FlowControl to enable per-stream flow control. When Enabled, NewWithOptions installs the send window (seeded with StreamWindow as initial credit) and the receive window (HighWater + GrantThreshold), so data writes become credit-gated and grants are emitted -- all driven by configuration rather than injected by tests. This is static, symmetric enablement: both peers must enable it with compatible sizes. Advertise-on-enable (safe under mixed enablement), the manager wiring that turns it on for a whole connection, and the connection-level window come in later commits. Co-Authored-By: roachdev-claude --- drpcstream/flow_control_option_test.go | 70 ++++++++++++++++++++++++++ drpcstream/stream.go | 33 ++++++++++++ 2 files changed, 103 insertions(+) create mode 100644 drpcstream/flow_control_option_test.go diff --git a/drpcstream/flow_control_option_test.go b/drpcstream/flow_control_option_test.go new file mode 100644 index 00000000..eac6dc40 --- /dev/null +++ b/drpcstream/flow_control_option_test.go @@ -0,0 +1,70 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "testing" + "time" + + "github.com/zeebo/assert" + + "storj.io/drpc/drpcwire" +) + +// Enabling flow control via Options installs the send and receive windows, +// seeding the send window with the configured stream window as initial credit. +func TestStream_FlowControlEnabledInstallsWindows(t *testing.T) { + mw := testMuxWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{ + SplitSize: 64 << 10, + FlowControl: FlowControl{ + Enabled: true, + StreamWindow: 256 << 10, + HighWater: 4 << 20, + GrantThreshold: 128 << 10, + }, + }) + assert.That(t, st.sendw != nil) + assert.That(t, st.recvw != nil) + assert.Equal(t, st.sendw.available(), int64(256<<10)) +} + +// Without flow control enabled, no windows are installed and behavior is +// unchanged (ungated). +func TestStream_FlowControlDisabledNoWindows(t *testing.T) { + mw := testMuxWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + assert.That(t, st.sendw == nil) + assert.That(t, st.recvw == nil) +} + +// End to end via the option: an option-installed window gates a data write that +// exceeds the initial credit, and an incoming grant resumes it. +func TestStream_FlowControlOptionGatesAndResumes(t *testing.T) { + mw := testMuxWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{ + SplitSize: 64 << 10, + FlowControl: FlowControl{Enabled: true, StreamWindow: 4, HighWater: 1 << 20, GrantThreshold: 2}, + }) + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("hello")) }() // 5 bytes > 4 + + select { + case <-done: + t.Fatal("write returned before sufficient credit") + case <-time.After(blockShort): + } + + // A grant from the peer tops up the send window and resumes the write. + assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 1))) + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("write did not resume after grant") + } +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index bd1e0b42..786c754a 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -30,10 +30,35 @@ type Options struct { // more allocations. 0 is unlimited. MaximumBufferSize int + // FlowControl configures per-stream flow control. It is disabled by default. + FlowControl FlowControl + // Internal contains options that are for internal use only. Internal drpcopts.Stream } +// FlowControl configures per-stream flow control. When Enabled, the stream +// installs a send window (data writes acquire credit before going on the wire) +// and a receive window (grants are returned as data is dispatched and +// consumed). It is the caller's responsibility to size these so the liveness +// rule StreamWindow >= 2*SplitSize holds. +type FlowControl struct { + // Enabled turns per-stream flow control on for the stream. + Enabled bool + + // StreamWindow is the per-stream send window: the initial send credit and + // the nominal window size. + StreamWindow int64 + + // HighWater is the receive-side buffered-byte high-water mark above which + // grants are withheld. + HighWater int64 + + // GrantThreshold is the amount of returnable credit that must accrue before + // a grant is emitted, coalescing many frames into one KindWindowUpdate. + GrantThreshold int64 +} + // Stream represents an rpc actively happening on a transport. type Stream struct { ctx streamCtx @@ -123,6 +148,14 @@ func NewWithOptions( s.recvQueue.init(pool) + // When flow control is enabled, install the per-stream windows. The send + // window starts with StreamWindow credit (statically agreed with the peer); + // the receive window drives grant emission. + if opts.FlowControl.Enabled { + s.sendw = newSendWindow(opts.FlowControl.StreamWindow) + s.recvw = newRecvWindow(opts.FlowControl.HighWater, opts.FlowControl.GrantThreshold) + } + return s } From b8e91f1f456e62edf121354845b7c468fb3c07a8 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 21:20:50 +0530 Subject: [PATCH 07/11] drpcmanager: end-to-end test for flow-control enablement The manager already threads drpcstream.Options through newStream, so enabling FlowControl in Options.Stream installs windows on every stream the manager creates (client and server) with no additional wiring. Add an end-to-end test over a real net.Pipe connection between two flow-control-enabled managers: a message larger than the send window completes only if credit grants flow back across the connection, exercising the full stream-level credit loop (gate -> grant-on-dispatch -> apply-grant) through the real read/write paths. Co-Authored-By: roachdev-claude --- drpcmanager/flow_control_test.go | 72 ++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 drpcmanager/flow_control_test.go diff --git a/drpcmanager/flow_control_test.go b/drpcmanager/flow_control_test.go new file mode 100644 index 00000000..1fb1c8fb --- /dev/null +++ b/drpcmanager/flow_control_test.go @@ -0,0 +1,72 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcmanager + +import ( + "context" + "errors" + "io" + "net" + "testing" + + "github.com/zeebo/assert" + + "storj.io/drpc/drpcstream" + "storj.io/drpc/drpctest" + "storj.io/drpc/drpcwire" +) + +// Enabling flow control through the manager's Stream options installs windows +// on every stream it creates (client and server), so a message larger than the +// send window only completes if credit grants flow across the real connection. +func TestManager_FlowControlEndToEnd(t *testing.T) { + ctx := drpctest.NewTracker(t) + defer ctx.Close() + + cconn, sconn := net.Pipe() + defer func() { _ = cconn.Close() }() + defer func() { _ = sconn.Close() }() + + opts := Options{Stream: drpcstream.Options{ + SplitSize: 64 << 10, + FlowControl: drpcstream.FlowControl{ + Enabled: true, + StreamWindow: 128 << 10, // 2 frames of initial credit + HighWater: 1 << 20, + GrantThreshold: 64 << 10, + }, + }} + + cman := NewWithOptions(cconn, Client, opts) + defer func() { _ = cman.Close() }() + sman := NewWithOptions(sconn, Server, opts) + defer func() { _ = sman.Close() }() + + // 256 KiB is four frames: two fit in the initial window, the rest can only + // be sent as the server dispatches frames and returns credit. + msg := make([]byte, 256<<10) + + ctx.Run(func(ctx context.Context) { + stream, err := cman.NewClientStream(ctx, "rpc") + assert.NoError(t, err) + assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc"))) + assert.NoError(t, stream.RawWrite(drpcwire.KindMessage, msg)) + assert.NoError(t, stream.Close()) + }) + + ctx.Run(func(ctx context.Context) { + stream, _, err := sman.NewServerStream(ctx) + assert.NoError(t, err) + defer func() { _ = stream.Close() }() + + got, err := stream.RawRecv() + assert.NoError(t, err) + assert.Equal(t, len(got), len(msg)) + + _, err = stream.RawRecv() + assert.That(t, errors.Is(err, io.EOF)) + }) + + ctx.Wait() +} From 482d67bbcde88323ee480b6e2adf72fb1d06fbae Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 21:50:20 +0530 Subject: [PATCH 08/11] drpcmanager: test that context cancellation wakes a parked sender Add an integration test for the production wake path of a sender parked on flow-control credit. A flow-control-enabled client sends a message larger than its window to a server that has no flow control (so it never returns credit); the send parks, and cancelling the RPC context wakes it. This exercises manageStream -> Cancel -> terminate -> sendWindow.close end to end, which the unit tests previously covered only in pieces (the primitive's ctx path, and stream-level Cancel). Co-Authored-By: roachdev-claude --- drpcmanager/flow_control_test.go | 69 ++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/drpcmanager/flow_control_test.go b/drpcmanager/flow_control_test.go index 1fb1c8fb..f7c45b60 100644 --- a/drpcmanager/flow_control_test.go +++ b/drpcmanager/flow_control_test.go @@ -9,6 +9,7 @@ import ( "io" "net" "testing" + "time" "github.com/zeebo/assert" @@ -70,3 +71,71 @@ func TestManager_FlowControlEndToEnd(t *testing.T) { ctx.Wait() } + +// Cancelling an RPC's context wakes a sender parked on flow-control credit. The +// client has flow control with a small window; the server does not, so it never +// returns credit and the client's oversized send parks. This exercises the +// production wake path: manageStream sees the cancellation and calls +// Cancel -> terminate -> sendWindow.close. +func TestManager_FlowControlContextCancelWakesParkedSend(t *testing.T) { + tr := drpctest.NewTracker(t) + defer tr.Close() + + cconn, sconn := net.Pipe() + defer func() { _ = cconn.Close() }() + defer func() { _ = sconn.Close() }() + + cman := NewWithOptions(cconn, Client, Options{Stream: drpcstream.Options{ + SplitSize: 64 << 10, + FlowControl: drpcstream.FlowControl{ + Enabled: true, + StreamWindow: 128 << 10, // two frames of initial credit + HighWater: 1 << 20, + GrantThreshold: 64 << 10, + }, + }}) + defer func() { _ = cman.Close() }() + // Server has no flow control, so it never returns credit. + sman := New(sconn, Server) + defer func() { _ = sman.Close() }() + + // Accept the stream and drain; the message never completes (the sender + // parks mid-message), so RawRecv just blocks until teardown. + tr.Run(func(ctx context.Context) { + stream, _, err := sman.NewServerStream(ctx) + if err != nil { + return + } + defer func() { _ = stream.Close() }() + for { + if _, err := stream.RawRecv(); err != nil { + return + } + } + }) + + streamCtx, cancel := context.WithCancel(context.Background()) + stream, err := cman.NewClientStream(streamCtx, "rpc") + assert.NoError(t, err) + assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc"))) + + // A 512 KiB message exceeds the 128 KiB window; once the initial credit is + // spent the send parks waiting for grants that never come. + done := make(chan error, 1) + go func() { done <- stream.RawWrite(drpcwire.KindMessage, make([]byte, 512<<10)) }() + + select { + case <-done: + t.Fatal("send returned before it could park on credit") + case <-time.After(50 * time.Millisecond): + } + + cancel() + + select { + case err := <-done: + assert.Error(t, err) + case <-time.After(time.Second): + t.Fatal("parked send did not wake on context cancellation") + } +} From fd0c66a0d068b3f262f28418c5c76bd7cbb23055 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 21:36:59 +0530 Subject: [PATCH 09/11] drpcmanager: cap concurrent streams with MaxStreams Add Options.MaxStreams, a per-connection cap on how many concurrent streams a peer may open against this manager. When the cap is reached an inbound stream is refused in the read path -- before any per-stream state is allocated -- with a stream-level KindError, leaving the connection up. 0 means unlimited. This bounds the per-stream bookkeeping (goroutines, stream-table entries, window counters) that the byte-based flow-control windows do not, giving DoS resistance against a peer that opens many zero-byte streams. It is the third flow-control bound alongside the per-stream and connection windows, and is independent of the credit path. Co-Authored-By: roachdev-claude --- drpcmanager/manager.go | 38 ++++++++++++++++++ drpcmanager/max_streams_test.go | 68 +++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 drpcmanager/max_streams_test.go diff --git a/drpcmanager/manager.go b/drpcmanager/manager.go index c0cd43f0..11c3d717 100644 --- a/drpcmanager/manager.go +++ b/drpcmanager/manager.go @@ -28,6 +28,10 @@ import ( var managerClosed = errs.Class("manager closed") +// maxStreamsExceeded is returned to a peer whose inbound stream is refused +// because the manager is already at its concurrent-stream cap. +var maxStreamsExceeded = errs.Class("too many concurrent streams") + // Options controls configuration settings for a manager. type Options struct { // Reader are passed to any readers the manager creates. @@ -46,6 +50,13 @@ type Options struct { // handling. When enabled, the server stream will decode incoming metadata // into grpc metadata in the context. GRPCMetadataCompatMode bool + + // MaxStreams caps the number of concurrent streams the peer may open against + // this manager. When the cap is reached, further inbound streams are refused + // with a stream-level error rather than admitted, bounding the per-stream + // bookkeeping (goroutines, stream-table entries) that byte windows do not. + // 0 means unlimited. + MaxStreams int } // Manager handles the logic of managing a transport for a drpc client or @@ -246,6 +257,16 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error { return nil } + // Enforce the concurrent-stream cap before admitting the stream. Refuse it + // here, in the read path, so the per-stream state the cap protects is never + // allocated; the peer learns via a stream-level error and the connection + // stays up. + if m.opts.MaxStreams > 0 && m.streams.Len() >= m.opts.MaxStreams { + m.rejectStream(pkt.ID.Stream) + delete(m.pendingStreams, fr.ID.Stream) + return nil + } + // Invoke packet completes the sequence. Send to NewServerStream. select { case m.invokes <- invokeInfo{sid: pkt.ID.Stream, data: pkt.Data, metadata: ps.metadata}: @@ -260,6 +281,23 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error { return nil } +// rejectStream refuses an inbound stream that would exceed MaxStreams by +// sending a stream-level KindError to the peer. The control bit is set so the +// frame is appended immediately without blocking the reader goroutine on write +// backpressure (as with an abortive cancel). No stream is created. +func (m *Manager) rejectStream(sid uint64) { + err := maxStreamsExceeded.New("refused: at most %d concurrent streams", m.opts.MaxStreams) + // Message id 1: this is the first (and only) frame the peer receives on the + // refused stream, and its receive assembler expects ids to start at 1. + _ = m.wr.WriteFrame(drpcwire.Frame{ + ID: drpcwire.ID{Stream: sid, Message: 1}, + Kind: drpcwire.KindError, + Data: drpcwire.MarshalError(err), + Done: true, + Control: true, + }, nil) +} + // // manage streams // diff --git a/drpcmanager/max_streams_test.go b/drpcmanager/max_streams_test.go new file mode 100644 index 00000000..f0636e33 --- /dev/null +++ b/drpcmanager/max_streams_test.go @@ -0,0 +1,68 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcmanager + +import ( + "context" + "net" + "strings" + "testing" + + "github.com/zeebo/assert" + + "storj.io/drpc/drpctest" + "storj.io/drpc/drpcwire" +) + +// A server with MaxStreams set refuses an inbound stream beyond the cap with a +// stream-level error, without tearing down the connection or admitting the +// stream. +func TestManager_MaxStreamsRefusesExcessInbound(t *testing.T) { + ctx := drpctest.NewTracker(t) + defer ctx.Close() + + cconn, sconn := net.Pipe() + defer func() { _ = cconn.Close() }() + defer func() { _ = sconn.Close() }() + + cman := New(cconn, Client) + defer func() { _ = cman.Close() }() + sman := NewWithOptions(sconn, Server, Options{MaxStreams: 1}) + defer func() { _ = sman.Close() }() + + accepted := make(chan struct{}) + release := make(chan struct{}) + + // Server accepts exactly one stream and holds it open (so it keeps counting + // toward the cap) until released. + ctx.Run(func(ctx context.Context) { + s1, _, err := sman.NewServerStream(ctx) + assert.NoError(t, err) + close(accepted) + <-release + _ = s1.Close() + }) + + ctx.Run(func(ctx context.Context) { + // First stream is admitted. + c1, err := cman.NewClientStream(ctx, "rpc") + assert.NoError(t, err) + assert.NoError(t, c1.RawWrite(drpcwire.KindInvoke, []byte("rpc"))) + <-accepted // ensure the first stream is active before opening the second + + // Second stream: the server is at its cap, so it is refused with a + // stream-level error surfaced on the next receive. + c2, err := cman.NewClientStream(ctx, "rpc") + assert.NoError(t, err) + assert.NoError(t, c2.RawWrite(drpcwire.KindInvoke, []byte("rpc"))) + _, err = c2.RawRecv() + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "concurrent streams")) + + close(release) + _ = c1.Close() + }) + + ctx.Wait() +} From 0d215e29cddd17f0ccb7df8fbb7003b8f57afa5f Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Thu, 2 Jul 2026 10:34:25 +0530 Subject: [PATCH 10/11] drpcstream: reject messages larger than the implicit max size Flow control bounds a single message to high_water + window (the implicit maximum message size): a larger message can never complete -- the sender runs out of credit while the receiver's buffer sits full with nothing to drain -- so it would deadlock. rawWriteLocked now rejects such a message up front with a clear error rather than parking the sender on a message that can never finish. The bound is checked against the local FlowControl config, which under static symmetric enablement matches the peer's; a zero bound means unlimited (windows not sized). Defending against a non-compliant peer that sends an oversized message is a separate, receiver-side fail-stop (a follow-up). Co-Authored-By: roachdev-claude --- drpcstream/flow_control_option_test.go | 29 ++++++++++++++++++++++++++ drpcstream/stream.go | 13 ++++++++++++ 2 files changed, 42 insertions(+) diff --git a/drpcstream/flow_control_option_test.go b/drpcstream/flow_control_option_test.go index eac6dc40..a857a7f7 100644 --- a/drpcstream/flow_control_option_test.go +++ b/drpcstream/flow_control_option_test.go @@ -5,6 +5,7 @@ package drpcstream import ( "context" + "strings" "testing" "time" @@ -68,3 +69,31 @@ func TestStream_FlowControlOptionGatesAndResumes(t *testing.T) { t.Fatal("write did not resume after grant") } } + +// A message larger than the implicit maximum (high_water + window) is rejected +// up front rather than parking the sender on a message that can never complete. +func TestStream_FlowControlRejectsOversizedMessage(t *testing.T) { + mw := testMuxWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{ + SplitSize: 64 << 10, + FlowControl: FlowControl{Enabled: true, StreamWindow: 256, HighWater: 1024, GrantThreshold: 128}, + }) + // maxMsg = HighWater + StreamWindow = 1280. + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, make([]byte, 1281)) }() + + select { + case err := <-done: + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "exceeds")) + case <-time.After(time.Second): + t.Fatal("oversized message did not fail fast (it parked on credit instead)") + } +} + +// With flow control disabled there is no implicit message-size limit. +func TestStream_NoFlowControlNoMessageSizeLimit(t *testing.T) { + mw := testMuxWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + assert.NoError(t, st.RawWrite(drpcwire.KindMessage, make([]byte, 256<<10))) +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 786c754a..94af4ab1 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -464,6 +464,19 @@ func (s *Stream) RawWrite(kind drpcwire.Kind, data []byte) (err error) { // appropriate locks. // TODO(shubham): can we merge this with sendPacketLocked? func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) { + // Flow control bounds a single message to the implicit maximum of + // high_water + window: a larger message can never complete (the sender runs + // out of credit with the receiver's buffer full and nothing to drain), so + // reject it up front rather than parking on a message that would deadlock. + // This is checked against the local configuration, which under static + // symmetric enablement matches the peer's; a non-compliant peer is a + // separate, receiver-side concern. + if kind == drpcwire.KindMessage && s.sendw != nil { + if maxMsg := s.opts.FlowControl.HighWater + s.opts.FlowControl.StreamWindow; maxMsg > 0 && int64(len(data)) > maxMsg { + return errs.New("flow control: message of %d bytes exceeds the maximum of %d bytes (high_water + window)", len(data), maxMsg) + } + } + fr := s.newFrameLocked(kind) n := s.opts.SplitSize From 2f79115af9026419a5cae60a54d0202023b7fee5 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Thu, 2 Jul 2026 10:53:37 +0530 Subject: [PATCH 11/11] drpcstream: fail the stream on an oversized inbound message The sender-side guard rejects an oversized message locally, but a peer that does not honor the max-message bound (a non-compliant or older peer, or an asymmetric configuration where the sender's window is larger than ours) can still put one on the wire. Such a message can never complete: once it reaches the receiver's implicit maximum of high_water + window the sender is out of credit with our buffer full and nothing to drain, so the reassembly would stall. HandleFrame now tracks the in-progress inbound message size and, when it grows past the maximum, fails the stream with a data-overflow error and notifies the peer with an abortive cancel. The connection stays up (stream-level isolation). SendCancel is safe from the reader goroutine because it terminates -- waking any parked local sender and releasing the write lock -- before it takes that lock, unlike SendError. Co-Authored-By: roachdev-claude --- drpcmanager/flow_control_test.go | 59 +++++++++++++++++++++++++++ drpcstream/recv_window_wiring_test.go | 29 +++++++++++++ drpcstream/stream.go | 25 ++++++++++++ 3 files changed, 113 insertions(+) diff --git a/drpcmanager/flow_control_test.go b/drpcmanager/flow_control_test.go index f7c45b60..7710588b 100644 --- a/drpcmanager/flow_control_test.go +++ b/drpcmanager/flow_control_test.go @@ -8,6 +8,7 @@ import ( "errors" "io" "net" + "strings" "testing" "time" @@ -139,3 +140,61 @@ func TestManager_FlowControlContextCancelWakesParkedSend(t *testing.T) { t.Fatal("parked send did not wake on context cancellation") } } + +// A peer that does not honor the receiver's implicit maximum message size is +// caught on receipt: the server's window is small, and a client with a much +// larger window sends a message that exceeds the server's maximum. The server +// fails that stream with a data-overflow error while the connection stays up. +func TestManager_FlowControlReceiverRejectsOversizedMessage(t *testing.T) { + ctx := drpctest.NewTracker(t) + defer ctx.Close() + + cconn, sconn := net.Pipe() + defer func() { _ = cconn.Close() }() + defer func() { _ = sconn.Close() }() + + // The client's window is large enough to send the whole message immediately + // without waiting for grants and without tripping its own sender-side limit; + // the server's window is small, so the message exceeds the server's implicit + // maximum (high_water + window) and must be rejected on receipt. + cman := NewWithOptions(cconn, Client, Options{Stream: drpcstream.Options{ + SplitSize: 64 << 10, + FlowControl: drpcstream.FlowControl{ + Enabled: true, StreamWindow: 1 << 20, HighWater: 8 << 20, GrantThreshold: 256 << 10, + }, + }}) + defer func() { _ = cman.Close() }() + sman := NewWithOptions(sconn, Server, Options{Stream: drpcstream.Options{ + SplitSize: 64 << 10, + FlowControl: drpcstream.FlowControl{ + Enabled: true, StreamWindow: 64 << 10, HighWater: 64 << 10, GrantThreshold: 32 << 10, + }, + }}) + defer func() { _ = sman.Close() }() + + // 256 KiB exceeds the server's 128 KiB implicit maximum (64 KiB high-water + + // 64 KiB window). + msg := make([]byte, 256<<10) + + ctx.Run(func(ctx context.Context) { + stream, err := cman.NewClientStream(ctx, "rpc") + assert.NoError(t, err) + assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc"))) + // Best-effort: the send may complete (credit is ample) or fail once the + // server's cancel arrives; either way the server must reject the message. + _ = stream.RawWrite(drpcwire.KindMessage, msg) + _ = stream.Close() + }) + + ctx.Run(func(ctx context.Context) { + stream, _, err := sman.NewServerStream(ctx) + assert.NoError(t, err) + defer func() { _ = stream.Close() }() + + _, err = stream.RawRecv() + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "exceeds")) + }) + + ctx.Wait() +} diff --git a/drpcstream/recv_window_wiring_test.go b/drpcstream/recv_window_wiring_test.go index 85107199..7b96b4be 100644 --- a/drpcstream/recv_window_wiring_test.go +++ b/drpcstream/recv_window_wiring_test.go @@ -6,6 +6,7 @@ package drpcstream import ( "context" "io" + "strings" "testing" "time" @@ -140,3 +141,31 @@ func TestStream_NoWindowsNoGrants(t *testing.T) { assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 100))) } + +// The receiver defends against a peer that does not honor the max-message +// bound: an incoming message that grows past the implicit maximum (high_water + +// window) can never complete, so the stream is failed with a data-overflow +// error and the peer is notified with an abortive cancel. The connection stays +// up (HandleFrame returns nil). +func TestStream_ReceiverRejectsOversizedMessage(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{ + SplitSize: 64 << 10, + FlowControl: FlowControl{Enabled: true, StreamWindow: 256, HighWater: 1024, GrantThreshold: 128}, + }) + sid := st.ID() + // maxMsg = HighWater + StreamWindow = 1280. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, make([]byte, 1024), false))) + // This frame pushes the in-progress message to 1324 > 1280 without a Done, + // so it can never complete: the receiver must fail the stream. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, make([]byte, 300), false))) + + // The peer is notified with a cancel. + fr := waitFrame(t, frames) + assert.Equal(t, fr.Kind, drpcwire.KindCancel) + + // The local stream is failed with the data-overflow error. + _, err := st.RawRecv() + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "exceeds")) +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 94af4ab1..e1964016 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -88,6 +88,12 @@ type Stream struct { // flow control is not enabled, in which case no grants are emitted. recvw *recvWindow + // recvMsgBytes tracks the size of the in-progress inbound message so an + // oversized message (one the sender could never legitimately complete) can + // be detected. It is touched only from the reader goroutine (HandleFrame) + // and reset when a message completes. + recvMsgBytes int64 + mu sync.Mutex // protects state transitions sigs struct { send drpcsignal.Signal // set when done sending messages @@ -269,7 +275,26 @@ func (s *Stream) HandleFrame(fr drpcwire.Frame) (err error) { // A data frame dispatched off the wire may return credit to the sender. if fr.Kind == drpcwire.KindMessage && s.recvw != nil { + s.recvMsgBytes += int64(len(fr.Data)) + + // Receiver-side defense against a peer that does not honor the implicit + // maximum message size (high_water + window): such a message can never + // complete -- once it reaches the maximum the sender is out of credit + // with our buffer full -- so we would otherwise buffer up to the limit + // and then stall. Fail the stream with a data-overflow error and notify + // the peer with an abortive cancel; the connection stays up. SendCancel + // is safe from the reader goroutine because it terminates before taking + // the write lock. A zero bound means unlimited (windows not sized). + if maxMsg := s.opts.FlowControl.HighWater + s.opts.FlowControl.StreamWindow; maxMsg > 0 && + (s.recvMsgBytes > maxMsg || (s.recvMsgBytes == maxMsg && !fr.Done)) { + _ = s.SendCancel(errs.New("flow control: received message exceeds the maximum of %d bytes (high_water + window)", maxMsg)) + return nil + } + s.emitGrant(s.recvw.dispatched(int64(len(fr.Data)))) + if fr.Done { + s.recvMsgBytes = 0 + } } if !packetReady {