diff --git a/drpcstream/send_window.go b/drpcstream/send_window.go index 660a109..bc65dc8 100644 --- a/drpcstream/send_window.go +++ b/drpcstream/send_window.go @@ -4,7 +4,6 @@ package drpcstream import ( - "context" "io" "math" "sync" @@ -35,23 +34,19 @@ func (w *sendWindow) available() int64 { } // acquire debits n bytes of credit, blocking until available, and returns nil. -// It returns early (consuming no credit) if the window is closed or ctx is -// canceled -func (w *sendWindow) acquire(ctx context.Context, n int64) error { +// 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. +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 ctx.Err() != nil: - w.mu.Unlock() - return ctx.Err() case n <= 0: - // Nothing to acquire; after the terminal cases (so a closed/canceled - // window still fails) and never debits (so a negative n cannot add credit). - w.mu.Unlock() + w.mu.Unlock() // nothing to acquire; never debits, so a negative n adds no credit return nil case w.avail >= n: w.avail -= n @@ -60,19 +55,15 @@ func (w *sendWindow) acquire(ctx context.Context, n int64) error { } // 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. + // 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() - - select { - case <-ch: - // Credit was granted or the window closed; loop and re-check. - case <-ctx.Done(): - return ctx.Err() - } + <-ch // credit was granted or the window closed; loop and re-check } } diff --git a/drpcstream/send_window_gate_test.go b/drpcstream/send_window_gate_test.go new file mode 100644 index 0000000..3b27592 --- /dev/null +++ b/drpcstream/send_window_gate_test.go @@ -0,0 +1,121 @@ +// Copyright (C) 2026 Cockroach Labs. +// See LICENSE for copying information. + +package drpcstream + +import ( + "context" + "errors" + "io" + "testing" + "time" + + "github.com/zeebo/assert" + "github.com/zeebo/errs" + + "storj.io/drpc/drpcwire" +) + +// newGateStream builds a stream writing to io.Discard with an explicit +// SplitSize so small payloads are a single frame. +func newGateStream(t *testing.T) *Stream { + mw := testMuxWriter(t) + return NewWithOptions(context.Background(), 1, mw, NewBufferPool(), Options{SplitSize: 64 << 10}) +} + +// By default no send window is installed, so data writes are ungated +// (unlimited) and behavior is unchanged. +func TestStream_SendWindowDefaultUngated(t *testing.T) { + st := newGateStream(t) + assert.That(t, st.sendw == nil) + assert.NoError(t, st.RawWrite(drpcwire.KindMessage, []byte("hello"))) +} + +// With a finite send window, a data write blocks until enough credit is +// granted. +func TestStream_SendWindowGatesDataWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(4) // 4 bytes of credit + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("hello")) }() // 5 bytes > 4 + + select { + case <-done: + t.Fatal("data write returned before sufficient credit") + case <-time.After(blockShort): + } + + st.sendw.grant(1) // 4 + 1 = 5 >= 5 + + select { + case err := <-done: + assert.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("data write did not complete after grant") + } +} + +// Control kinds (here, invoke) are not flow-controlled: they proceed even with +// zero send credit. +func TestStream_SendWindowControlKindsBypassGate(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // no credit at all + + assert.NoError(t, st.WriteInvoke("service.Method", nil)) +} + +// SendCancel preempts a send parked on credit: it terminates (closing the +// window) before taking the write lock, so the parked write wakes, releases the +// lock, and the cancel frame goes out. +func TestStream_SendWindowSendCancelPreemptsParkedWrite(t *testing.T) { + st := newGateStream(t) + st.sendw = newSendWindow(0) // send will park immediately + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }() + + select { + case <-done: + t.Fatal("data write returned before cancel") + case <-time.After(blockShort): + } + + assert.NoError(t, st.SendCancel(context.Canceled)) + + select { + case err := <-done: + // Same error as a send parked in WriteFrame or a later send would see. + assert.That(t, errors.Is(err, io.EOF)) + case <-time.After(time.Second): + t.Fatal("parked data write was not preempted by SendCancel") + } + + // A subsequent send observes the same error as the parked one. + assert.That(t, errors.Is(st.RawWrite(drpcwire.KindMessage, []byte("more")), io.EOF)) +} + +// 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 + + done := make(chan error, 1) + go func() { done <- st.RawWrite(drpcwire.KindMessage, []byte("data")) }() + + select { + case <-done: + t.Fatal("data write returned before termination") + case <-time.After(blockShort): + } + + st.Cancel(errs.New("boom")) + + select { + case err := <-done: + // Cancel pre-sets sigs.send to io.EOF; the window closes with it. + assert.That(t, errors.Is(err, io.EOF)) + case <-time.After(time.Second): + t.Fatal("parked data write did not wake on termination") + } +} diff --git a/drpcstream/send_window_test.go b/drpcstream/send_window_test.go index 84c12a3..f1140df 100644 --- a/drpcstream/send_window_test.go +++ b/drpcstream/send_window_test.go @@ -4,7 +4,6 @@ package drpcstream import ( - "context" "errors" "io" "math" @@ -23,10 +22,10 @@ func TestSendWindowAcquireImmediate(t *testing.T) { w := newSendWindow(1000) assert.Equal(t, w.available(), int64(1000)) - assert.NoError(t, w.acquire(context.Background(), 400)) + assert.NoError(t, w.acquire(400)) assert.Equal(t, w.available(), int64(600)) - assert.NoError(t, w.acquire(context.Background(), 600)) + assert.NoError(t, w.acquire(600)) assert.Equal(t, w.available(), int64(0)) } @@ -36,14 +35,14 @@ func TestSendWindowGrantsAccumulate(t *testing.T) { w.grant(50) assert.Equal(t, w.available(), int64(150)) - assert.NoError(t, w.acquire(context.Background(), 150)) + assert.NoError(t, w.acquire(150)) assert.Equal(t, w.available(), int64(0)) } func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) { w := newSendWindow(100) done := make(chan error, 1) - go func() { done <- w.acquire(context.Background(), 300) }() + go func() { done <- w.acquire(300) }() // Not enough credit yet: acquire must block. select { @@ -63,30 +62,6 @@ func TestSendWindowAcquireBlocksUntilGrant(t *testing.T) { assert.Equal(t, w.available(), int64(50)) } -func TestSendWindowAcquireContextCancel(t *testing.T) { - w := newSendWindow(0) - ctx, cancel := context.WithCancel(context.Background()) - done := make(chan error, 1) - go func() { done <- w.acquire(ctx, 100) }() - - select { - case <-done: - t.Fatal("acquire returned before cancellation") - case <-time.After(blockShort): - } - - cancel() - - select { - case err := <-done: - assert.That(t, errors.Is(err, context.Canceled)) - case <-time.After(time.Second): - t.Fatal("acquire did not wake on context cancellation") - } - // Credit was not consumed by a failed acquire. - assert.Equal(t, w.available(), int64(0)) -} - func TestSendWindowGrantSaturates(t *testing.T) { // Adding to a near-max balance saturates at MaxInt64 instead of wrapping. w := newSendWindow(math.MaxInt64 - 10) @@ -121,22 +96,11 @@ func TestSendWindowGrantSaturates(t *testing.T) { assert.Equal(t, w5.available(), int64(math.MaxInt64)) } -func TestSendWindowAcquireCanceledCtxWithCredit(t *testing.T) { - w := newSendWindow(1000) - ctx, cancel := context.WithCancel(context.Background()) - cancel() // already canceled, with credit available - - err := w.acquire(ctx, 100) - assert.That(t, errors.Is(err, context.Canceled)) - // A canceled context must not consume credit even though it was available. - assert.Equal(t, w.available(), int64(1000)) -} - func TestSendWindowCloseWakesAcquire(t *testing.T) { w := newSendWindow(0) closeErr := errs.New("terminated") done := make(chan error, 1) - go func() { done <- w.acquire(context.Background(), 100) }() + go func() { done <- w.acquire(100) }() select { case <-done: @@ -160,32 +124,28 @@ func TestSendWindowAcquireAfterClose(t *testing.T) { w.close(closeErr) // Even though credit is available, a closed window returns the close error. - assert.That(t, errors.Is(w.acquire(context.Background(), 1), closeErr)) + 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(context.Background(), 1), io.EOF)) + assert.That(t, errors.Is(w.acquire(1), io.EOF)) } func TestSendWindowAcquireNonPositive(t *testing.T) { w := newSendWindow(100) - assert.NoError(t, w.acquire(context.Background(), 0)) - assert.NoError(t, w.acquire(context.Background(), -5)) + 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 TestSendWindowAcquireZeroObservesTerminalState(t *testing.T) { - // A zero-length acquire (empty frame) must still observe a canceled context... - ctx, cancel := context.WithCancel(context.Background()) - cancel() - assert.That(t, errors.Is(newSendWindow(100).acquire(ctx, 0), context.Canceled)) - - // ...and a closed window, rather than reporting success. +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(context.Background(), 0), closeErr)) + assert.That(t, errors.Is(w.acquire(0), closeErr)) } diff --git a/drpcstream/stream.go b/drpcstream/stream.go index 60e1409..ce9c73a 100644 --- a/drpcstream/stream.go +++ b/drpcstream/stream.go @@ -61,6 +61,10 @@ type Stream struct { cbuf []byte // compression scratch buffer dbuf []byte // decompression scratch buffer + // sendw is the per-stream send-side flow-control window. It is nil when + // flow control is not enabled, in which case data writes are ungated. + sendw *sendWindow + mu sync.Mutex // protects state transitions sigs struct { send drpcsignal.Signal // set when done sending messages @@ -361,6 +365,12 @@ 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()) + } s.checkFinished() } @@ -417,6 +427,15 @@ func (s *Stream) rawWriteLocked(kind drpcwire.Kind, data []byte) (err error) { fr.Data, data = drpcwire.SplitData(data, n) fr.Done = len(data) == 0 + // Only data frames consume send credit; a nil window (flow control + // disabled) leaves sends ungated. acquire parks until credit arrives or + // the window closes (stream termination) -- the latter is the abort path. + if kind == drpcwire.KindMessage && s.sendw != nil { + if err := s.sendw.acquire(int64(len(fr.Data))); err != nil { + return err + } + } + drpcopts.GetStreamStats(&s.opts.Internal).AddWritten(uint64(len(fr.Data))) s.log("SEND", fr.String)