drpcstream: gate data-frame sends on the per-stream send window

Wire the sendWindow credit gate into rawWriteLocked: a KindMessage frame
now acquires len(frame) bytes of per-stream send credit before it is
handed to the writer. Control frames (invoke/metadata) bypass the gate,
and write stats are recorded only after credit is acquired, next to the
WriteFrame they describe.

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.

Close and SendError now close the send window before taking the write
lock, mirroring SendCancel's signal-first ordering. Without this, both
would block on the write lock behind a send parked on credit and never
reach the terminate that frees it. SendError pre-closes with io.EOF to
match the send signal it sets moments later; Close pre-closes with
termClosed for the same reason.

CloseSend is a graceful half-close, so it waits for a parked send to
complete rather than aborting it -- but it must not hold s.mu while
waiting for the write lock: the parked writer is only freed by
terminate, and every terminate path (Cancel/Close/SendError/an inbound
terminal frame) needs s.mu first. CloseSend now releases s.mu before
waiting, then re-locks and re-checks state before sending
KindCloseSend.

Per-stream only; the connection-level window and the enablement path
that installs the window come in later commits.

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

Author: suj-krishnan

drpcstream/send_window_gate_test.go

@@ -0,0 +1,269 @@
+// 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: terminate (which closes the
+// window) runs before SendCancel takes 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))
+}
+
+// Close preempts a send parked on credit the same way: it closes the send
+// window before taking the write lock rather than waiting on a grant that
+// may never come.
+func TestStream_SendWindowClosePreemptsParkedWrite(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 close")
+	case <-time.After(blockShort):
+	}
+
+	assert.NoError(t, st.Close())
+
+	select {
+	case err := <-done:
+		// Same error later sends see: terminate sets sigs.send to termClosed.
+		assert.That(t, errors.Is(err, termClosed))
+	case <-time.After(time.Second):
+		t.Fatal("parked data write was not preempted by Close")
+	}
+}
+
+// SendError preempts a send parked on credit, like Close: it closes the send
+// window before taking the write lock so reporting an error is never stuck
+// behind a slow consumer.
+func TestStream_SendWindowSendErrorPreemptsParkedWrite(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 error")
+	case <-time.After(blockShort):
+	}
+
+	assert.NoError(t, st.SendError(errs.New("boom")))
+
+	select {
+	case err := <-done:
+		// io.EOF, matching sigs.send: parked and later sends agree.
+		assert.That(t, errors.Is(err, io.EOF))
+	case <-time.After(time.Second):
+		t.Fatal("parked data write was not preempted by SendError")
+	}
+}
+
+// CloseSend is a graceful half-close, not a termination: it must NOT preempt
+// a send parked on credit. It waits for the write lock; once credit arrives
+// the parked write completes successfully and CloseSend follows it out.
+func TestStream_SendWindowCloseSendWaitsForParkedWrite(t *testing.T) {
+	st := newGateStream(t)
+	st.sendw = newSendWindow(0) // send will park immediately
+
+	write := make(chan error, 1)
+	go func() { write <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }()
+
+	select {
+	case <-write:
+		t.Fatal("data write returned before credit")
+	case <-time.After(blockShort):
+	}
+
+	closeSend := make(chan error, 1)
+	go func() { closeSend <- st.CloseSend() }()
+
+	// Neither may make progress yet: the write is parked on credit and
+	// CloseSend is parked behind it on the write lock.
+	select {
+	case <-write:
+		t.Fatal("data write returned without credit")
+	case <-closeSend:
+		t.Fatal("CloseSend preempted a parked data write")
+	case <-time.After(blockShort):
+	}
+
+	st.sendw.grant(uint64(len("data")))
+
+	select {
+	case err := <-write:
+		assert.NoError(t, err) // the parked write completed, not aborted
+	case <-time.After(time.Second):
+		t.Fatal("parked data write did not complete after grant")
+	}
+	select {
+	case err := <-closeSend:
+		assert.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("CloseSend did not complete after the parked write finished")
+	}
+}
+
+// While CloseSend waits behind a credit-parked send, termination must still
+// be able to proceed: Cancel needs s.mu to terminate, so CloseSend must not
+// hold s.mu while waiting for the write lock.
+func TestStream_SendWindowCancelUnwedgesCloseSendBehindParkedWrite(t *testing.T) {
+	st := newGateStream(t)
+	st.sendw = newSendWindow(0) // send will park immediately
+
+	write := make(chan error, 1)
+	go func() { write <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }()
+
+	select {
+	case <-write:
+		t.Fatal("data write returned before credit")
+	case <-time.After(blockShort):
+	}
+
+	closeSend := make(chan error, 1)
+	go func() { closeSend <- st.CloseSend() }()
+
+	select {
+	case <-closeSend:
+		t.Fatal("CloseSend preempted a parked data write")
+	case <-time.After(blockShort):
+	}
+
+	// Cancel terminates the stream, closing the send window: the parked write
+	// wakes with an error and CloseSend unblocks as a no-op.
+	canceled := make(chan struct{})
+	go func() { st.Cancel(errs.New("boom")); close(canceled) }()
+
+	select {
+	case <-canceled:
+	case <-time.After(time.Second):
+		t.Fatal("Cancel blocked behind CloseSend waiting for the write lock")
+	}
+	select {
+	case err := <-write:
+		assert.That(t, errors.Is(err, io.EOF))
+	case <-time.After(time.Second):
+		t.Fatal("parked data write did not wake on termination")
+	}
+	select {
+	case err := <-closeSend:
+		assert.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("CloseSend did not unblock after termination")
+	}
+}
+
+// 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")
+	}
+}

drpcstream/stream.go

@@ -55,6 +55,10 @@ type Stream struct {
 	recvQueue ringBuffer
 	wbuf      []byte
 
+	// 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
@@ -352,6 +356,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()
 }
 
@@ -402,6 +412,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)
 
@@ -502,6 +521,13 @@ func (s *Stream) SendError(serr error) (err error) {
 	}
 
 	defer s.checkFinished()
+
+	// Close the send window before taking the write lock so a send parked on
+	// credit wakes and releases the lock (same ordering as Close). io.EOF to
+	// match the sigs.send error set below, so parked and later sends agree.
+	if s.sendw != nil {
+		s.sendw.close(io.EOF)
+	}
 	s.write.Lock()
 	defer s.write.Unlock()
 
@@ -558,6 +584,13 @@ func (s *Stream) Close() (err error) {
 	}
 
 	defer s.checkFinished()
+
+	// Close the send window before taking the write lock so a send parked on
+	// credit wakes and releases the lock, instead of stalling the close on a
+	// grant that may never come (mirrors SendCancel's signal-first ordering).
+	if s.sendw != nil {
+		s.sendw.close(termClosed)
+	}
 	s.write.Lock()
 	defer s.write.Unlock()
 
@@ -581,11 +614,22 @@ func (s *Stream) CloseSend() (err error) {
 		s.mu.Unlock()
 		return nil
 	}
+	s.mu.Unlock()
 
 	defer s.checkFinished()
+
+	// Wait for the write lock without holding s.mu: the writer holding it may
+	// be parked on send credit, and every path that frees it (terminate, via
+	// Cancel/Close/SendError or an inbound terminal frame) needs s.mu first.
 	s.write.Lock()
 	defer s.write.Unlock()
 
+	s.mu.Lock()
+	// Re-check: sending may have ended while we waited for the write lock.
+	if s.sigs.send.IsSet() || s.sigs.term.IsSet() {
+		s.mu.Unlock()
+		return nil
+	}
 	s.sigs.send.Set(sendClosed)
 	s.terminateIfBothClosed()
 	s.mu.Unlock()