diff --git a/drpcmanager/flow_control_test.go b/drpcmanager/flow_control_test.go new file mode 100644 index 00000000..7710588b --- /dev/null +++ b/drpcmanager/flow_control_test.go @@ -0,0 +1,200 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcmanager + +import ( + "context" + "errors" + "io" + "net" + "strings" + "testing" + "time" + + "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() +} + +// 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") + } +} + +// 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/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() +} diff --git a/drpcstream/flow_control_option_test.go b/drpcstream/flow_control_option_test.go new file mode 100644 index 00000000..a857a7f7 --- /dev/null +++ b/drpcstream/flow_control_option_test.go @@ -0,0 +1,99 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "strings" + "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") + } +} + +// 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/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)) +} diff --git a/drpcstream/recv_window_wiring_test.go b/drpcstream/recv_window_wiring_test.go new file mode 100644 index 00000000..7b96b4be --- /dev/null +++ b/drpcstream/recv_window_wiring_test.go @@ -0,0 +1,171 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "io" + "strings" + "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))) +} + +// 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/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_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/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)) +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index fe7b0b48..e1964016 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 @@ -55,6 +80,20 @@ 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 + + // 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 + + // 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 @@ -115,6 +154,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 } @@ -207,10 +254,49 @@ 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.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 { return nil } @@ -276,6 +362,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 // @@ -352,6 +450,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() } @@ -388,6 +489,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 @@ -405,6 +519,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 @@ -429,8 +554,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 } @@ -470,9 +600,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 } 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) + } + } +}