cockroachdb / cockroachdb/cockroach

drpc: context cancellation improvements

Open
#160,900 0 comments 0 reactions 1 assignee Claimed by @suj-krishnan View on GitHub
C-enhancement T-db-server
Dominant language
Go
Stars
32.5k
Forks
4.1k
PR merge metrics
PR metrics pending

Description

When a drpc call is made, the context that is passed to the call (via `Invoke`/`NewStream`) eventually becomes the stream context ([Ref 1](https://github.com/cockroachdb/drpc/blob/b4a71786571655e377ebe8f392db64a2ed5f1a18/drpcconn/conn.go#L112), [Ref 2](https://github.com/cockroachdb/drpc/blob/b4a71786571655e377ebe8f392db64a2ed5f1a18/drpcmanager/manager.go#L311)) . The drpc framework code waits for cancelations on this context in [monitorStream](https://github.com/cockroachdb/drpc/blob/b4a71786571655e377ebe8f392db64a2ed5f1a18/drpcpool/conn.go#L106) and [manageStream](https://github.com/cockroachdb/drpc/blob/b4a71786571655e377ebe8f392db64a2ed5f1a18/drpcmanager/manager.go#L353), so that the connection can be returned to the pool and the stream cleanly terminated during cancellation scenarios.
In the case of a pooled drpc connection, the returned `poolConn` is instantiated via [DialDRPC](https://github.com/cockroachdb/cockroach/blob/8c5e0401c583061e4191d2598f3191f9ebd62713/pkg/rpc/drpc.go#L61). Cancellation and cleanup for this connection is handled in the [runHeartbeatUntilFailure](https://github.com/cockroachdb/cockroach/blob/8c5e0401c583061e4191d2598f3191f9ebd62713/pkg/rpc/peer.go#L553) routine. When the context associated with the `pooledConn` (parent context) is different from the stream context (child context), cancellation of this context does not cancel the associated active streams. The drpc `monitorStream`/`manageStream` goroutines also never terminate and therefore are reported as leaked.

Possible solutions:
- Save the parent context in the `poolConn` object and pass it to `drpc.Conn::Invoke/NewStream`, so that streams are associated with the same context as the `poolConn`.
- Save the parent context in the `poolConn` object and pass it to `drpc.Conn::Invoke/NewStream`, which in turn, passes it to the manager object. The `manageStream` code can then wait on this parent context cancelation in addition to the stream context.

Jira issue: CRDB-58609

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.