fix: don't leak server-side state when a future is cancelled - #1374
Open
bikeshedder wants to merge 4 commits into
Open
bikeshedder wants to merge 4 commits into
bikeshedder wants to merge 4 commits into
Conversation
Dropping an in-flight future can strand server-side state. Three sites
share the same shape: the message that creates the object is written
before the future first yields, but the value whose Drop cleans it up is
only constructed after the last response has been read.
- prepare() allocates a statement name and sends Parse/Describe/Sync,
but Statement - whose Drop is the only emitter of Close(b'S', ...) -
is built only at the end. The statement stays for the session.
- bind() has the identical shape for portals and Close(b'P', ...).
- _savepoint() sends SAVEPOINT and builds the nested Transaction that
owns the ROLLBACK TO only afterwards.
prepare() is reached by every method taking a query string, since
ToStatement for str prepares under the hood, so this fires for a plain
timeout(d, client.query("SELECT ...", &[])) - and by prepare() itself
via the recursive typeinfo lookup.
These tests fail. They are committed separately from the fix so the
failures can be reproduced by checking out this commit.
Controls are included so a nonzero count cannot be explained by the
catalog views lagging: a completed prepare/bind whose Statement/Portal
is then dropped leaves nothing behind, and a future dropped before its
first poll leaks nothing because send() never ran.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
StatementInner::drop and Portal's Inner::drop build the same Close message with only the target byte differing. Pull it into one place so there is a single definition of how a Close is framed - close, sync, and send as one request - rather than two copies that can drift apart. Pure refactor, no behaviour change. The early return for unnamed statements stays where it was, in StatementInner::drop. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Own each server-side object from the moment it might exist, rather than from the moment the value that owns it is constructed. prepare() and bind() both allocate a name and put the message that creates the object on the wire before awaiting anything, but Statement and Portal - whose Drop impls are the only emitters of Close - are built only after the last response has been read. Dropping the future in between, or returning early via `?`, stranded the object with nothing left that knows its name. _savepoint() has the same shape: it sends SAVEPOINT and builds the nested Transaction that owns the cleanup only afterwards. Add DropGuard, a general arm/disarm guard that runs a closure unless disarmed, and use it at all three sites: armed right after the request is sent, disarmed once the owning value takes over. For prepare() and bind(), arming before the creating message is confirmed is deliberate. Waiting for confirmation would miss the common case, cancelled after the request went out but before any response came back. It is safe because closing a name that does not exist is explicitly not an error in the protocol, so the cost is at most one redundant Close on a failed prepare. Ordering is safe too, since the Close travels the same request channel as the message that created the object and the connection processes requests FIFO. Because the guard covers `?` exits as well as cancellation, and because each recursive prepare arms its own, this also fixes the typeinfo lookups that prepare() performs for non-builtin types. _savepoint() needs the opposite error policy: RELEASE against a savepoint that does not exist is an error and would poison the enclosing transaction, so a spurious cleanup is not free the way a spurious Close is, and the guard disarms on failure rather than propagating first. Arming before the send is still safe because no suspension point separates constructing the guard from the send inside batch_execute - if the guard is ever dropped armed then the SAVEPOINT was queued and will run. RELEASE rather than ROLLBACK TO, since nothing ran inside the savepoint and ROLLBACK TO would leave it defined. All tests added two commits ago now pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit added a general arm/disarm guard. TransactionBuilder ::start already hand-rolls the same thing as a local RollbackIfNotDone struct, so fold it in. Behaviour is unchanged, including the error policy: an error from batch_execute still propagates with the guard armed, so a failed START TRANSACTION is still followed by a fire-and-forget ROLLBACK. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
Author
|
btw. please don't squash and merge this PR. The commits are stacked in a meaningful way:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
While working on deadpool-postgres I discovered memory leaks caused by cancelled futures. Most of the code - was written by Claude Code but meticulously reviewed by me.
The only code where I let Claude do it's thing and didn't bother improving upon were the tests. This part of the code is AI slop at its best but shows the problem to the point.
You can just check out the test(tokio): add failing tests for cancellation resource leaks commit and see the leak by running the test suite.
This PR also includes a bit of refactoring which made patching the actual bug a bit more straight forward.
Fun fact: This kind of error plagued the first few versions of
deadpool. I learned the hard way that every.awaitcould mean that the code following this expression will never be executed if the future gets dropped. The only way to clean up reliably is by utilizing theDroptrait.