From 6649f347402d94f6b2a9feec7bf69073defe4817 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. 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_gate_test.go | 121 ++++++++++++++++++++++++++++ drpcstream/stream.go | 19 +++++ 2 files changed, 140 insertions(+) create mode 100644 drpcstream/send_window_gate_test.go diff --git a/drpcstream/send_window_gate_test.go b/drpcstream/send_window_gate_test.go new file mode 100644 index 0000000..3b27592 --- /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/stream.go b/drpcstream/stream.go index 60e1409..031668c 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, + // the ctx is canceled, or the window closes (stream termination). + if kind == drpcwire.KindMessage && s.sendw != nil { + if err := s.sendw.acquire(s.Context(), int64(len(fr.Data))); err != nil { + return err + } + } + drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(fr.Data))) s.log("SEND", fr.String)