drpcstream: mutex-free sendWindow flow-control primitive

drpcstream/send_window.go

@@ -4,78 +4,81 @@
 package drpcstream
 
 import (
-	"io"
 	"math"
-	"sync"
+	"sync/atomic"
 )
 
 // 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.
+// tracks how many more bytes the stream may put on the wire. acquire spends
+// credit, blocking while the balance is negative (insufficient); grant adds
+// credit and wakes a blocked acquire once the deficit is repaid.
+//
+// Termination is delegated to the done channel supplied at
+// construction -- the stream's send signal.
 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{} // lazily allocated by parkers; closed+nilled to wake them
+	avail atomic.Int64    // credit; negative means an acquire is waiting for repayment
+	ch    chan struct{}   // grant signals here, acquire waits on it
+	done  <-chan struct{} // closed to abort a blocked (or future) acquire
+	err   func() error    // terminal error, consulted only when done is closed
 }
 
-// newSendWindow returns a sendWindow seeded with initial credit.
-func newSendWindow(initial int64) *sendWindow {
-	return &sendWindow{avail: initial}
+// newSendWindow returns a sendWindow seeded with initial credit. Closing done
+// aborts acquire, which then returns err().
+func newSendWindow(initial int64, done <-chan struct{}, err func() error) *sendWindow {
+	w := &sendWindow{ch: make(chan struct{}, 1), done: done, err: err}
+	w.avail.Store(initial)
+	return w
 }
 
 // available returns the current credit balance.
-func (w *sendWindow) available() int64 {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	return w.avail
-}
+func (w *sendWindow) available() int64 { return w.avail.Load() }
 
-// acquire debits n bytes of credit, blocking until available, and returns nil.
-// 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.
+// acquire debits n bytes of credit, blocking while the resulting balance is
+// negative, and returns nil once it is covered. A blocked acquire returns err()
+// when done is closed. n <= 0 is a no-op.
 func (w *sendWindow) acquire(n int64) error {
+	if n <= 0 || w.avail.Add(-n) >= 0 {
+		return nil
+	}
 	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 n <= 0:
-			w.mu.Unlock() // nothing to acquire; never debits, so a negative n adds no credit
-			return nil
-		case w.avail >= n:
-			w.avail -= n
-			w.mu.Unlock()
-			return nil
+		select {
+		case <-w.done:
+			return w.err()
+		case <-w.ch:
+			if w.avail.Load() >= 0 {
+				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. Allocated
-		// 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()
-		<-ch // credit was granted or the window closed; loop and re-check
 	}
 }
 
-// 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.
+// grant raises the balance by n and, when that repays a waiting acquire's
+// deficit, wakes it. n is the wire delta (unsigned), so a grant never lowers the
+// balance.
 func (w *sendWindow) grant(n uint64) {
-	w.mu.Lock()
-	if !w.closed {
-		w.avail = applyGrant(w.avail, n)
-		w.wakeLocked()
+	if n == 0 {
+		return
+	}
+	for {
+		old := w.avail.Load()
+		next := applyGrant(old, n)
+		if !w.avail.CompareAndSwap(old, next) {
+			continue
+		}
+		if old < 0 && next >= 0 {
+			// The deficit just cleared, so unblock a blocked acquire.
+			// A signal with no waiter yet stays buffered until it is consumed,
+			// so the wakeup is never lost. A later grant that also tries to
+			// deposit a token before it is consumed, just drops it. This is
+			// okay because when acquire finally wakes, it reads the current
+			// avail, which reflects every grant that happened so far.
+			select {
+			case w.ch <- struct{}{}:
+			default:
+			}
+		}
+		return
 	}
-	w.mu.Unlock()
 }
 
 // applyGrant returns avail + n with an upper bound of math.MaxInt64.
@@ -98,28 +101,3 @@ func applyGrant(avail int64, n uint64) int64 {
 	}
 	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 notify channel;
-// Requires w.mu.
-func (w *sendWindow) wakeLocked() {
-	if w.notify != nil {
-		close(w.notify)
-		w.notify = nil
-	}
-}

drpcstream/send_window_gate_test.go

@@ -35,7 +35,7 @@ func TestStream_SendWindowDefaultUngated(t *testing.T) {
 // granted.
 func TestStream_SendWindowGatesDataWrite(t *testing.T) {
 	st := newGateStream(t)
-	st.sendw = newSendWindow(4) // 4 bytes of credit
+	st.sendw = newSendWindow(4, st.sigs.send.Signal(), st.sigs.send.Err) // 4 bytes of credit
 
 	done := make(chan error, 1)
 	go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("hello")) }() // 5 bytes > 4
@@ -60,7 +60,7 @@ func TestStream_SendWindowGatesDataWrite(t *testing.T) {
 // zero send credit.
 func TestStream_SendWindowControlKindsBypassGate(t *testing.T) {
 	st := newGateStream(t)
-	st.sendw = newSendWindow(0) // no credit at all
+	st.sendw = newSendWindow(0, st.sigs.send.Signal(), st.sigs.send.Err) // no credit at all
 
 	assert.NoError(t, st.WriteInvoke("service.Method", nil))
 }
@@ -70,7 +70,7 @@ func TestStream_SendWindowControlKindsBypassGate(t *testing.T) {
 // lock, and the cancel frame goes out.
 func TestStream_SendWindowSendCancelPreemptsParkedWrite(t *testing.T) {
 	st := newGateStream(t)
-	st.sendw = newSendWindow(0) // send will park immediately
+	st.sendw = newSendWindow(0, st.sigs.send.Signal(), st.sigs.send.Err) // send will park immediately
 
 	done := make(chan error, 1)
 	go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }()
@@ -98,7 +98,7 @@ func TestStream_SendWindowSendCancelPreemptsParkedWrite(t *testing.T) {
 // 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
+	st.sendw = newSendWindow(0, st.sigs.send.Signal(), st.sigs.send.Err) // send will park immediately
 
 	done := make(chan error, 1)
 	go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }()

drpcstream/send_window_test.go

@@ -5,7 +5,6 @@ package drpcstream
 
 import (
 	"errors"
-	"io"
 	"math"
 	"testing"
 	"time"
@@ -18,8 +17,16 @@ import (
 // blocked before we release it.
 const blockShort = 20 * time.Millisecond
 
+// newTestWindow returns a window seeded with initial credit plus a terminate
+// func that closes its done channel; a terminated acquire returns termErr.
+func newTestWindow(initial int64, termErr error) (*sendWindow, func()) {
+	done := make(chan struct{})
+	w := newSendWindow(initial, done, func() error { return termErr })
+	return w, func() { close(done) }
+}
+
 func TestSendWindowAcquireImmediate(t *testing.T) {
-	w := newSendWindow(1000)
+	w, _ := newTestWindow(1000, nil)
 	assert.Equal(t, w.available(), int64(1000))
 
 	assert.NoError(t, w.acquire(400))
@@ -30,7 +37,7 @@ func TestSendWindowAcquireImmediate(t *testing.T) {
 }
 
 func TestSendWindowGrantsAccumulate(t *testing.T) {
-	w := newSendWindow(0)
+	w, _ := newTestWindow(0, nil)
 	w.grant(100)
 	w.grant(50)
 	assert.Equal(t, w.available(), int64(150))
@@ -40,18 +47,19 @@ func TestSendWindowGrantsAccumulate(t *testing.T) {
 }
 
 func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) {
-	w := newSendWindow(100)
+	w, _ := newTestWindow(100, nil)
 	done := make(chan error, 1)
 	go func() { done <- w.acquire(300) }()
 
-	// Not enough credit yet: acquire must block.
+	// Not enough credit yet: acquire debits into a deficit and blocks.
 	select {
 	case <-done:
 		t.Fatal("acquire returned before sufficient credit")
 	case <-time.After(blockShort):
 	}
+	assert.Equal(t, w.available(), int64(-200))
 
-	w.grant(250) // 100 + 250 = 350 >= 300
+	w.grant(250) // -200 + 250 = 50 >= 0
 
 	select {
 	case err := <-done:
@@ -62,90 +70,96 @@ func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) {
 	assert.Equal(t, w.available(), int64(50))
 }
 
+// A grant that only partially repays the deficit does not wake the acquire; the
+// grant that clears it does.
+func TestSendWindowPartialGrantsThenClear(t *testing.T) {
+	w, _ := newTestWindow(0, nil)
+	done := make(chan error, 1)
+	go func() { done <- w.acquire(300) }()
+
+	select {
+	case <-done:
+		t.Fatal("acquire returned before any credit")
+	case <-time.After(blockShort):
+	}
+
+	w.grant(250) // -300 + 250 = -50, still negative -> stays blocked
+	select {
+	case <-done:
+		t.Fatal("acquire returned while still in deficit")
+	case <-time.After(blockShort):
+	}
+
+	w.grant(100) // -50 + 100 = 50 >= 0 -> wakes
+	select {
+	case err := <-done:
+		assert.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("acquire did not return after the deficit cleared")
+	}
+	assert.Equal(t, w.available(), int64(50))
+}
+
 func TestSendWindowGrantSaturates(t *testing.T) {
 	// Adding to a near-max balance saturates at MaxInt64 instead of wrapping.
-	w := newSendWindow(math.MaxInt64 - 10)
+	w, _ := newTestWindow(math.MaxInt64-10, nil)
 	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, _ := newTestWindow(0, nil)
 	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, _ := newTestWindow(0, nil)
+	w3.avail.Store(-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, _ := newTestWindow(0, nil)
+	w4.avail.Store(-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, _ := newTestWindow(0, nil)
+	w5.avail.Store(-300)
 	w5.grant(math.MaxUint64)
 	assert.Equal(t, w5.available(), int64(math.MaxInt64))
 }
 
-func TestSendWindowCloseWakesAcquire(t *testing.T) {
-	w := newSendWindow(0)
+func TestSendWindowDoneWakesAcquire(t *testing.T) {
 	closeErr := errs.New("terminated")
+	w, terminate := newTestWindow(0, closeErr)
 	done := make(chan error, 1)
 	go func() { done <- w.acquire(100) }()
 
 	select {
 	case <-done:
-		t.Fatal("acquire returned before close")
+		t.Fatal("acquire returned before termination")
 	case <-time.After(blockShort):
 	}
 
-	w.close(closeErr)
+	terminate()
 
 	select {
 	case err := <-done:
 		assert.That(t, errors.Is(err, closeErr))
 	case <-time.After(time.Second):
-		t.Fatal("acquire did not wake on close")
+		t.Fatal("acquire did not wake on termination")
 	}
 }
 
-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(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(1), io.EOF))
-}
-
 func TestSendWindowAcquireNonPositive(t *testing.T) {
-	w := newSendWindow(100)
+	w, _ := newTestWindow(100, nil)
 	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 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(0), closeErr))
-}