drpcstream: install per-stream windows from an internal flow-control option

Streams install their send and receive windows from a FlowControl option
carried in the internal (drpcopts) stream options. Keeping enablement
internal-only makes accidental activation impossible for module
consumers: bumping the drpc dependency mid-implementation cannot turn
flow control on, because no public API reaches it. A later change
promotes the option to the public Options once the version-gated
enablement in CockroachDB is ready.

Construction validates the configuration and panics on one that cannot
make progress (silently disabling would drop the protection instead):
all sizes positive, SplitSize bounded (negative means unsplit frames, so
no bound holds; zero uses SplitData's 64 KiB default), and
GrantThreshold + frame size <= StreamWindow -- up to GrantThreshold-1
bytes stay withheld at quiescence, and the remaining window must still
cover the sender's next frame.

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

Author: suj-krishnan

drpcstream/flow_control_option_test.go

@@ -0,0 +1,119 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc/drpcwire"
+	"storj.io/drpc/internal/drpcopts"
+)
+
+// fcOptions returns stream Options with flow control enabled via the internal
+// option, the only way to enable it until it is promoted to a public option.
+func fcOptions(window, highWater, threshold int64) Options {
+	opts := Options{SplitSize: 64 << 10}
+	drpcopts.SetStreamFlowControl(&opts.Internal, drpcopts.FlowControl{
+		Enabled:        true,
+		StreamWindow:   window,
+		HighWater:      highWater,
+		GrantThreshold: threshold,
+	})
+	return opts
+}
+
+// Enabling flow control 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(), fcOptions(256<<10, 4<<20, 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)
+}
+
+// An invalid flow-control configuration panics at construction rather than
+// deadlocking later or silently disabling the protection.
+func TestStream_FlowControlValidation(t *testing.T) {
+	mustPanic := func(name string, opts Options) {
+		t.Helper()
+		defer func() {
+			if recover() == nil {
+				t.Fatalf("%s: expected construction to panic", name)
+			}
+		}()
+		NewWithOptions(context.Background(), 1, testMuxWriter(t), NewBufferPool(), opts)
+	}
+
+	mustPanic("zero window", fcOptions(0, 4<<20, 128<<10))
+	mustPanic("zero high water", fcOptions(256<<10, 0, 128<<10))
+	mustPanic("zero threshold", fcOptions(256<<10, 4<<20, 0))
+	mustPanic("negative window", fcOptions(-1, 4<<20, 128<<10))
+
+	// Threshold + frame must fit in the window, or a quiescent receiver can
+	// strand the sender below the next frame's cost.
+	mustPanic("threshold+frame exceeds window", fcOptions(128<<10, 4<<20, 128<<10))
+
+	// A near-max threshold must not sneak past via int64 overflow of the
+	// threshold+frame sum.
+	mustPanic("near-max threshold overflow", fcOptions(256<<10, 4<<20, math.MaxInt64))
+
+	// SplitSize < 0 means unsplit (unbounded) frames: no bound holds.
+	unbounded := fcOptions(256<<10, 4<<20, 64<<10)
+	unbounded.SplitSize = -1
+	mustPanic("unbounded SplitSize", unbounded)
+
+	// SplitSize 0 uses SplitData's 64 KiB default as the frame size.
+	def := fcOptions(128<<10, 4<<20, 64<<10)
+	def.SplitSize = 0
+	NewWithOptions(context.Background(), 1, testMuxWriter(t), NewBufferPool(), def) // 64K+64K <= 128K: ok
+	def2 := fcOptions(96<<10, 4<<20, 64<<10)
+	def2.SplitSize = 0
+	mustPanic("default frame exceeds window headroom", def2)
+}
+
+// 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. SplitSize
+// 2 with a 4-byte window means "hello" (frames 2+2+1) parks on its last frame.
+func TestStream_FlowControlOptionGatesAndResumes(t *testing.T) {
+	mw := testMuxWriter(t)
+	opts := Options{SplitSize: 2}
+	drpcopts.SetStreamFlowControl(&opts.Internal, drpcopts.FlowControl{
+		Enabled: true, StreamWindow: 4, HighWater: 1 << 20, GrantThreshold: 2,
+	})
+	st := NewWithOptions(context.Background(), 1, mw, NewBufferPool(), opts)
+
+	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")
+	}
+}

drpcstream/stream.go

@@ -34,10 +34,33 @@ type Options struct {
 	// All KindMessage data is compressed on send and decompressed on receive.
 	Compression drpc.Compression
 
-	// Internal contains options that are for internal use only.
+	// Internal contains options that are for internal use only. Per-stream
+	// flow control is configured here (drpcopts.SetStreamFlowControl) and is
+	// internal-only until the version-gated enablement work is ready.
 	Internal drpcopts.Stream
 }
 
+// validateFlowControl panics on a flow-control configuration that cannot make
+// progress: an invalid window is a programmer error, and disabling silently
+// would silently drop the protection. The frame size is SplitSize, or
+// SplitData's 64 KiB default when SplitSize is 0.
+func validateFlowControl(fc drpcopts.FlowControl, splitSize int) {
+	frame := int64(splitSize)
+	if splitSize == 0 {
+		frame = 64 << 10
+	}
+	switch {
+	case splitSize < 0:
+		panic("drpcstream: flow control requires a bounded SplitSize")
+	case fc.StreamWindow <= 0 || fc.HighWater <= 0 || fc.GrantThreshold <= 0:
+		panic(fmt.Sprintf("drpcstream: flow-control sizes must be positive: %+v", fc))
+	case frame > fc.StreamWindow || fc.GrantThreshold > fc.StreamWindow-frame:
+		panic(fmt.Sprintf(
+			"drpcstream: GrantThreshold (%d) + frame size (%d) exceeds StreamWindow (%d)",
+			fc.GrantThreshold, frame, fc.StreamWindow))
+	}
+}
+
 // Stream represents an rpc actively happening on a transport.
 type Stream struct {
 	ctx  streamCtx
@@ -132,6 +155,15 @@ func NewWithOptions(
 		onDequeue: drpcopts.GetStreamOnReceiveQueueDequeue(&opts.Internal),
 	})
 
+	// 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 fc := drpcopts.GetStreamFlowControl(&opts.Internal); fc.Enabled {
+		validateFlowControl(fc, opts.SplitSize)
+		s.sendw = newSendWindow(fc.StreamWindow)
+		s.recvw = newRecvWindow(fc.HighWater, fc.GrantThreshold)
+	}
+
 	return s
 }
 

internal/drpcopts/stream.go

@@ -17,8 +17,34 @@ type Stream struct {
 	stats                 *drpcstats.Stats
 	onReceiveQueueEnqueue func(int64)
 	onReceiveQueueDequeue func(int64)
+	flowControl           FlowControl
 }
 
+// FlowControl configures per-stream flow control. Internal-only until the
+// CockroachDB version-gated enablement is ready; a later change promotes it to
+// a public option. The installation site validates it (see drpcstream).
+type FlowControl struct {
+	// Enabled turns per-stream flow control on.
+	Enabled bool
+
+	// StreamWindow is the sender's initial and nominal per-stream credit.
+	StreamWindow int64
+
+	// HighWater is the receive-side buffered-byte mark at or above which
+	// grants are withheld.
+	HighWater int64
+
+	// GrantThreshold is the credit that must accrue before a grant is emitted,
+	// coalescing many frames into one KindWindowUpdate.
+	GrantThreshold int64
+}
+
+// GetStreamFlowControl returns the FlowControl stored in the options.
+func GetStreamFlowControl(opts *Stream) FlowControl { return opts.flowControl }
+
+// SetStreamFlowControl sets the FlowControl stored in the options.
+func SetStreamFlowControl(opts *Stream, fc FlowControl) { opts.flowControl = fc }
+
 // GetStreamTransport returns the drpc.Transport stored in the options.
 func GetStreamTransport(opts *Stream) drpc.Transport { return opts.transport }