drpcstream: wire the receive-window grant policy into the read path

Install a per-stream recvWindow and drive it from the read/recv path:

  - HandleFrame intercepts an incoming KindWindowUpdate before the
    assembler and applies it to the send window. Intercepting early keeps
    a grant that interleaves with an in-progress message (grants are
    emitted off the write lock) from disturbing reassembly.
  - A dispatched KindMessage frame returns credit via recvWindow, emitting
    a KindWindowUpdate once the accrued amount crosses the threshold and
    the high-water gate is open.
  - RawRecv/MsgRecv call recvWindow after dequeuing; consuming can reopen a
    gate the high-water mark had closed and flush the accrued grant.

Grants are control frames, appended without blocking, so emitting them
from the reader goroutine is safe. The window is opt-in: with no recvw
installed no grants are emitted, and an incoming grant with no send
window is ignored, so default behavior is unchanged.

Stream-level only; the connection-level receive budget and the enablement
path that installs the windows come in later commits.

Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Author: suj-krishnan

drpcstream/recv_window_wiring_test.go

@@ -0,0 +1,142 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc/drpcwire"
+)
+
+// captureWriter returns a MuxWriter whose output frames are decoded and
+// delivered on the returned channel, so tests can observe emitted grants.
+func captureWriter(t *testing.T) (*drpcwire.MuxWriter, <-chan drpcwire.Frame) {
+	t.Helper()
+	pr, pw := io.Pipe()
+	mw := drpcwire.NewMuxWriter(pw, func(error) {})
+	frames := make(chan drpcwire.Frame, 64)
+	rd := drpcwire.NewReader(pr)
+	go func() {
+		for {
+			fr, err := rd.ReadFrame()
+			if err != nil {
+				return
+			}
+			frames <- fr
+		}
+	}()
+	t.Cleanup(func() {
+		mw.Stop(nil)
+		<-mw.Done()
+		_ = pw.Close()
+		_ = pr.Close()
+	})
+	return mw, frames
+}
+
+func waitFrame(t *testing.T, frames <-chan drpcwire.Frame) drpcwire.Frame {
+	t.Helper()
+	select {
+	case fr := <-frames:
+		return fr
+	case <-time.After(time.Second):
+		t.Fatal("expected a frame to be written, got none")
+		return drpcwire.Frame{}
+	}
+}
+
+func assertNoFrame(t *testing.T, frames <-chan drpcwire.Frame) {
+	t.Helper()
+	select {
+	case fr := <-frames:
+		t.Fatalf("unexpected frame written: %v", fr)
+	case <-time.After(blockShort):
+	}
+}
+
+func msgFrame(sid, mid uint64, data []byte, done bool) drpcwire.Frame {
+	return drpcwire.Frame{
+		ID:   drpcwire.ID{Stream: sid, Message: mid},
+		Kind: drpcwire.KindMessage,
+		Data: data,
+		Done: done,
+	}
+}
+
+// An incoming KindWindowUpdate is intercepted and applied to the send window.
+func TestStream_IncomingGrantAppliesToSendWindow(t *testing.T) {
+	st := newGateStream(t)
+	st.sendw = newSendWindow(0)
+
+	assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 500)))
+	assert.Equal(t, st.sendw.available(), int64(500))
+}
+
+// A grant interleaved with an in-progress message is intercepted before the
+// assembler, so reassembly is undisturbed.
+func TestStream_GrantDoesNotDisturbReassembly(t *testing.T) {
+	st := newGateStream(t)
+	sid := st.ID()
+
+	assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, []byte("foo"), false)))
+	assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(sid, 100))) // interleaved
+	assert.NoError(t, st.HandleFrame(msgFrame(sid, 1, []byte("bar"), true)))
+
+	got, err := st.RawRecv()
+	assert.NoError(t, err)
+	assert.Equal(t, string(got), "foobar")
+}
+
+// Dispatching a data frame returns credit once the threshold is met.
+func TestStream_DispatchEmitsGrant(t *testing.T) {
+	mw, frames := captureWriter(t)
+	st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10})
+	st.recvw = newRecvWindow(1<<20, 100)
+
+	assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 150), true)))
+
+	fr := waitFrame(t, frames)
+	sid, delta, ok := drpcwire.ParseWindowUpdate(fr)
+	assert.That(t, ok)
+	assert.Equal(t, sid, st.ID())
+	assert.Equal(t, delta, uint64(150))
+}
+
+// Above the high-water mark, dispatch withholds; consuming resumes and flushes
+// the accrued credit.
+func TestStream_ConsumeEmitsGrantOnResume(t *testing.T) {
+	mw, frames := captureWriter(t)
+	st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10})
+	st.recvw = newRecvWindow(100, 50)
+
+	// 120 bytes buffered >= high-water 100 -> withheld.
+	assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 120), true)))
+	assertNoFrame(t, frames)
+
+	// Consume -> buffered drops below high-water -> flush accrued 120.
+	_, err := st.RawRecv()
+	assert.NoError(t, err)
+
+	fr := waitFrame(t, frames)
+	_, delta, ok := drpcwire.ParseWindowUpdate(fr)
+	assert.That(t, ok)
+	assert.Equal(t, delta, uint64(120))
+}
+
+// With no receive window installed, no grants are emitted and an incoming grant
+// with no send window is harmlessly ignored.
+func TestStream_NoWindowsNoGrants(t *testing.T) {
+	mw, frames := captureWriter(t)
+	st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10})
+
+	assert.NoError(t, st.HandleFrame(msgFrame(st.ID(), 1, make([]byte, 150), true)))
+	assertNoFrame(t, frames)
+
+	assert.NoError(t, st.HandleFrame(drpcwire.WindowUpdateFrame(st.ID(), 100)))
+}

drpcstream/stream.go

@@ -59,6 +59,10 @@ type Stream struct {
 	// flow control is not enabled, in which case data writes are ungated.
 	sendw *sendWindow
 
+	// recvw is the per-stream receive-side flow-control policy. It is nil when
+	// flow control is not enabled, in which case no grants are emitted.
+	recvw *recvWindow
+
 	mu   sync.Mutex // protects state transitions
 	sigs struct {
 		send drpcsignal.Signal // set when done sending messages
@@ -211,10 +215,30 @@ func (s *Stream) HandleFrame(fr drpcwire.Frame) (err error) {
 		return nil
 	}
 
+	// Flow-control grants are out-of-band signaling, not part of the message
+	// stream, so they are intercepted before the assembler. Grants are emitted
+	// off the write lock, so one can arrive interleaved with an in-progress
+	// message; intercepting here keeps it from disturbing reassembly. Without a
+	// send window (flow control not enabled) the grant is simply ignored.
+	if fr.Kind == drpcwire.KindWindowUpdate {
+		if s.sendw != nil {
+			if _, delta, ok := drpcwire.ParseWindowUpdate(fr); ok {
+				s.sendw.grant(int64(delta))
+			}
+		}
+		return nil
+	}
+
 	packet, packetReady, err := s.pa.AppendFrame(fr)
 	if err != nil {
 		return err
 	}
+
+	// A data frame dispatched off the wire may return credit to the sender.
+	if fr.Kind == drpcwire.KindMessage && s.recvw != nil {
+		s.emitGrant(s.recvw.dispatched(int64(len(fr.Data))))
+	}
+
 	if !packetReady {
 		return nil
 	}
@@ -280,6 +304,18 @@ func (s *Stream) handlePacket(pkt drpcwire.Packet) (err error) {
 	}
 }
 
+// emitGrant sends a per-stream flow-control credit grant of delta bytes to the
+// peer, and is a no-op when delta is not positive. Grants are control frames,
+// which WriteFrame appends immediately without blocking on backpressure, so
+// this is safe to call from the reader goroutine. It is best-effort: a write
+// error means the stream is going away.
+func (s *Stream) emitGrant(delta int64) {
+	if delta <= 0 {
+		return
+	}
+	_ = s.wr.WriteFrame(drpcwire.WindowUpdateFrame(s.id.Stream, uint64(delta)), nil)
+}
+
 //
 // helpers
 //
@@ -447,8 +483,13 @@ func (s *Stream) RawRecv() (data []byte, err error) {
 		return nil, err
 	}
 	data = append([]byte(nil), b...)
+	n := len(b)
 	s.recvQueue.Done()
 
+	// Consuming buffered data can return credit to the sender.
+	if s.recvw != nil {
+		s.emitGrant(s.recvw.consumed(int64(n)))
+	}
 	return data, nil
 }
 
@@ -488,9 +529,14 @@ func (s *Stream) MsgRecv(msg drpc.Message, enc drpc.Encoding) (err error) {
 	if err != nil {
 		return err
 	}
+	n := len(b)
 	err = enc.Unmarshal(b, msg)
 	s.recvQueue.Done()
 
+	// Consuming buffered data can return credit to the sender.
+	if s.recvw != nil {
+		s.emitGrant(s.recvw.consumed(int64(n)))
+	}
 	return err
 }