drpcstream,drpcmanager: enable per-stream flow control (static)

drpcmanager/flow_control_test.go

@@ -0,0 +1,218 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc"
+	"storj.io/drpc/drpcstream"
+	"storj.io/drpc/drpctest"
+	"storj.io/drpc/drpcwire"
+	"storj.io/drpc/internal/drpcopts"
+)
+
+// fcManagerOptions returns manager Options whose streams have flow control
+// enabled via the internal stream option.
+func fcManagerOptions(window, highWater, threshold int64) Options {
+	stream := drpcstream.Options{SplitSize: 64 << 10}
+	drpcopts.SetStreamFlowControl(&stream.Internal, drpcopts.FlowControl{
+		Enabled:        true,
+		StreamWindow:   window,
+		HighWater:      highWater,
+		GrantThreshold: threshold,
+	})
+	return Options{Stream: stream}
+}
+
+// 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() }()
+
+	// 128 KiB window = 2 frames of initial credit.
+	opts := fcManagerOptions(128<<10, 1<<20, 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", drpc.CompressionNone)
+		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, fcManagerOptions(128<<10, 1<<20, 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", drpc.CompressionNone)
+	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")
+	}
+}
+
+// sniffConn tees every byte read off the connection into w, where a frame
+// parser classifies it.
+type sniffConn struct {
+	net.Conn
+	w io.Writer
+}
+
+func (s sniffConn) Read(p []byte) (int, error) {
+	n, err := s.Conn.Read(p)
+	if n > 0 {
+		_, _ = s.w.Write(p[:n])
+	}
+	return n, err
+}
+
+// Tripwire for accidental enablement: a default-configuration connection must
+// never put a KindWindowUpdate on the wire. Both directions are sniffed; the
+// KindMessage assertion proves the sniffer actually parsed the traffic.
+func TestManager_DefaultConfigEmitsNoWindowUpdates(t *testing.T) {
+	tr := drpctest.NewTracker(t)
+	defer tr.Close()
+
+	var sawMessage, sawWindowUpdate atomic.Bool
+	sniff := func(c net.Conn) net.Conn {
+		pr, pw := io.Pipe()
+		go func() {
+			rd := drpcwire.NewReader(pr)
+			for {
+				fr, err := rd.ReadFrame()
+				if err != nil {
+					_, _ = io.Copy(io.Discard, pr) // keep the tee unblocked
+					return
+				}
+				switch fr.Kind {
+				case drpcwire.KindMessage:
+					sawMessage.Store(true)
+				case drpcwire.KindWindowUpdate:
+					sawWindowUpdate.Store(true)
+				}
+			}
+		}()
+		t.Cleanup(func() { _ = pw.Close(); _ = pr.Close() })
+		return sniffConn{Conn: c, w: pw}
+	}
+
+	cconn, sconn := net.Pipe()
+	defer func() { _ = cconn.Close() }()
+	defer func() { _ = sconn.Close() }()
+
+	cman := NewWithOptions(sniff(cconn), Client, Options{})
+	defer func() { _ = cman.Close() }()
+	sman := NewWithOptions(sniff(sconn), Server, Options{})
+	defer func() { _ = sman.Close() }()
+
+	tr.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), 256<<10)
+		_, err = stream.RawRecv()
+		assert.That(t, errors.Is(err, io.EOF))
+	})
+
+	stream, err := cman.NewClientStream(context.Background(), "rpc", drpc.CompressionNone)
+	assert.NoError(t, err)
+	assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
+	assert.NoError(t, stream.RawWrite(drpcwire.KindMessage, make([]byte, 256<<10)))
+	assert.NoError(t, stream.Close())
+	tr.Wait()
+
+	assert.That(t, sawMessage.Load())
+	assert.That(t, !sawWindowUpdate.Load())
+}

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
 }