Skip to content

Commit 48c98b5

Browse files
drpcstream: add per-stream receive-window grant policy
Add recvWindow, the per-stream receive side of flow control. It tracks buffered bytes (in-progress reassembly plus completed, not-yet-consumed data) and decides when to return credit to the sender: grants are withheld while buffered is at or above the high-water mark, and the accrued returnable credit is otherwise released once it reaches the threshold, coalescing many frames into a single grant. Consuming reopens a gate that the high-water mark had closed. dispatched/consumed return the credit delta the caller should emit as a KindWindowUpdate. Emitting the frame, wiring into the read path (HandleFrame/MsgRecv), and the connection-level receive budget come in later commits. Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>
1 parent 10b863f commit 48c98b5

2 files changed

Lines changed: 124 additions & 0 deletions

File tree

drpcstream/recv_window.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
// Copyright (C) 2026 Cockroach Labs.
2+
// See LICENSE for copying information.
3+
4+
package drpcstream
5+
6+
import "sync"
7+
8+
// recvWindow decides when the receive side returns credit to the sender: it
9+
// withholds grants while buffered bytes are at or above the high-water mark
10+
// (credit keeps accruing), and otherwise releases accrued credit once it
11+
// reaches the threshold, coalescing many frames into one grant. dispatched
12+
// (reader goroutine) and consumed (application goroutine) return the credit
13+
// delta to send as a KindWindowUpdate and may run concurrently.
14+
type recvWindow struct {
15+
mu sync.Mutex
16+
buffered int64 // in-progress reassembly + completed, not-yet-consumed bytes
17+
pending int64 // returnable credit accrued but not yet granted
18+
highWater int64 // withhold grants while buffered is at or above this
19+
threshold int64 // release accrued credit once it reaches this
20+
}
21+
22+
// newRecvWindow returns a recvWindow seeded with highWater and threshold.
23+
func newRecvWindow(highWater, threshold int64) *recvWindow {
24+
return &recvWindow{highWater: highWater, threshold: threshold}
25+
}
26+
27+
// dispatched records that n bytes were dispatched off the wire into the
28+
// stream's buffers and returns the credit delta to grant (0 if none).
29+
func (w *recvWindow) dispatched(n int64) int64 {
30+
w.mu.Lock()
31+
defer w.mu.Unlock()
32+
w.buffered += n
33+
w.pending += n
34+
return w.maybeGrantLocked()
35+
}
36+
37+
// consumed records that n bytes were consumed by the application and returns
38+
// the credit delta to grant (0 if none).
39+
func (w *recvWindow) consumed(n int64) int64 {
40+
w.mu.Lock()
41+
defer w.mu.Unlock()
42+
w.buffered -= n
43+
return w.maybeGrantLocked()
44+
}
45+
46+
// bufferedBytes returns the current buffered byte count.
47+
func (w *recvWindow) bufferedBytes() int64 {
48+
w.mu.Lock()
49+
defer w.mu.Unlock()
50+
return w.buffered
51+
}
52+
53+
// maybeGrantLocked releases accrued credit when the gate is open (buffered
54+
// below high-water) and pending has reached the threshold. Requires w.mu.
55+
func (w *recvWindow) maybeGrantLocked() int64 {
56+
if w.buffered < w.highWater && w.pending >= w.threshold {
57+
g := w.pending
58+
w.pending = 0
59+
return g
60+
}
61+
return 0
62+
}

drpcstream/recv_window_test.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
// Copyright (C) 2026 Cockroach Labs.
2+
// See LICENSE for copying information.
3+
4+
package drpcstream
5+
6+
import (
7+
"testing"
8+
9+
"github.com/zeebo/assert"
10+
)
11+
12+
// Below the high-water mark, accrued returnable credit is released as a grant
13+
// once it reaches the threshold, and withheld (0) below it.
14+
func TestRecvWindowGrantsAtThreshold(t *testing.T) {
15+
w := newRecvWindow(1000, 100)
16+
assert.Equal(t, w.dispatched(60), int64(0)) // pending 60 < 100
17+
assert.Equal(t, w.dispatched(60), int64(120)) // pending 120 >= 100, buffered 120 < 1000
18+
assert.Equal(t, w.bufferedBytes(), int64(120))
19+
assert.Equal(t, w.dispatched(50), int64(0)) // pending reset by the grant; 50 < 100
20+
}
21+
22+
// At or above the high-water mark, grants are withheld even when the accrued
23+
// credit is past the threshold.
24+
func TestRecvWindowWithholdsAtOrAboveHighWater(t *testing.T) {
25+
w := newRecvWindow(100, 50)
26+
assert.Equal(t, w.dispatched(120), int64(0)) // buffered 120 >= 100 -> withhold
27+
assert.Equal(t, w.bufferedBytes(), int64(120))
28+
}
29+
30+
// Consuming buffered data drops below the high-water mark and flushes the
31+
// credit accrued while the gate was closed (resume-on-consume).
32+
func TestRecvWindowResumesOnConsume(t *testing.T) {
33+
w := newRecvWindow(100, 50)
34+
assert.Equal(t, w.dispatched(120), int64(0)) // withheld; pending 120
35+
assert.Equal(t, w.consumed(30), int64(120)) // buffered 90 < 100 -> flush pending
36+
assert.Equal(t, w.bufferedBytes(), int64(90))
37+
}
38+
39+
// Consuming while the accrued credit is below the threshold emits no grant.
40+
func TestRecvWindowConsumeBelowThresholdNoGrant(t *testing.T) {
41+
w := newRecvWindow(1000, 100)
42+
assert.Equal(t, w.dispatched(50), int64(0)) // pending 50 < 100
43+
assert.Equal(t, w.consumed(20), int64(0)) // buffered 30, pending 50 < 100
44+
assert.Equal(t, w.bufferedBytes(), int64(30))
45+
}
46+
47+
// Many small dispatches coalesce into few grants, each carrying the accrued
48+
// amount.
49+
func TestRecvWindowCoalescesGrants(t *testing.T) {
50+
w := newRecvWindow(1_000_000, 100)
51+
grants, granted := 0, int64(0)
52+
for i := 0; i < 10; i++ {
53+
if g := w.dispatched(30); g > 0 {
54+
grants++
55+
granted += g
56+
}
57+
}
58+
// 10 dispatches of 30 (300 bytes) with threshold 100 coalesce into 2 grants
59+
// of 120 each; the remaining 60 stays accrued.
60+
assert.Equal(t, grants, 2)
61+
assert.Equal(t, granted, int64(240))
62+
}

0 commit comments

Comments
 (0)