cockroachdb / cockroachdb/cockroach

drpc: `pool.Close` should make the pool unusable

Open
#144,838 1 comment 0 reactions 0 assignees View on GitHub
A-server-networking C-enhancement T-db-server
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.