drpcstream: add per-stream send-window credit primitive

Add sendWindow, the per-stream flow-control credit balance on the sender.
acquire spends credit, blocking until enough is available; grant adds
credit; close terminates the window. acquire returns early on close (with
the close error) or context cancellation (with ctx.Err()), consuming no
credit in either case; both are checked before any debit, so a canceled
context never consumes credit or lets a frame proceed even when credit is
available.

Grants are no-revoke: grant takes an unsigned delta (the wire type), so a
grant can only ever raise the balance, and the addition saturates at
math.MaxInt64 so a large or malicious delta can never wrap it negative.
The balance is a signed int64; acquire never drives it below zero, but the
enablement layer may debit credit for bytes sent before flow control
begins enforcing, leaving it negative, in which case acquire blocks until
grants restore it.

This is the primitive only -- it is not yet wired into the stream write
path (that gate comes in a later commit).

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

Author: suj-krishnan

drpcstream/send_window.go

@@ -0,0 +1,133 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"io"
+	"math"
+	"sync"
+)
+
+// sendWindow is a per-stream flow-control credit balance on the sender. It
+// tracks how many more bytes the stream is allowed to put on the wire right
+// now. acquire spends credit (blocking until enough is available), grant adds
+// credit, and close terminates the window.
+//
+// avail is a signed int64;
+// acquire never drives it below zero, but the enablement layer may debit credit
+// for bytes sent before flow control begins enforcing, leaving it negative, in
+// which case acquire blocks until grants restore it.
+type sendWindow struct {
+	mu     sync.Mutex
+	avail  int64         // available credit; signed (maybe negative during enablement)
+	closed bool          // set once by close; no further acquires succeed
+	err    error         // terminal error returned by acquire after close
+	notify chan struct{} // closed+replaced to wake parked acquirers
+}
+
+// newSendWindow returns a sendWindow seeded with initial credit.
+func newSendWindow(initial int64) *sendWindow {
+	return &sendWindow{avail: initial, notify: make(chan struct{})}
+}
+
+// available returns the current credit balance.
+func (w *sendWindow) available() int64 {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.avail
+}
+
+// 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 {
+	for {
+		w.mu.Lock()
+		switch {
+		case w.closed:
+			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()
+			return nil
+		case w.avail >= n:
+			w.avail -= n
+			w.mu.Unlock()
+			return nil
+		}
+		// Snapshot the notify channel under the lock before parking, so a grant
+		// or close that fires the instant we unlock is not missed.
+		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()
+		}
+	}
+}
+
+// grant raises the balance by n and wakes any parked acquirer. n is unsigned,
+// so a grant never lowers the balance. Grants after close are ignored.
+func (w *sendWindow) grant(n uint64) {
+	w.mu.Lock()
+	if !w.closed {
+		w.avail = applyGrant(w.avail, n)
+		w.wakeLocked()
+	}
+	w.mu.Unlock()
+}
+
+// applyGrant returns avail + n with an upper bound of math.MaxInt64.
+// n is the wire delta (unsigned), so the result never drops below avail.
+func applyGrant(avail int64, n uint64) int64 {
+	if avail >= 0 {
+		if n > uint64(math.MaxInt64-avail) {
+			return math.MaxInt64
+		}
+		return avail + int64(n)
+	}
+	// A negative avail (before FC enablement) is repaid before any positive
+	// credit starts accruing.
+	deficit := uint64(-avail) // |avail| as uint64; -math.MinInt64 wraps to its magnitude
+	if n <= deficit {
+		return -int64(deficit - n) // debt only partly repaid; result still <= 0
+	}
+	if rem := n - deficit; rem <= uint64(math.MaxInt64) {
+		return int64(rem)
+	}
+	return math.MaxInt64
+}
+
+// close terminates the window with err, waking every parked acquirer, which
+// then returns err. Subsequent acquires also return err. It is a no-op if the
+// window is already closed.
+func (w *sendWindow) close(err error) {
+	if err == nil {
+		err = io.EOF // never nil: acquire must not report success on a closed window
+	}
+	w.mu.Lock()
+	if !w.closed {
+		w.closed = true
+		w.err = err
+		w.wakeLocked()
+	}
+	w.mu.Unlock()
+}
+
+// wakeLocked broadcasts to all parked acquirers by closing the current notify
+// channel and installing a fresh one. It must be called with w.mu held.
+func (w *sendWindow) wakeLocked() {
+	close(w.notify)
+	w.notify = make(chan struct{})
+}

drpcstream/send_window_test.go

@@ -0,0 +1,191 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"errors"
+	"io"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/zeebo/assert"
+	"github.com/zeebo/errs"
+)
+
+// blockShort is how long we wait to conclude that an acquire is (correctly)
+// blocked before we release it.
+const blockShort = 20 * time.Millisecond
+
+func TestSendWindowAcquireImmediate(t *testing.T) {
+	w := newSendWindow(1000)
+	assert.Equal(t, w.available(), int64(1000))
+
+	assert.NoError(t, w.acquire(context.Background(), 400))
+	assert.Equal(t, w.available(), int64(600))
+
+	assert.NoError(t, w.acquire(context.Background(), 600))
+	assert.Equal(t, w.available(), int64(0))
+}
+
+func TestSendWindowGrantsAccumulate(t *testing.T) {
+	w := newSendWindow(0)
+	w.grant(100)
+	w.grant(50)
+	assert.Equal(t, w.available(), int64(150))
+
+	assert.NoError(t, w.acquire(context.Background(), 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) }()
+
+	// Not enough credit yet: acquire must block.
+	select {
+	case <-done:
+		t.Fatal("acquire returned before sufficient credit")
+	case <-time.After(blockShort):
+	}
+
+	w.grant(250) // 100 + 250 = 350 >= 300
+
+	select {
+	case err := <-done:
+		assert.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("acquire did not return after grant")
+	}
+	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)
+	w.grant(100)
+	assert.Equal(t, w.available(), int64(math.MaxInt64))
+
+	// A single delta larger than MaxInt64 (the wire delta is uint64) saturates
+	// rather than becoming a negative balance.
+	w2 := newSendWindow(0)
+	w2.grant(math.MaxUint64)
+	assert.Equal(t, w2.available(), int64(math.MaxInt64))
+	assert.That(t, w2.available() > 0) // never wrapped negative
+
+	// A grant that raises a negative balance is applied exactly (no saturation).
+	w3 := newSendWindow(0)
+	w3.avail = -300
+	w3.grant(500)
+	assert.Equal(t, w3.available(), int64(200))
+
+	// A very large grant against a negative balance repays the debt before
+	// clamping: the true sum -300 + MaxInt64 fits, so it must not saturate to
+	// MaxInt64 (which would erase the 300 of pre-enforcement debt).
+	w4 := newSendWindow(0)
+	w4.avail = -300
+	w4.grant(math.MaxInt64)
+	assert.Equal(t, w4.available(), int64(math.MaxInt64-300))
+
+	// Only when the true sum truly overflows does it clamp.
+	w5 := newSendWindow(0)
+	w5.avail = -300
+	w5.grant(math.MaxUint64)
+	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) }()
+
+	select {
+	case <-done:
+		t.Fatal("acquire returned before close")
+	case <-time.After(blockShort):
+	}
+
+	w.close(closeErr)
+
+	select {
+	case err := <-done:
+		assert.That(t, errors.Is(err, closeErr))
+	case <-time.After(time.Second):
+		t.Fatal("acquire did not wake on close")
+	}
+}
+
+func TestSendWindowAcquireAfterClose(t *testing.T) {
+	w := newSendWindow(1000)
+	closeErr := errs.New("closed")
+	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))
+}
+
+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))
+}
+
+func TestSendWindowAcquireNonPositive(t *testing.T) {
+	w := newSendWindow(100)
+	assert.NoError(t, w.acquire(context.Background(), 0))
+	assert.NoError(t, w.acquire(context.Background(), -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.
+	w := newSendWindow(100)
+	closeErr := errs.New("terminated")
+	w.close(closeErr)
+	assert.That(t, errors.Is(w.acquire(context.Background(), 0), closeErr))
+}