drpcstream: optimize sendWindow

drpcstream/send_window.go

@@ -7,75 +7,80 @@ 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.
+// credit, and close terminates a blocked acquire. There must be at most one
+// goroutine calling acquire and one goroutine calling grant.
 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 // available credit; negative while acquire is waiting
+
+	mu   sync.Mutex
+	cond *sync.Cond // allocated only if acquire has to wait
+	err  error      // terminal error returned by a blocked acquire after close
 }
 
 // newSendWindow returns a sendWindow seeded with initial credit.
 func newSendWindow(initial int64) *sendWindow {
-	return &sendWindow{avail: initial}
+	w := &sendWindow{}
+	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
+	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 and blocks while the resulting balance is
+// negative. If the window is closed while blocked, it returns the close error.
+// n <= 0 is a no-op.
 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 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
-		}
-		// 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{})
+	if n <= 0 {
+		return nil
+	}
+	if w.avail.Add(-n) >= 0 {
+		return nil
+	}
+
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	for w.avail.Load() < 0 && w.err == nil {
+		if w.cond == nil {
+			w.cond = sync.NewCond(&w.mu)
 		}
-		ch := w.notify
-		w.mu.Unlock()
-		<-ch // credit was granted or the window closed; loop and re-check
+		w.cond.Wait()
 	}
+
+	return w.err // can be nil
 }
 
-// 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 wakes the acquirer if the grant satisfies
+// its deficit. n is 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 {
+		oldCredits := w.avail.Load()
+		newCredits := applyGrant(oldCredits, n)
+		if !w.avail.CompareAndSwap(oldCredits, newCredits) {
+			continue
+		}
+		if oldCredits < 0 && newCredits >= 0 {
+			// Taking mu makes the predicate check in acquire and this signal
+			// atomic with respect to cond.Wait, preventing a lost wakeup.
+			w.mu.Lock()
+			if w.cond != nil {
+				w.cond.Signal()
+			}
+			w.mu.Unlock()
+		}
+		return
 	}
-	w.mu.Unlock()
 }
 
 // applyGrant returns avail + n with an upper bound of math.MaxInt64.
@@ -100,26 +105,17 @@ func applyGrant(avail int64, n uint64) int64 {
 }
 
 // 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.
+// then returns 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
+		err = io.EOF
 	}
 	w.mu.Lock()
-	if !w.closed {
-		w.closed = true
+	if w.err == nil {
 		w.err = err
-		w.wakeLocked()
+		if w.cond != nil {
+			w.cond.Signal()
+		}
 	}
 	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_test.go

@@ -50,8 +50,9 @@ func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) {
 		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
 
 	select {
 	case err := <-done:
@@ -77,21 +78,21 @@ func TestSendWindowGrantSaturates(t *testing.T) {
 
 	// A grant that raises a negative balance is applied exactly (no saturation).
 	w3 := newSendWindow(0)
-	w3.avail = -300
+	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.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.avail.Store(-300)
 	w5.grant(math.MaxUint64)
 	assert.Equal(t, w5.available(), int64(math.MaxInt64))
 }
@@ -120,16 +121,16 @@ func TestSendWindowCloseWakesAcquire(t *testing.T) {
 
 func TestSendWindowAcquireAfterClose(t *testing.T) {
 	w := newSendWindow(1000)
-	closeErr := errs.New("closed")
-	w.close(closeErr)
+	w.close(errs.New("closed"))
 
-	// Even though credit is available, a closed window returns the close error.
-	assert.That(t, errors.Is(w.acquire(1), closeErr))
+	// The fast path does not consult closed when enough credit is available.
+	assert.NoError(t, w.acquire(1))
+	assert.Equal(t, w.available(), int64(999))
 }
 
 func TestSendWindowCloseNilError(t *testing.T) {
-	w := newSendWindow(1000)
-	w.close(nil) // closing with nil must not let a later acquire report success
+	w := newSendWindow(0)
+	w.close(nil)
 	assert.That(t, errors.Is(w.acquire(1), io.EOF))
 }
 
@@ -141,11 +142,8 @@ func TestSendWindowAcquireNonPositive(t *testing.T) {
 	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).
+func TestSendWindowAcquireZeroAfterClose(t *testing.T) {
 	w := newSendWindow(100)
-	closeErr := errs.New("terminated")
-	w.close(closeErr)
-	assert.That(t, errors.Is(w.acquire(0), closeErr))
+	w.close(errs.New("terminated"))
+	assert.NoError(t, w.acquire(0))
 }