diff --git a/drpcstream/send_window.go b/drpcstream/send_window.go index bc65dc8..d166ef8 100644 --- a/drpcstream/send_window.go +++ b/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 - } -} diff --git a/drpcstream/send_window_test.go b/drpcstream/send_window_test.go index f1140df..4c385ae 100644 --- a/drpcstream/send_window_test.go +++ b/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,7 +78,7 @@ 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)) @@ -85,13 +86,13 @@ func TestSendWindowGrantSaturates(t *testing.T) { // 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)) }