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,115 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"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.
+//
+// Grants are no-revoke: grant is strictly additive, and the receiver never
+// takes back credit it has already issued. 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.
+type sendWindow struct {
+	mu     sync.Mutex
+	avail  int64         // available credit; signed (may be negative — see doc)
+	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 blocks until n bytes of credit are available, then debits them and
+// returns nil. It returns early if the window is closed (with the close error)
+// or ctx is canceled (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.
+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 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 adds n bytes of credit and wakes any parked acquirer. n is unsigned
+// (the wire delta type), so a grant can only ever raise the balance — no-revoke.
+// The addition saturates at math.MaxInt64 so a large or malicious delta can
+// never wrap the balance negative. Grants after close are ignored (the window
+// is dead).
+func (w *sendWindow) grant(n uint64) {
+	w.mu.Lock()
+	if !w.closed {
+		if n >= uint64(math.MaxInt64) {
+			w.avail = math.MaxInt64
+		} else if sum := w.avail + int64(n); sum < w.avail {
+			w.avail = math.MaxInt64 // overflowed past MaxInt64
+		} else {
+			w.avail = sum
+		}
+		w.wakeLocked()
+	}
+	w.mu.Unlock()
+}
+
+// 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) {
+	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,149 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcstream
+
+import (
+	"context"
+	"errors"
+	"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))
+}
+
+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))
+}