diff --git a/drpcstream/send_window.go b/drpcstream/send_window.go index bc65dc8..05f0315 100644 --- a/drpcstream/send_window.go +++ b/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 - } -} diff --git a/drpcstream/send_window_gate_test.go b/drpcstream/send_window_gate_test.go index 3b27592..6e41d3b 100644 --- a/drpcstream/send_window_gate_test.go +++ b/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")) }() diff --git a/drpcstream/send_window_test.go b/drpcstream/send_window_test.go index f1140df..bac2001 100644 --- a/drpcstream/send_window_test.go +++ b/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)) -} diff --git a/drpcstream/stream.go b/drpcstream/stream.go index ce9c73a..44c3975 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -365,12 +365,10 @@ func (s *Stream) terminate(err error) { s.sigs.recv.Set(err) s.sigs.term.Set(err) s.recvQueue.Close(err) - if s.sendw != nil { - // Close with the send-side error: sigs.send is first-wins, so when a - // caller pre-set it (io.EOF for cancel/error), a send parked on credit - // returns the same error as one parked in WriteFrame or a later send. - s.sendw.close(s.sigs.send.Err()) - } + // The send window's done channel is sigs.send, so setting the signal above + // already wakes any acquire parked on credit; it returns sigs.send.Err() -- + // the same error a send parked in WriteFrame or a later send sees. No + // separate close is needed. s.checkFinished() } diff --git a/drpcwire/mux_writer.go b/drpcwire/mux_writer.go index 9fa21df..f318026 100644 --- a/drpcwire/mux_writer.go +++ b/drpcwire/mux_writer.go @@ -168,6 +168,16 @@ func (mw *MuxWriter) WriteFrame(fr Frame, cancel *drpcsignal.Signal) (err error) mw.mu.Unlock() return mw.closeErr } + // Honor the cancel signal on the non-blocking path too, not only when + // parking: otherwise a data frame whose send was canceled/terminated + // (sigs.send set) between the caller's check and here would still be + // appended whenever the buffer has room. IsSet is an atomic read, so the + // fast path stays free of the signal's lazy channel allocation. Terminal + // frames pass a nil cancel, so they are never dropped here. + if cancel != nil && cancel.IsSet() { + mw.mu.Unlock() + return cancel.Err() + } // A control-bit frame (e.g. an abortive KindCancel) is appended even past // the cap so it is never delayed by backpressure. This overshoots the // high-water mark by at most one small control frame per terminating diff --git a/drpcwire/mux_writer_test.go b/drpcwire/mux_writer_test.go index 2bee335..80723ae 100644 --- a/drpcwire/mux_writer_test.go +++ b/drpcwire/mux_writer_test.go @@ -434,6 +434,20 @@ func TestMuxWriter_WriteFrameCanceledWhileBlocked(t *testing.T) { <-mw.Done() } +// An already-set cancel signal is honored on the non-blocking path: a data +// frame is not appended even when the buffer has room. +func TestMuxWriter_WriteFrameCanceledWithRoom(t *testing.T) { + mw := NewMuxWriter(io.Discard, func(error) {}) + defer func() { mw.Stop(nil); <-mw.Done() }() + + var cancel drpcsignal.Signal + cancelErr := errors.New("canceled") + cancel.Set(cancelErr) + + // Buffer is empty (room available), but the cancel is already set. + assert.Equal(t, mw.WriteFrame(RandFrame(), &cancel), cancelErr) +} + // A control-bit frame is appended immediately even when the buffer is full, so // an abortive cancel is never delayed by backpressure. func TestMuxWriter_ControlFrameBypassesFullBuffer(t *testing.T) {