Skip to content

PYTHON-6136 Make pool checkout accounting transactional #3089

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
blink1073 merged 13 commits into from
Oct 2, 2026
Merged

PYTHON-6136 Make pool checkout accounting transactional #3089

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
d16d9b2
PYTHON-6136 Fix pool checkout accounting rollback on BaseException
blink1073 Oct 1, 2026
04dccc5
PYTHON-6136 Always roll back the size gate on checkout unwind
blink1073 Oct 1, 2026
2191a2f
PYTHON-6136 Add regression tests for interrupted checkout cleanup
blink1073 Oct 2, 2026
a909572
PYTHON-6136 Bind pool before the race test's try block
blink1073 Oct 2, 2026
dbbc615
PYTHON-6136 Roll back operation_count on every failed checkout
blink1073 Oct 2, 2026
40762fa
PYTHON-6136 Lock-protect the operation_count rollback on the gate path
blink1073 Oct 2, 2026
397e7a4
PYTHON-6136 Replace checkout rollback flags with an undo bitmask
blink1073 Oct 2, 2026
85e2bc3
PYTHON-6136 Wake maxConnecting waiters when a failed checkout frees a…
blink1073 Oct 2, 2026
e10f2db
PYTHON-6136 Retry interrupted gate notifications and fix comment mang…
blink1073 Oct 2, 2026
8bd3fa5
PYTHON-6136 Retry the maxConnecting cleanup notify after a kill
blink1073 Oct 2, 2026
52f72be
PYTHON-6136 Reword retry comment to avoid the synchro mangler
blink1073 Oct 2, 2026
72423b5
PYTHON-6136 Assert operation_count drains in the stress test
blink1073 Oct 2, 2026
1b577fb
PYTHON-6136 Document why _restore_counters always reacquires the lock
blink1073 Oct 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 92 additions & 36 deletions pymongo/asynchronous/pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,13 @@

_IS_SYNC = False

# Flags recording the counters a checkout has incremented, so a failed
# checkout can restore them exactly once (PYTHON-6136).
_UNDO_OPERATION_COUNT = 1
_UNDO_REQUESTS = 2
_UNDO_SOCKETS = 4
_UNDO_PENDING = 8


class AsyncConnection(_ConnectionTelemetryInfo):
"""Store a connection with some metadata.
Expand Down Expand Up @@ -969,8 +976,12 @@ async def _get_conn(
"Attempted to check out a connection from closed connection pool"
)

# Every counter increment sets a flag in ``applied`` so a failed
# checkout can restore them exactly once (PYTHON-6136).
applied = 0
async with self.lock:
self.operation_count += 1
applied |= _UNDO_OPERATION_COUNT

# Get a free socket or create one.
if _csot.get_timeout():
Expand All @@ -980,28 +991,32 @@ async def _get_conn(
else:
deadline = None

async with self.size_cond:
self._raise_if_not_ready(checkout_started_time, emit_event=True)
while not (self.requests < self.max_pool_size):
timeout = deadline - time.monotonic() if deadline else None
if not await _async_cond_wait(self.size_cond, timeout):
# Timed out, notify the next thread to ensure a
# timeout doesn't consume the condition.
if self.requests < self.max_pool_size:
self.size_cond.notify()
self._raise_wait_queue_timeout(checkout_started_time)
try:
async with self.size_cond:
self._raise_if_not_ready(checkout_started_time, emit_event=True)
self.requests += 1
while not (self.requests < self.max_pool_size):
timeout = deadline - time.monotonic() if deadline else None
if not await _async_cond_wait(self.size_cond, timeout):
# Timed out, notify the next thread to ensure a
# timeout doesn't consume the condition.
if self.requests < self.max_pool_size:
self.size_cond.notify()
self._raise_wait_queue_timeout(checkout_started_time)
self._raise_if_not_ready(checkout_started_time, emit_event=True)
self.requests += 1
applied |= _UNDO_REQUESTS
except BaseException:
await self._restore_counters(applied)
raise

# We've now acquired the semaphore and must release it on error.
conn = None
incremented = False
emitted_event = False
is_new_conn = False
try:
async with self.lock:
self.active_sockets += 1
incremented = True
applied |= _UNDO_SOCKETS
while conn is None:
# CMAP: we MUST wait for either maxConnecting OR for a socket
# to be checked back into the pool.
Expand All @@ -1022,6 +1037,7 @@ async def _get_conn(
conn = self.conns.popleft()
except IndexError:
self._pending += 1
applied |= _UNDO_PENDING
if conn: # We got a socket from the pool
if await self._perished(conn):
conn = None
Expand All @@ -1033,7 +1049,17 @@ async def _get_conn(
finally:
async with self._max_connecting_cond:
self._pending -= 1
self._max_connecting_cond.notify()
applied &= ~_UNDO_PENDING
notified = False
try:
self._max_connecting_cond.notify()
notified = True
finally:
if not notified:
# A kill landed inside notify() (a gevent
# yield point); retry so a waiting
# checkout is not stranded (PYTHON-6136).
self._max_connecting_cond.notify()

conn.active = True
# connect() already adds cancel_context for new connections; only add
Expand All @@ -1043,27 +1069,14 @@ async def _get_conn(
self.active_contexts.add(conn.cancel_context)
# Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
except BaseException:
if conn:
# We checked out a socket but authentication failed.
await conn.close_conn(ConnectionClosedReason.ERROR)
# Re-apply the accounting if a GreenletExit interrupts
# during the size_cond acquisition; during unwind gevent
# lets the re-acquire complete (PYTHON-6074).
accounted = False
try:
async with self.size_cond:
self.requests -= 1
if incremented:
self.active_sockets -= 1
accounted = True
self.size_cond.notify()
if conn is not None:
# We checked out a socket but authentication failed.
await conn.close_conn(ConnectionClosedReason.ERROR)
finally:
if not accounted:
async with self.size_cond:
self.requests -= 1
if incremented:
self.active_sockets -= 1
self.size_cond.notify()
# Restore the counters even if the cleanup above was
# interrupted (PYTHON-6136).
await self._restore_counters(applied)

if not emitted_event:
self._telemetry.checkout_failed(
Expand All @@ -1075,6 +1088,51 @@ async def _get_conn(

return conn

def _restore_applied(self, applied: int) -> None:
"""Restore the counters flagged in ``applied``. Caller holds ``size_cond``."""
if applied & _UNDO_OPERATION_COUNT:
self.operation_count -= 1
if applied & _UNDO_REQUESTS:
self.requests -= 1
if applied & _UNDO_SOCKETS:
self.active_sockets -= 1
if applied & _UNDO_PENDING:
self._pending -= 1

async def _restore_counters(self, applied: int) -> None:
"""Restore the counters a failed checkout incremented (PYTHON-6136).

Gevent grants the re-acquire during unwind (PYTHON-6074). A kill can
also land inside notify() (a yield point), so notifications are
tracked and retried separately from the counter restore.
"""
accounted = False
notified = 0
try:
async with self.size_cond:
self._restore_applied(applied)
accounted = True
if applied & _UNDO_REQUESTS:
# A pool slot was freed; wake the next waiting thread.
self.size_cond.notify()
Comment thread
Copilot marked this conversation as resolved.
notified |= _UNDO_REQUESTS
if applied & _UNDO_PENDING:
# A maxConnecting slot was freed; wake the next waiting thread.
self._max_connecting_cond.notify()
notified |= _UNDO_PENDING
finally:
# Always reacquired: `applied` keeps restore-only flags
# `notified` never tracks, so skipping needs a mask synced
# to the notify sites; uncontended acquires don't yield (PYTHON-6136).
async with self.size_cond:

@aclark4life aclark4life Oct 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like this locks whether it needs to or not, how about something like this instead:

      finally:
          # Both conditions below read only local variables, so the lock
          # is not needed to evaluate them; skip the acquire when there
          # is no leftover work.
          if not accounted or applied & ~notified:
              async with self.size_cond:
                  if not accounted:
                      self._restore_applied(applied)
                  missing = applied & ~notified
                  if missing & _UNDO_REQUESTS:
                      self.size_cond.notify()
                  if missing & _UNDO_PENDING:
                      self._max_connecting_cond.notify()

@blink1073 blink1073 Oct 2, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I verified that the current structure is correct and added a comment.

if not accounted:
self._restore_applied(applied)
Comment thread
blink1073 marked this conversation as resolved.
missing = applied & ~notified
if missing & _UNDO_REQUESTS:
self.size_cond.notify()
if missing & _UNDO_PENDING:
self._max_connecting_cond.notify()

def _checkin_apply(
self, conn: AsyncConnection, txn: bool, cursor: bool, forked: bool
) -> tuple[Optional[str], bool, bool]:
Expand Down Expand Up @@ -1121,9 +1179,7 @@ async def checkin(self, conn: AsyncConnection) -> None:
conn.pinned_cursor = False
self._pinned_sockets.discard(conn)
forked = self.pid != os.getpid()
# Re-apply the accounting if a gevent GreenletExit interrupts during
# the size_cond acquisition; gevent lets the re-acquire complete while
# unwinding (PYTHON-6074).
# Re-apply the accounting if a BaseException interrupts here (PYTHON-6074).
close_conn_reason: Optional[str] = None
emit_closed = False
accounted = False
Expand Down
128 changes: 92 additions & 36 deletions pymongo/synchronous/pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,13 @@

_IS_SYNC = True

# Flags recording the counters a checkout has incremented, so a failed
# checkout can restore them exactly once (PYTHON-6136).
_UNDO_OPERATION_COUNT = 1
_UNDO_REQUESTS = 2
_UNDO_SOCKETS = 4
_UNDO_PENDING = 8


class Connection(_ConnectionTelemetryInfo):
"""Store a connection with some metadata.
Expand Down Expand Up @@ -965,8 +972,12 @@ def _get_conn(
"Attempted to check out a connection from closed connection pool"
)

# Every counter increment sets a flag in ``applied`` so a failed
# checkout can restore them exactly once (PYTHON-6136).
applied = 0
with self.lock:
self.operation_count += 1
applied |= _UNDO_OPERATION_COUNT

# Get a free socket or create one.
if _csot.get_timeout():
Expand All @@ -976,28 +987,32 @@ def _get_conn(
else:
deadline = None

with self.size_cond:
self._raise_if_not_ready(checkout_started_time, emit_event=True)
while not (self.requests < self.max_pool_size):
timeout = deadline - time.monotonic() if deadline else None
if not _cond_wait(self.size_cond, timeout):
# Timed out, notify the next thread to ensure a
# timeout doesn't consume the condition.
if self.requests < self.max_pool_size:
self.size_cond.notify()
self._raise_wait_queue_timeout(checkout_started_time)
try:
with self.size_cond:
self._raise_if_not_ready(checkout_started_time, emit_event=True)
self.requests += 1
while not (self.requests < self.max_pool_size):
timeout = deadline - time.monotonic() if deadline else None
if not _cond_wait(self.size_cond, timeout):
# Timed out, notify the next thread to ensure a
# timeout doesn't consume the condition.
if self.requests < self.max_pool_size:
self.size_cond.notify()
self._raise_wait_queue_timeout(checkout_started_time)
self._raise_if_not_ready(checkout_started_time, emit_event=True)
self.requests += 1
applied |= _UNDO_REQUESTS
except BaseException:
self._restore_counters(applied)
raise

# We've now acquired the semaphore and must release it on error.
conn = None
incremented = False
emitted_event = False
is_new_conn = False
try:
with self.lock:
self.active_sockets += 1
incremented = True
applied |= _UNDO_SOCKETS
while conn is None:
# CMAP: we MUST wait for either maxConnecting OR for a socket
# to be checked back into the pool.
Expand All @@ -1018,6 +1033,7 @@ def _get_conn(
conn = self.conns.popleft()
except IndexError:
self._pending += 1
applied |= _UNDO_PENDING
if conn: # We got a socket from the pool
if self._perished(conn):
conn = None
Expand All @@ -1029,7 +1045,17 @@ def _get_conn(
finally:
with self._max_connecting_cond:
self._pending -= 1
self._max_connecting_cond.notify()
applied &= ~_UNDO_PENDING
notified = False
try:
self._max_connecting_cond.notify()
notified = True
finally:
if not notified:
# A kill landed inside notify() (a gevent
# yield point); retry so a waiting
# checkout is not stranded (PYTHON-6136).
self._max_connecting_cond.notify()

conn.active = True
# connect() already adds cancel_context for new connections; only add
Expand All @@ -1039,27 +1065,14 @@ def _get_conn(
self.active_contexts.add(conn.cancel_context)
# Catch KeyboardInterrupt, CancelledError, etc. and cleanup.
except BaseException:
if conn:
# We checked out a socket but authentication failed.
conn.close_conn(ConnectionClosedReason.ERROR)
# Re-apply the accounting if a GreenletExit interrupts
# during the size_cond acquisition; during unwind gevent
# lets the re-acquire complete (PYTHON-6074).
accounted = False
try:
with self.size_cond:
self.requests -= 1
if incremented:
self.active_sockets -= 1
accounted = True
self.size_cond.notify()
if conn is not None:
# We checked out a socket but authentication failed.
conn.close_conn(ConnectionClosedReason.ERROR)
finally:
if not accounted:
with self.size_cond:
self.requests -= 1
if incremented:
self.active_sockets -= 1
self.size_cond.notify()
# Restore the counters even if the cleanup above was
# interrupted (PYTHON-6136).
self._restore_counters(applied)

if not emitted_event:
self._telemetry.checkout_failed(
Expand All @@ -1071,6 +1084,51 @@ def _get_conn(

return conn

def _restore_applied(self, applied: int) -> None:
"""Restore the counters flagged in ``applied``. Caller holds ``size_cond``."""
if applied & _UNDO_OPERATION_COUNT:
self.operation_count -= 1
if applied & _UNDO_REQUESTS:
self.requests -= 1
if applied & _UNDO_SOCKETS:
self.active_sockets -= 1
if applied & _UNDO_PENDING:
self._pending -= 1

def _restore_counters(self, applied: int) -> None:
"""Restore the counters a failed checkout incremented (PYTHON-6136).

Gevent grants the re-acquire during unwind (PYTHON-6074). A kill can
also land inside notify() (a yield point), so notifications are
tracked and retried separately from the counter restore.
"""
accounted = False
notified = 0
try:
with self.size_cond:
self._restore_applied(applied)
accounted = True
if applied & _UNDO_REQUESTS:
# A pool slot was freed; wake the next waiting thread.
self.size_cond.notify()
notified |= _UNDO_REQUESTS
if applied & _UNDO_PENDING:
# A maxConnecting slot was freed; wake the next waiting thread.
self._max_connecting_cond.notify()
notified |= _UNDO_PENDING
finally:
# Always reacquired: `applied` keeps restore-only flags
# `notified` never tracks, so skipping needs a mask synced
# to the notify sites; uncontended acquires don't yield (PYTHON-6136).
with self.size_cond:
if not accounted:
self._restore_applied(applied)
Comment thread
blink1073 marked this conversation as resolved.
missing = applied & ~notified
if missing & _UNDO_REQUESTS:
self.size_cond.notify()
if missing & _UNDO_PENDING:
self._max_connecting_cond.notify()

def _checkin_apply(
self, conn: Connection, txn: bool, cursor: bool, forked: bool
) -> tuple[Optional[str], bool, bool]:
Expand Down Expand Up @@ -1117,9 +1175,7 @@ def checkin(self, conn: Connection) -> None:
conn.pinned_cursor = False
self._pinned_sockets.discard(conn)
forked = self.pid != os.getpid()
# Re-apply the accounting if a gevent GreenletExit interrupts during
# the size_cond acquisition; gevent lets the re-acquire complete while
# unwinding (PYTHON-6074).
# Re-apply the accounting if a BaseException interrupts here (PYTHON-6074).
close_conn_reason: Optional[str] = None
emit_closed = False
accounted = False
Expand Down
Loading
Loading