diff --git a/drpcstream/recv_window.go b/drpcstream/recv_window.go new file mode 100644 index 0000000..c916de7 --- /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 0000000..3331c1c --- /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 0000000..7858eef --- /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 ce9c73a..84cc235 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 2cf7fae..f319eb9 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 41cf70d..8c19667 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