Skip to content

Commit 482d67b

Browse files
drpcmanager: test that context cancellation wakes a parked sender
Add an integration test for the production wake path of a sender parked on flow-control credit. A flow-control-enabled client sends a message larger than its window to a server that has no flow control (so it never returns credit); the send parks, and cancelling the RPC context wakes it. This exercises manageStream -> Cancel -> terminate -> sendWindow.close end to end, which the unit tests previously covered only in pieces (the primitive's ctx path, and stream-level Cancel). Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>
1 parent b8e91f1 commit 482d67b

1 file changed

Lines changed: 69 additions & 0 deletions

File tree

drpcmanager/flow_control_test.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"io"
1010
"net"
1111
"testing"
12+
"time"
1213

1314
"github.com/zeebo/assert"
1415

@@ -70,3 +71,71 @@ func TestManager_FlowControlEndToEnd(t *testing.T) {
7071

7172
ctx.Wait()
7273
}
74+
75+
// Cancelling an RPC's context wakes a sender parked on flow-control credit. The
76+
// client has flow control with a small window; the server does not, so it never
77+
// returns credit and the client's oversized send parks. This exercises the
78+
// production wake path: manageStream sees the cancellation and calls
79+
// Cancel -> terminate -> sendWindow.close.
80+
func TestManager_FlowControlContextCancelWakesParkedSend(t *testing.T) {
81+
tr := drpctest.NewTracker(t)
82+
defer tr.Close()
83+
84+
cconn, sconn := net.Pipe()
85+
defer func() { _ = cconn.Close() }()
86+
defer func() { _ = sconn.Close() }()
87+
88+
cman := NewWithOptions(cconn, Client, Options{Stream: drpcstream.Options{
89+
SplitSize: 64 << 10,
90+
FlowControl: drpcstream.FlowControl{
91+
Enabled: true,
92+
StreamWindow: 128 << 10, // two frames of initial credit
93+
HighWater: 1 << 20,
94+
GrantThreshold: 64 << 10,
95+
},
96+
}})
97+
defer func() { _ = cman.Close() }()
98+
// Server has no flow control, so it never returns credit.
99+
sman := New(sconn, Server)
100+
defer func() { _ = sman.Close() }()
101+
102+
// Accept the stream and drain; the message never completes (the sender
103+
// parks mid-message), so RawRecv just blocks until teardown.
104+
tr.Run(func(ctx context.Context) {
105+
stream, _, err := sman.NewServerStream(ctx)
106+
if err != nil {
107+
return
108+
}
109+
defer func() { _ = stream.Close() }()
110+
for {
111+
if _, err := stream.RawRecv(); err != nil {
112+
return
113+
}
114+
}
115+
})
116+
117+
streamCtx, cancel := context.WithCancel(context.Background())
118+
stream, err := cman.NewClientStream(streamCtx, "rpc")
119+
assert.NoError(t, err)
120+
assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("rpc")))
121+
122+
// A 512 KiB message exceeds the 128 KiB window; once the initial credit is
123+
// spent the send parks waiting for grants that never come.
124+
done := make(chan error, 1)
125+
go func() { done <- stream.RawWrite(drpcwire.KindMessage, make([]byte, 512<<10)) }()
126+
127+
select {
128+
case <-done:
129+
t.Fatal("send returned before it could park on credit")
130+
case <-time.After(50 * time.Millisecond):
131+
}
132+
133+
cancel()
134+
135+
select {
136+
case err := <-done:
137+
assert.Error(t, err)
138+
case <-time.After(time.Second):
139+
t.Fatal("parked send did not wake on context cancellation")
140+
}
141+
}

0 commit comments

Comments
 (0)