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

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
 	}
 }
 

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/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))
 }

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)