diff --git a/drpcmanager/flow_control_test.go b/drpcmanager/flow_control_test.go new file mode 100644 index 0000000..be8ed78 --- /dev/null +++ b/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()) +} diff --git a/drpcstream/flow_control_option_test.go b/drpcstream/flow_control_option_test.go new file mode 100644 index 0000000..702f485 --- /dev/null +++ b/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") + } +} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 84cc235..e87ba7e 100644 --- a/drpcstream/stream.go +++ b/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 } diff --git a/internal/drpcopts/stream.go b/internal/drpcopts/stream.go index d0eb29d..be7b2f7 100644 --- a/internal/drpcopts/stream.go +++ b/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 }