cockroachdb / cockroachdb/cockroach
drpc: `pool.Close` should make the pool unusable
- Dominant language
- Go
- Stars
- 32.5k
- Forks
- 4.1k
- PR merge metrics
- PR metrics pending
Description
(_logging here to make sure we don't lose the track_)
We ran into a case where calling `pool.Close()` didn't fully drain or shut down the pool. Specifically, we observed that `pool.Len()` was non-zero even after calling `pool.Close()`. This was unexpected.
Turns out, this happens if a client does `stream.CloseSend()` but doesn't wait on `stream.Context().Done()`. If we immediately call `pool.Close()` after that, the pool gets reset, but in parallel, `poolConn.manageStream` (via `monitorStream`) may still run and put the connection back into the pool.
Here's a simplified repro:
```go
for i := range numStreams {
streamID := i
stream, err := client.PingS(ctx)
require.NoError(t, err, fmt.Errorf("stream %d failed to start", streamID))
defer func() { _ = stream.Close() }()
for k := range numMessages {
payload := fmt.Appendf(nil, "stream-%d-msg-%d", streamID, k)
err = stream.Send(&pb.Ping{Payload: payload})
require.NoError(t, err, fmt.Errorf("stream %d failed to send", streamID))
pong, err := stream.Recv()
require.NoError(t, err, fmt.Errorf("stream %d failed to recv", streamID))
require.True(t, bytes.Equal(pong.Payload, payload),
"stream %d payload mismatch: got %q, want %q",
streamID, pong.Payload, payload)
}
err = stream.CloseSend()
require.NoError(t, err, fmt.Errorf("stream %d CloseSend error", streamID))
// <-stream.Context().Done() // <-- this would make the test pass
require.Equal(t, 1, pool.Len())
}
require.Equal(t, 1, pool.Len())
require.NoError(t, pool.Close())
require.Equal(t, 0, pool.Len()) // <-- this assert can fail
```
Relevant code from the pool:
```go
func (p *poolConn[K, V]) monitorStream(stream drpc.Stream, conn V, done *drpcsignal.Chan) {
<-stream.Context().Done()
p.pool.Put(p.key, conn)
done.Close()
}
```
This means even after closing the pool, a race can put a connection back into it.
We probably need to gate the `Put` call in `monitorStream` behind a check to see if the pool is still valid/open, or add some kind of guard in `Put()` itself. Either way, `Close()` should make it impossible for future `Put()`s to succeed.
Jira issue: CRDB-49568
Contributor guide
Research direction
Read poolConn.monitorStream, Put, and Close, then run the simplified reproduction involving CloseSend without waiting on stream.Context().Done(). Trace the concurrent lifecycle and verify that after pool.Close() returns, pool.Len() stays zero even if a monitorStream finishes later.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- go
- Domain
- backend, distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100