From 995bd817a0a2793e74c2d8449f120bbdec42e2fd Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Thu, 23 Jul 2026 07:03:46 +0530 Subject: [PATCH 1/2] drpcstream: make sendWindow mutex-free with a channel and done signal Rework sendWindow into a lock-free credit balance in the style of grpc-go's writeQuota: an atomic balance, a cap-1 buffered channel for grant wakeups, and a caller-supplied done channel for termination. There is at most one acquirer (serialized by the stream write lock) and one granter (the connection reader), so no mutex is needed. acquire debits credit atomically and, if that leaves a deficit, waits on the channel; grant repays the balance and edge-signals the channel only when the deficit clears. A cap-1 non-blocking send makes the wakeup lost-wakeup-free without a condition variable. Termination is delegated to the done channel -- the stream's send signal -- so the window holds no error or closed state. terminate no longer calls into the window; setting sigs.send both wakes a blocked acquire and supplies its error (sigs.send.Err()), which matches the error a send parked in WriteFrame or a later send reads. acquire is a pure credit gate: with credit available it returns without consulting done, so a send racing termination can be admitted. That is harmless -- rawWriteLocked holds the stream write lock, so the frame is ordered before the terminal frame, and the write path drops it when the send signal is already set (companion change). Only a blocked acquire observes done, which is what wakes a parked sender on termination. Co-Authored-By: roachdev-claude --- drpcstream/send_window.go | 134 ++++++++++++---------------- drpcstream/send_window_gate_test.go | 8 +- drpcstream/send_window_test.go | 102 ++++++++++++--------- drpcstream/stream.go | 10 +-- 4 files changed, 122 insertions(+), 132 deletions(-) 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() } From fdfd41e2c4cebf29bddf9a038edefdf94b3f215a Mon Sep 17 00:00:00 2001 From: Sujatha Sivaramakrishnan Date: Thu, 23 Jul 2026 13:39:32 +0530 Subject: [PATCH 2/2] drpcwire: honor the cancel signal on WriteFrame's non-blocking path WriteFrame consulted its cancel signal only when it had to park on a full buffer; on the common non-blocking path (buffer has room) it appended unconditionally. A data frame whose send was canceled or the stream terminated -- sigs.send set after the write path's own check but before WriteFrame -- would therefore still reach the wire. Check cancel at the top of the loop, before appending. IsSet is an atomic read and does not force the signal's lazy channel allocation, so the fast path stays allocation-free. Terminal frames pass a nil cancel and are never affected. This lets the send window stay a pure credit gate: acquire need not be absorbing on its fast path, because the write path now drops a data frame whose stream has terminated -- the drpc analogue of grpc-go's loopyWriter checking live stream membership before writing. It also closes the pre-existing race (independent of flow control) where a data frame could be appended just after the send signal was set. Co-Authored-By: roachdev-claude --- drpcwire/mux_writer.go | 10 ++++++++++ drpcwire/mux_writer_test.go | 14 ++++++++++++++ 2 files changed, 24 insertions(+) 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) {