From 8d3dafc7d00c733ecbd61d08728bfbb2b4d58797 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 20:35:47 +0530 Subject: [PATCH 1/2] 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 | 62 ++++++++++++++++++++++++++++++++++ drpcstream/recv_window_test.go | 62 ++++++++++++++++++++++++++++++++++ 2 files changed, 124 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..c916de77 --- /dev/null +++ b/drpcstream/recv_window.go @@ -0,0 +1,62 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import "sync" + +// recvWindow decides when the receive side returns credit to the sender: it +// withholds grants while buffered bytes are at or above the high-water mark +// (credit keeps accruing), and otherwise releases accrued credit once it +// reaches the threshold, coalescing many frames into one grant. dispatched +// (reader goroutine) and consumed (application goroutine) return the credit +// delta to send as a KindWindowUpdate and may run concurrently. +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 seeded with highWater and 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). +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 accrued credit when the gate is open (buffered +// below high-water) and pending has reached the threshold. Requires w.mu. +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 3b297032336af5639ff9017b85e74582a4f74285 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Wed, 1 Jul 2026 20:45:38 +0530 Subject: [PATCH 2/2] drpcstream,drpcwire: 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. The packet assembler can silently drop an unfinished message when a higher message id supersedes it (a legitimate sequence: a send preempted mid-message, then a cancel). Those bytes were counted into the receive window at dispatch and would inflate buffered forever, wedging the high-water gate shut and deadlocking the sender once its credit ran out. The assembler now records the drop (TakeDiscarded) and HandleFrame releases the bytes as consumed, letting their credit flow back. The release runs last (deferred past the replacement frame's accounting and handlePacket), so the grant decision sees the final buffered state rather than the transient dip, and a terminal packet that supersedes a partial suppresses the grant via termination. emitGrant is also a no-op once the stream has terminated: the ring drains queued messages after close, and returning credit behind the terminal frame would invite the peer to send more doomed data. Co-Authored-By: roachdev-claude --- drpcstream/recv_window_wiring_test.go | 317 ++++++++++++++++++++++++++ drpcstream/stream.go | 66 +++++- drpcwire/packet_assembler.go | 23 +- drpcwire/packet_assembler_test.go | 63 +++++ 4 files changed, 464 insertions(+), 5 deletions(-) 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..7858eef0 --- /dev/null +++ b/drpcstream/recv_window_wiring_test.go @@ -0,0 +1,317 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "io" + "testing" + "time" + + "github.com/zeebo/assert" + "github.com/zeebo/errs" + + "storj.io/drpc" + "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))) +} + +// Nonconforming grant frames are dropped by the intercept without crediting. +func TestStream_MalformedGrantIgnored(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) + + missing := drpcwire.WindowUpdateFrame(st.ID(), 1) + missing.Data = nil // no delta payload + assert.NoError(t, st.HandleFrame(missing)) + assert.Equal(t, st.sendw.available(), int64(0)) + + zero := drpcwire.WindowUpdateFrame(st.ID(), 1) + zero.Data = drpcwire.AppendVarint(nil, 0) // zero delta + assert.NoError(t, st.HandleFrame(zero)) + assert.Equal(t, st.sendw.available(), int64(0)) +} + +// A grant arriving after termination is dropped. +func TestStream_GrantAfterTerminationDropped(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) + + st.Cancel(errs.New("boom")) + assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 100))) + assert.Equal(t, st.sendw.available(), int64(0)) +} + +// An unfinished message superseded by a higher message id must release its +// buffered bytes, or the high-water gate sticks shut and the withheld grant +// is never emitted (sender deadlock once its credit is exhausted). +func TestStream_DiscardedPartialMessageReleasesCredit(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + st.recvw = newRecvWindow(100, 50) + sid := st.ID() + + // Partial message 1 reaches the high-water mark: grants withheld. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, make([]byte, 120), false))) + assertNoFrame(t, frames) + + // Message 2 supersedes it; the discarded 120 bytes leave the buffer after + // the replacement is counted, the gate reopens, and all accrued credit + // (120 discarded + 10 replacement) flushes. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 2, make([]byte, 10), true))) + + _, delta, ok := drpcwire.ParseWindowUpdate(waitFrame(t, frames)) + assert.That(t, ok) + assert.Equal(t, delta, uint64(130)) + + // The completed message is intact and its own credit accounting holds. + got, err := st.RawRecv() + assert.NoError(t, err) + assert.Equal(t, len(got), 10) + assert.Equal(t, st.recvw.bufferedBytes(), int64(0)) +} + +// The discard release must not decide on the transient dip before the +// replacement frame is counted: a large replacement that lands above +// high-water withholds everything, including the discarded bytes' credit. +func TestStream_DiscardReleaseSeesReplacementFrame(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + st.recvw = newRecvWindow(100, 50) + sid := st.ID() + + // Partial message 1 pins buffered at the high-water mark. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, make([]byte, 120), false))) + assertNoFrame(t, frames) + + // Message 2 (1000 bytes) supersedes it. Final buffered = 1000 >= 100, so + // no grant may be emitted -- releasing the 120 on the dip would be wrong. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 2, make([]byte, 1000), true))) + assertNoFrame(t, frames) + + // Consuming the completed message reopens the gate and flushes everything. + _, err := st.RawRecv() + assert.NoError(t, err) + _, delta, ok := drpcwire.ParseWindowUpdate(waitFrame(t, frames)) + assert.That(t, ok) + assert.Equal(t, delta, uint64(1120)) +} + +// A terminal frame that supersedes a partial message must not flush the +// discarded bytes' credit: termination is applied first, and emitGrant +// no-ops on a terminated stream. +func TestStream_TerminalSupersedeEmitsNoGrant(t *testing.T) { + mw, frames := captureWriter(t) + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) + st.recvw = newRecvWindow(100, 50) + sid := st.ID() + + // Partial message 1: grants withheld above high-water. + assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, make([]byte, 120), false))) + assertNoFrame(t, frames) + + // A remote cancel with a higher message id discards the partial and + // terminates the stream; the withheld credit must stay unemitted. + assert.NoError(t, st.HandleFrame(drpcwire.Frame{ + ID: drpcwire.ID{Stream: sid, Message: 2}, Kind: drpcwire.KindCancel, Done: true, + })) + assert.That(t, st.IsTerminated()) + assertNoFrame(t, frames) +} + +// Draining queued messages after termination must not emit grants: credit +// returned behind the terminal frame would invite more doomed data. +func TestStream_NoGrantsAfterTermination(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: grant withheld while live. + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 120), true))) + assertNoFrame(t, frames) + + st.Cancel(errs.New("boom")) + + // The queued message still drains (ring close keeps buffered data + // readable), which on a live stream would flush the accrued 120 bytes of + // credit; after termination it must stay unemitted. + got, err := st.RawRecv() + assert.NoError(t, err) + assert.Equal(t, len(got), 120) + assertNoFrame(t, frames) +} + +// rawEnc round-trips raw byte slices for MsgRecv tests. +type rawEnc struct{} + +func (rawEnc) Marshal(msg drpc.Message) ([]byte, error) { return *(msg.(*[]byte)), nil } +func (rawEnc) Unmarshal(buf []byte, msg drpc.Message) error { + *(msg.(*[]byte)) = append([]byte(nil), buf...) + return nil +} + +// The MsgRecv consume path also returns credit, like RawRecv. +func TestStream_MsgRecvEmitsGrant(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 on dispatch. + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 120), true))) + assertNoFrame(t, frames) + + var msg []byte + assert.NoError(t, st.MsgRecv(&msg, rawEnc{})) + assert.Equal(t, len(msg), 120) + + _, delta, ok := drpcwire.ParseWindowUpdate(waitFrame(t, frames)) + assert.That(t, ok) + assert.Equal(t, delta, uint64(120)) +} + +// A message that fails to decompress still releases its bytes to the receive +// window. The bytes have left the queue, so buffered must not stay inflated -- +// otherwise the high-water gate could wedge credit shut after a bad frame. +func TestStream_DecompressErrorReleasesCredit(t *testing.T) { + mw := testMuxWriter(t) + opts := Options{SplitSize: 64 << 10, Compression: drpc.CompressionSnappy} + st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), opts) + st.recvw = newRecvWindow(1<<20, 1<<20) // high threshold: no grant to mask the release + + // A KindMessage payload that is not valid snappy. + bad := []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff} + assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, bad, true))) + assert.Equal(t, st.recvw.bufferedBytes(), int64(len(bad))) // dispatched, not yet consumed + + // RawRecv fails to decompress, but the bytes must still be released. + _, err := st.RawRecv() + assert.Error(t, err) + assert.Equal(t, st.recvw.bufferedBytes(), int64(0)) +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index ce9c73aa..84cc235f 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -65,6 +65,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 @@ -220,10 +224,38 @@ func (s *Stream) HandleFrame(fr drpcwire.Frame) (err error) { return nil } + // Grants are out-of-band signaling: intercept before the assembler so one + // arriving mid-message cannot disturb reassembly. Without a send window + // (flow control disabled) the grant is 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 } + + // Release bytes the assembler discarded (an unfinished message superseded + // by a higher id). Deferred so the gate decision sees the final buffered + // state and a terminal packet suppresses the grant (emitGrant checks term). + defer func() { + if kind, n := s.pa.TakeDiscarded(); n > 0 && kind == drpcwire.KindMessage && s.recvw != nil { + s.emitGrant(s.recvw.consumed(int64(n))) + } + }() + + // Accrue credit for a dispatched data frame; granted now or when the + // high-water gate reopens. + if fr.Kind == drpcwire.KindMessage && s.recvw != nil { + s.emitGrant(s.recvw.dispatched(int64(len(fr.Data)))) + } + if !packetReady { return nil } @@ -289,6 +321,17 @@ func (s *Stream) handlePacket(pkt drpcwire.Packet) (err error) { } } +// emitGrant sends a credit grant of delta bytes; no-op unless delta is +// positive and the stream is live (no credit behind a terminal frame). +// Control frames append without blocking, so the reader may call this. +// Best-effort: a write error means the stream is going away. +func (s *Stream) emitGrant(delta int64) { + if delta <= 0 || s.sigs.term.IsSet() { + return + } + _ = s.wr.WriteFrame(drpcwire.WindowUpdateFrame(s.id.Stream, uint64(delta)), nil) +} + // // helpers // @@ -463,6 +506,15 @@ func (s *Stream) RawRecv() (data []byte, err error) { return nil, err } defer s.recvQueue.Done() + n := len(b) // wire bytes, pre-decompression: matches dispatch accounting + + // The bytes have left the queue; release them to the receive window on every + // return path (including a decode failure below), or buffered stays inflated + // and the high-water gate can wedge credit shut. Consuming may also reopen + // the gate, releasing credit accrued at dispatch. + if s.recvw != nil { + defer func() { s.emitGrant(s.recvw.consumed(int64(n))) }() + } if s.opts.Compression != drpc.CompressionNone { s.dbuf, err = drpcwire.Decompress(s.opts.Compression, s.dbuf[:0], b) @@ -472,7 +524,6 @@ func (s *Stream) RawRecv() (data []byte, err error) { b = s.dbuf } data = append([]byte(nil), b...) - return data, nil } @@ -515,6 +566,15 @@ func (s *Stream) MsgRecv(msg drpc.Message, enc drpc.Encoding) (err error) { return err } defer s.recvQueue.Done() + n := len(b) // wire bytes, pre-decompression: matches dispatch accounting + + // The bytes have left the queue; release them to the receive window on every + // return path (including a decode failure below), or buffered stays inflated + // and the high-water gate can wedge credit shut. Consuming may also reopen + // the gate, releasing credit accrued at dispatch. + if s.recvw != nil { + defer func() { s.emitGrant(s.recvw.consumed(int64(n))) }() + } if s.opts.Compression != drpc.CompressionNone { s.dbuf, err = drpcwire.Decompress(s.opts.Compression, s.dbuf[:0], b) @@ -523,9 +583,7 @@ func (s *Stream) MsgRecv(msg drpc.Message, enc drpc.Encoding) (err error) { } b = s.dbuf } - err = enc.Unmarshal(b, msg) - - return err + return enc.Unmarshal(b, msg) } // diff --git a/drpcwire/packet_assembler.go b/drpcwire/packet_assembler.go index 2cf7fae3..f319eb93 100644 --- a/drpcwire/packet_assembler.go +++ b/drpcwire/packet_assembler.go @@ -16,6 +16,9 @@ type PacketAssembler struct { pk Packet assembling bool streamInitialized bool + + discardedKind Kind + discardedLen int } // NewPacketAssembler returns a new PacketAssembler ready to assemble frames. @@ -41,12 +44,26 @@ func (pa *PacketAssembler) Reset() { } pa.assembling = false pa.streamInitialized = false + pa.discardedKind, pa.discardedLen = 0, 0 +} + +// TakeDiscarded returns the kind and payload byte count of an unfinished +// message dropped by the most recent AppendFrame (a higher message id +// superseded it), and clears the record. n is 0 when nothing was dropped. +func (pa *PacketAssembler) TakeDiscarded() (kind Kind, n int) { + kind, n = pa.discardedKind, pa.discardedLen + pa.discardedKind, pa.discardedLen = 0, 0 + return kind, n } // AppendFrame adds a frame to the in-progress packet. It returns the completed // packet and true when a frame with Done=true is received. It returns false // when more frames are needed to complete the packet. func (pa *PacketAssembler) AppendFrame(fr Frame) (packet Packet, packetReady bool, err error) { + // A discard record is scoped to the most recent AppendFrame (see + // TakeDiscarded); clear any stale one before possibly setting a fresh one. + pa.discardedKind, pa.discardedLen = 0, 0 + // Enforce stream ID consistency: infer from first frame or reject mismatches. if !pa.streamInitialized { pa.pk.ID.Stream = fr.ID.Stream @@ -60,7 +77,11 @@ func (pa *PacketAssembler) AppendFrame(fr Frame) (packet Packet, packetReady boo return Packet{}, false, drpc.ProtocolError.New( "message id monotonicity violation: got %v, expected >= %v", fr.ID.Message, pa.pk.ID.Message) } else if fr.ID.Message > pa.pk.ID.Message || !pa.assembling { - // New message: reset the buffer and start assembling. + // New message: reset and start assembling. Record any dropped + // unfinished bytes so byte-accounting callers can release them. + if pa.assembling && len(pa.pk.Data) > 0 { + pa.discardedKind, pa.discardedLen = pa.pk.Kind, len(pa.pk.Data) + } pa.pk.Data = pa.pk.Data[:0] pa.assembling = true pa.pk.ID.Message = fr.ID.Message diff --git a/drpcwire/packet_assembler_test.go b/drpcwire/packet_assembler_test.go index 41cf70da..8c19667a 100644 --- a/drpcwire/packet_assembler_test.go +++ b/drpcwire/packet_assembler_test.go @@ -87,6 +87,69 @@ func TestPacketAssembler_HigherMsgDiscardsInProgress(t *testing.T) { assert.NoError(t, err) assert.That(t, ready) assert.DeepEqual(t, pkt.Data, []byte("kept")) + + // The dropped bytes are reported once so byte accounting can release them. + kind, n := pa.TakeDiscarded() + assert.Equal(t, kind, KindMessage) + assert.Equal(t, n, len("discard")) + _, n = pa.TakeDiscarded() + assert.Equal(t, n, 0) +} + +// A completed message followed by the next message id is not a discard. +func TestPacketAssembler_CompletedMessageNotReportedDiscarded(t *testing.T) { + pa := NewPacketAssembler() + pa.SetStreamID(1) + + _, ready, err := pa.AppendFrame(Frame{ + ID: ID{Stream: 1, Message: 1}, Kind: KindMessage, Data: []byte("done"), Done: true, + }) + assert.NoError(t, err) + assert.That(t, ready) + + _, ready, err = pa.AppendFrame(Frame{ + ID: ID{Stream: 1, Message: 2}, Kind: KindMessage, Data: []byte("next"), Done: true, + }) + assert.NoError(t, err) + assert.That(t, ready) + + _, n := pa.TakeDiscarded() + assert.Equal(t, n, 0) +} + +// Discard state is scoped to the most recent AppendFrame: a later +// non-discarding append clears an undrained record, and Reset drops it too. +func TestPacketAssembler_DiscardStateNotStale(t *testing.T) { + pa := NewPacketAssembler() + pa.SetStreamID(1) + + _, _, err := pa.AppendFrame(Frame{ + ID: ID{Stream: 1, Message: 1}, Kind: KindMessage, Data: []byte("old"), Done: false, + }) + assert.NoError(t, err) + _, _, err = pa.AppendFrame(Frame{ // supersedes m1: discard recorded + ID: ID{Stream: 1, Message: 2}, Kind: KindMessage, Data: []byte("a"), Done: false, + }) + assert.NoError(t, err) + _, _, err = pa.AppendFrame(Frame{ // non-discarding: undrained record cleared + ID: ID{Stream: 1, Message: 2}, Kind: KindMessage, Data: []byte("b"), Done: true, + }) + assert.NoError(t, err) + _, n := pa.TakeDiscarded() + assert.Equal(t, n, 0) + + // Reset clears a pending record as well. + _, _, err = pa.AppendFrame(Frame{ + ID: ID{Stream: 1, Message: 3}, Kind: KindMessage, Data: []byte("old"), Done: false, + }) + assert.NoError(t, err) + _, _, err = pa.AppendFrame(Frame{ + ID: ID{Stream: 1, Message: 4}, Kind: KindMessage, Data: []byte("new"), Done: true, + }) + assert.NoError(t, err) + pa.Reset() + _, n = pa.TakeDiscarded() + assert.Equal(t, n, 0) } // Continuation frames (same message ID, mid-assembly) must carry the same