From 48cda986bb58be0e0f43c451c14216315238d578 Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Tue, 21 Jul 2026 19:53:20 +0530 Subject: [PATCH] drpcstream: gate data-frame sends on the per-stream send window Wire the sendWindow credit gate into rawWriteLocked: a KindMessage frame acquires len(frame) bytes of per-stream send credit before it is handed to the writer. Control frames (invoke/metadata) bypass the gate. The window is opt-in: a stream has no send window by default, so data writes stay ungated (unlimited) and behavior is unchanged until one is installed. terminate closes the window with the send-side error (sigs.send is first-wins, holding io.EOF when a cancel/error path pre-set it), so a send parked on credit returns the same error as one parked in WriteFrame or a later send. SendCancel and Cancel already terminate before taking the write lock, so they wake a credit-parked writer. acquire no longer takes a context: the stream's context Done channel is only closed once the stream is finished (after all ops complete), which cannot happen while an acquire is parked (the parked write is itself an in-flight op), so it never fired there. close is the sole abort path, and every termination route reaches it. Per-stream only. NOTE: Close/SendError/CloseSend still take the write lock while holding s.mu, so a send parked on credit can deadlock them; the lock-ordering rework that makes those paths preempt a parked writer is handled separately. Co-Authored-By: roachdev-claude --- drpcstream/send_window.go | 27 +++---- drpcstream/send_window_gate_test.go | 121 ++++++++++++++++++++++++++++ drpcstream/send_window_test.go | 66 +++------------ drpcstream/stream.go | 19 +++++ 4 files changed, 162 insertions(+), 71 deletions(-) create mode 100644 drpcstream/send_window_gate_test.go diff --git a/drpcstream/send_window.go b/drpcstream/send_window.go index 660a1090..bc65dc85 100644 --- a/drpcstream/send_window.go +++ b/drpcstream/send_window.go @@ -4,7 +4,6 @@ package drpcstream import ( - "context" "io" "math" "sync" @@ -35,23 +34,19 @@ func (w *sendWindow) available() int64 { } // acquire debits n bytes of credit, blocking until available, and returns nil. -// It returns early (consuming no credit) if the window is closed or ctx is -// canceled -func (w *sendWindow) acquire(ctx context.Context, n int64) error { +// It returns early (consuming no credit) with the close error if the window is +// closed. n <= 0 is a no-op unless the window is already closed. +func (w *sendWindow) acquire(n int64) error { for { w.mu.Lock() switch { case w.closed: + // Checked before n <= 0 so a closed window still fails an empty frame. err := w.err w.mu.Unlock() return err - case ctx.Err() != nil: - w.mu.Unlock() - return ctx.Err() case n <= 0: - // Nothing to acquire; after the terminal cases (so a closed/canceled - // window still fails) and never debits (so a negative n cannot add credit). - w.mu.Unlock() + w.mu.Unlock() // nothing to acquire; never debits, so a negative n adds no credit return nil case w.avail >= n: w.avail -= n @@ -60,19 +55,15 @@ func (w *sendWindow) acquire(ctx context.Context, n int64) error { } // Snapshot the notify channel under the lock before parking, so a grant // or close that fires the instant we unlock is not missed. Allocated - // here, by the first parker, so wakes with no waiters stay free. + // here, by the first parker, so wakes with no waiters stay free. The + // only wakeups are grant and close; abort is close's job (the stream + // terminates the window), so acquire needs no context. if w.notify == nil { w.notify = make(chan struct{}) } ch := w.notify w.mu.Unlock() - - select { - case <-ch: - // Credit was granted or the window closed; loop and re-check. - case <-ctx.Done(): - return ctx.Err() - } + <-ch // credit was granted or the window closed; loop and re-check } } diff --git a/drpcstream/send_window_gate_test.go b/drpcstream/send_window_gate_test.go new file mode 100644 index 00000000..3b275924 --- /dev/null +++ b/drpcstream/send_window_gate_test.go @@ -0,0 +1,121 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "errors" + "io" + "testing" + "time" + + "github.com/zeebo/assert" + "github.com/zeebo/errs" + + "storj.io/drpc/drpcwire" +) + +// newGateStream builds a stream writing to io.Discard with an explicit +// SplitSize so small payloads are a single frame. +func newGateStream(t *testing.T) *Stream { + mw := testMuxWriter(t) + return NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) +} + +// By default no send window is installed, so data writes are ungated +// (unlimited) and behavior is unchanged. +func TestStream_SendWindowDefaultUngated(t *testing.T) { + st := newGateStream(t) + assert.That(t, st.sendw == nil) + assert.NoError(t, st.RawWrite(drpcwire.KindMessage, []byte("hello"))) +} + +// With a finite send window, a data write blocks until enough credit is +// granted. +func TestStream_SendWindowGatesDataWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(4) // 4 bytes of credit + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("hello")) }() // 5 bytes > 4 + + select { + case <-done: + t.Fatal("data write returned before sufficient credit") + case <-time.After(blockShort): + } + + st.sendw.grant(1) // 4 + 1 = 5 >= 5 + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("data write did not complete after grant") + } +} + +// Control kinds (here, invoke) are not flow-controlled: they proceed even with +// zero send credit. +func TestStream_SendWindowControlKindsBypassGate(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // no credit at all + + assert.NoError(t, st.WriteInvoke("service.Method", nil)) +} + +// SendCancel preempts a send parked on credit: it terminates (closing the +// window) before taking the write lock, so the parked write wakes, releases the +// lock, and the cancel frame goes out. +func TestStream_SendWindowSendCancelPreemptsParkedWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // send will park immediately + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }() + + select { + case <-done: + t.Fatal("data write returned before cancel") + case <-time.After(blockShort): + } + + assert.NoError(t, st.SendCancel(context.Canceled)) + + select { + case err := <-done: + // Same error as a send parked in WriteFrame or a later send would see. + assert.That(t, errors.Is(err, io.EOF)) + case <-time.After(time.Second): + t.Fatal("parked data write was not preempted by SendCancel") + } + + // A subsequent send observes the same error as the parked one. + assert.That(t, errors.Is(st.RawWrite(drpcwire.KindMessage, []byte("more")), io.EOF)) +} + +// Terminating the stream wakes a send parked on credit. +func TestStream_SendWindowTerminateWakesParkedWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // send will park immediately + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }() + + select { + case <-done: + t.Fatal("data write returned before termination") + case <-time.After(blockShort): + } + + st.Cancel(errs.New("boom")) + + select { + case err := <-done: + // Cancel pre-sets sigs.send to io.EOF; the window closes with it. + assert.That(t, errors.Is(err, io.EOF)) + case <-time.After(time.Second): + t.Fatal("parked data write did not wake on termination") + } +} diff --git a/drpcstream/send_window_test.go b/drpcstream/send_window_test.go index 84c12a34..f1140dfe 100644 --- a/drpcstream/send_window_test.go +++ b/drpcstream/send_window_test.go @@ -4,7 +4,6 @@ package drpcstream import ( - "context" "errors" "io" "math" @@ -23,10 +22,10 @@ func TestSendWindowAcquireImmediate(t *testing.T) { w := newSendWindow(1000) assert.Equal(t, w.available(), int64(1000)) - assert.NoError(t, w.acquire(context.Background(), 400)) + assert.NoError(t, w.acquire(400)) assert.Equal(t, w.available(), int64(600)) - assert.NoError(t, w.acquire(context.Background(), 600)) + assert.NoError(t, w.acquire(600)) assert.Equal(t, w.available(), int64(0)) } @@ -36,14 +35,14 @@ func TestSendWindowGrantsAccumulate(t *testing.T) { w.grant(50) assert.Equal(t, w.available(), int64(150)) - assert.NoError(t, w.acquire(context.Background(), 150)) + assert.NoError(t, w.acquire(150)) assert.Equal(t, w.available(), int64(0)) } func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) { w := newSendWindow(100) done := make(chan error, 1) - go func() { done <- w.acquire(context.Background(), 300) }() + go func() { done <- w.acquire(300) }() // Not enough credit yet: acquire must block. select { @@ -63,30 +62,6 @@ func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) { assert.Equal(t, w.available(), int64(50)) } -func TestSendWindowAcquireContextCancel(t *testing.T) { - w := newSendWindow(0) - ctx, cancel := context.WithCancel(context.Background()) - done := make(chan error, 1) - go func() { done <- w.acquire(ctx, 100) }() - - select { - case <-done: - t.Fatal("acquire returned before cancellation") - case <-time.After(blockShort): - } - - cancel() - - select { - case err := <-done: - assert.That(t, errors.Is(err, context.Canceled)) - case <-time.After(time.Second): - t.Fatal("acquire did not wake on context cancellation") - } - // Credit was not consumed by a failed acquire. - assert.Equal(t, w.available(), int64(0)) -} - func TestSendWindowGrantSaturates(t *testing.T) { // Adding to a near-max balance saturates at MaxInt64 instead of wrapping. w := newSendWindow(math.MaxInt64 - 10) @@ -121,22 +96,11 @@ func TestSendWindowGrantSaturates(t *testing.T) { assert.Equal(t, w5.available(), int64(math.MaxInt64)) } -func TestSendWindowAcquireCanceledCtxWithCredit(t *testing.T) { - w := newSendWindow(1000) - ctx, cancel := context.WithCancel(context.Background()) - cancel() // already canceled, with credit available - - err := w.acquire(ctx, 100) - assert.That(t, errors.Is(err, context.Canceled)) - // A canceled context must not consume credit even though it was available. - assert.Equal(t, w.available(), int64(1000)) -} - func TestSendWindowCloseWakesAcquire(t *testing.T) { w := newSendWindow(0) closeErr := errs.New("terminated") done := make(chan error, 1) - go func() { done <- w.acquire(context.Background(), 100) }() + go func() { done <- w.acquire(100) }() select { case <-done: @@ -160,32 +124,28 @@ func TestSendWindowAcquireAfterClose(t *testing.T) { w.close(closeErr) // Even though credit is available, a closed window returns the close error. - assert.That(t, errors.Is(w.acquire(context.Background(), 1), closeErr)) + assert.That(t, errors.Is(w.acquire(1), closeErr)) } func TestSendWindowCloseNilError(t *testing.T) { w := newSendWindow(1000) w.close(nil) // closing with nil must not let a later acquire report success - assert.That(t, errors.Is(w.acquire(context.Background(), 1), io.EOF)) + assert.That(t, errors.Is(w.acquire(1), io.EOF)) } func TestSendWindowAcquireNonPositive(t *testing.T) { w := newSendWindow(100) - assert.NoError(t, w.acquire(context.Background(), 0)) - assert.NoError(t, w.acquire(context.Background(), -5)) + assert.NoError(t, w.acquire(0)) + assert.NoError(t, w.acquire(-5)) // Non-positive acquire consumes nothing; a negative one must not add credit. assert.Equal(t, w.available(), int64(100)) } -func TestSendWindowAcquireZeroObservesTerminalState(t *testing.T) { - // A zero-length acquire (empty frame) must still observe a canceled context... - ctx, cancel := context.WithCancel(context.Background()) - cancel() - assert.That(t, errors.Is(newSendWindow(100).acquire(ctx, 0), context.Canceled)) - - // ...and a closed window, rather than reporting success. +func TestSendWindowAcquireZeroObservesClose(t *testing.T) { + // A zero-length acquire (empty frame) must still observe a closed window + // rather than reporting success (closed is checked before n <= 0). w := newSendWindow(100) closeErr := errs.New("terminated") w.close(closeErr) - assert.That(t, errors.Is(w.acquire(context.Background(), 0), closeErr)) + assert.That(t, errors.Is(w.acquire(0), closeErr)) } diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 60e14096..ce9c73aa 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -61,6 +61,10 @@ type Stream struct { cbuf []byte // compression scratch buffer dbuf []byte // decompression scratch buffer + // sendw is the per-stream send-side flow-control window. It is nil when + // flow control is not enabled, in which case data writes are ungated. + sendw *sendWindow + mu sync.Mutex // protects state transitions sigs struct { send drpcsignal.Signal // set when done sending messages @@ -361,6 +365,12 @@ func (s *Stream) terminate(err error) { s.sigs.recv.Set(err) s.sigs.term.Set(err) s.recvQueue.Close(err) + if s.sendw != nil { + // Close with the send-side error: sigs.send is first-wins, so when a + // caller pre-set it (io.EOF for cancel/error), a send parked on credit + // returns the same error as one parked in WriteFrame or a later send. + s.sendw.close(s.sigs.send.Err()) + } s.checkFinished() } @@ -417,6 +427,15 @@ func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) { fr.Data, data = drpcwire.SplitData(data, n) fr.Done = len(data) == 0 + // Only data frames consume send credit; a nil window (flow control + // disabled) leaves sends ungated. acquire parks until credit arrives or + // the window closes (stream termination) -- the latter is the abort path. + if kind == drpcwire.KindMessage && s.sendw != nil { + if err := s.sendw.acquire(int64(len(fr.Data))); err != nil { + return err + } + } + drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(fr.Data))) s.log("SEND", fr.String)