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

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")
+	}
+}

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)