Skip to content

Fix redis connection leak - #5711

Merged
msureshkumar88 merged 12 commits into
IBM:mainfrom
aidbutlr:fix_redis_connection_leak
Jul 27, 2026
Merged

Fix redis connection leak#5711
msureshkumar88 merged 12 commits into
IBM:mainfrom
aidbutlr:fix_redis_connection_leak

Conversation

@aidbutlr

@aidbutlr aidbutlr commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

The redis.pubsub is not closing the connection to redis. This can consume all of the connections in the redis connection pool over time. Adding aclose to be called on exit.

Pull Request

🔗 Related Issue

Closes #5918


📝 Summary

The _plugin_invalidation_listener() function runs a long-lived loop that subscribes to a Redis pub/sub channel. Each iteration calls client.pubsub(), which acquires a dedicated connection from the shared pool (max_connections=50). When the iteration fails for any reason — network error, timeout, or Redis blip — the exception handler immediately slept and restarted the loop without ever calling pubsub.aclose().


📏 Reviewability

  • This PR has one clear purpose
  • The linked issue is not labeled triage
  • Unrelated bugs or improvements are tracked in separate issues/PRs
  • Tests are included with the code they validate
  • If AI-assisted, I understand and can explain the generated changes

🏷️ Type of Change

  • Bug fix
  • Feature / Enhancement
  • Documentation
  • Refactor
  • Chore (deps, CI, tooling)
  • Other (describe below)

🧪 Verification

List exact commands, screenshots, videos, logs, reproduction steps, or manual validation. If evidence is not feasible, explain why.

Check Command Status
Lint suite make lint
Unit tests make test
Coverage ≥ 80% make coverage

✅ Checklist

  • Code formatted (make black isort pre-commit)
  • Tests added/updated for changes
  • Documentation updated (if applicable)
  • No secrets or credentials committed

📓 Notes (optional)

Screenshots, design decisions, or additional context.

@aidbutlr
aidbutlr force-pushed the fix_redis_connection_leak branch from 0425545 to bdb9cba Compare July 20, 2026 16:47
@aidbutlr

Copy link
Copy Markdown
Contributor Author
image

@aidbutlr
aidbutlr marked this pull request as ready for review July 21, 2026 10:19
@aidbutlr
aidbutlr marked this pull request as draft July 21, 2026 10:59
@aidbutlr

Copy link
Copy Markdown
Contributor Author
image

@aidbutlr
aidbutlr force-pushed the fix_redis_connection_leak branch from 43a6ca8 to d1f7d23 Compare July 21, 2026 12:17
@aidbutlr
aidbutlr marked this pull request as ready for review July 21, 2026 12:24
@jonpspri jonpspri added this to the v1.0.7 milestone Jul 21, 2026
@gandhipratik203

Copy link
Copy Markdown
Collaborator

Thanks for making this change. aclose() in a finally with the wait_for(timeout=2.0) bound is in line with the fix we need.

F1: add a regression test. Nothing exercises the new cleanup yet. The existing listener tests use pubsub = MagicMock(), so pubsub.aclose() returns a non-awaitable, wait_for(...) raises, and except Exception swallows it. The cleanup runs as a silent no-op in tests, so the fix ships unverified.

A small pin closes it: AsyncMock for aclose(), force the reconnect path (and ideally a cancel case), assert aclose() was awaited.

Once that's in we're good to go.

@gandhipratik203 gandhipratik203 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The regression test could be a good addiiton. Otherwise, looks good to me!

@msureshkumar88 msureshkumar88 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for tracking down this connection leak — the root cause analysis is spot on, and the fix itself (reset pubsub per iteration, close it in a finally, guard the close with a timeout) is correct. I traced it through redis-py's PubSub.aclose()/__aexit__ and confirmed the finally runs on every exit path (happy CancelledError, retry-on-exception, and the early continue when Redis is unavailable), so the leak is genuinely fixed. The asyncio.wait_for(..., timeout=2.0) guard is a nice touch too — aclose() can stall if a concurrent reader still holds the connection's transport (see redis/redis-py#3941), so this isn't just a test-hang workaround, it's real production hardening.

Two small things I'd like to see before merging, both scoped to the same block:

1. The new except Exception: pass is a silent swallow.
Every other except Exception in this file logs and carries # noqa: BLE001 (lines 231, 282, 409, 481, 547) — this is the one exception. Given the whole point of this PR is visibility into connection lifecycle issues, it'd help future debugging to at least log it:

except Exception as exc:  # noqa: BLE001
    _logger.debug("Pub/sub: failed to close pubsub cleanly (%s)", exc)

2. No test actually exercises the new cleanup path.
TestListenerLoopBranches::test_listener_subscribes_and_dispatches_one_message_then_cancels uses a bare MagicMock() for pubsub with no aclose mock set. I checked empirically — pubsub.aclose() returns a non-awaitable MagicMock, asyncio.wait_for raises TypeError on it, and that's silently caught by the very except Exception: pass above. The test passes identically whether or not the finally block exists, so the fix's core behavior is currently unverified. Could you add:

  • An AsyncMock-backed pubsub.aclose + an exception raised inside listen() → assert aclose.assert_awaited_once()
  • Same assertion on the CancelledError/break exit path
  • A case where aclose() hangs, to confirm the 2s wait_for timeout actually bounds it (this is the exact scenario the "reworked to avoid test execution hanging" commit was addressing, and it'd be good to have it pinned down)

Nothing else blocking — no breaking change, implementation matches the PR description, and the unrelated import-reorder in the diff is harmless formatting. Happy to re-review once these are in.

@aidbutlr

Copy link
Copy Markdown
Contributor Author

I will make the requested updates. Thanks

@gandhipratik203 gandhipratik203 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for making the suggested changes. LGTM!

@msureshkumar88 msureshkumar88 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for iterating on this — the root cause analysis holds up, and the final approach (reset pubsub per iteration, close it in a finally, guard the close with asyncio.wait_for(timeout=2.0)) matches the same pattern already used in registry_cache.py, session_registry.py, and runtime_state.py for the same known redis-py stall (redis/redis-py#3941). Good that the earlier async with client.pubsub() attempt was reworked once it caused test hangs — the guarded manual close is the more correct fix here, not just a workaround. The added logging on close failure and the two new tests (cancellation path, reconnect-after-exception path) are solid.

One blocking item, carried over from earlier feedback on this PR:

Missing regression test for the timeout guard itself. The whole reason for the "Reworked fix to avoid test execution hanging" commit was that aclose() can hang — that's exactly what wait_for(timeout=2.0) is there to bound. Neither new test exercises that path; both use an AsyncMock that resolves immediately. Without a test where aclose() never completes, the timeout guard's behavior is unverified — the finally block would pass its tests identically whether the timeout parameter is 2.0 or removed entirely. Could you add a case where pubsub.aclose hangs (e.g. await asyncio.sleep(999)) and assert the listener proceeds past the finally within a few seconds?

A few smaller, non-blocking suggestions while in this area:

  1. The finally block's comment — "Using wait_for shields aclose() from an in-flight CancelledError" — isn't quite accurate. asyncio.wait_for doesn't shield the wrapped coroutine from external cancellation (that's what asyncio.shield() is for); if the enclosing task is cancelled while inside this wait_for, the cancellation propagates into aclose() too. The real value here is bounding a coroutine that can otherwise hang (the redis-py stall), which is a timeout guard, not a cancellation shield — worth tightening the comment so it doesn't mislead the next reader.
  2. runtime_state.py's _cleanup_pubsub() logs asyncio.TimeoutError at debug and other exceptions separately, while this new code logs everything (including the expected timeout case) at _logger.error. Might be worth splitting those two cases the same way, so a routine stall during shutdown doesn't read as an ERROR-level incident every time.

None of this is a design concern — the fix itself is right. Happy to re-review once the timeout test is in.

aidbutlr added 3 commits July 27, 2026 12:06
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
aidbutlr added 4 commits July 27, 2026 12:06
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
@aidbutlr
aidbutlr force-pushed the fix_redis_connection_leak branch from 456b3d6 to ca39a83 Compare July 27, 2026 11:06
@aidbutlr
aidbutlr force-pushed the fix_redis_connection_leak branch from a92e52e to 3b3ed4b Compare July 27, 2026 12:08
aidbutlr added 2 commits July 27, 2026 13:09
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com>
@aidbutlr
aidbutlr force-pushed the fix_redis_connection_leak branch from 3b3ed4b to 20ed493 Compare July 27, 2026 12:10

@gandhipratik203 gandhipratik203 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM!

msureshkumar88
msureshkumar88 previously approved these changes Jul 27, 2026

@msureshkumar88 msureshkumar88 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

E2E verification

I built and ran a live end-to-end test against a real redis:7 container (Docker, no mocking of Redis, no mocked pool internals) to reproduce and verify the fix for the connection leak in _plugin_invalidation_listener().

Setup

  • docker run -d --rm --name e2e_redis_5711 -p 16399:6379 redis:7
  • Started the real _plugin_invalidation_listener() coroutine with set_shared_redis_provider(get_redis_client) wired to the real client, REDIS_MAX_CONNECTIONS=5 to make pool exhaustion observable quickly.
  • Compared PR head (20ed493ff) against the pre-fix baseline (merge-base 3fe3035e4, checked out in a separate git worktree) with everything else held identical.

Fault injection — what actually reaches the buggy branch

I first tried killing the listener's subscriber connection via CLIENT KILL to simulate a network blip. That does not exercise the bug: redis-py's own connection-level retry (PubSub._execute_reconnect) transparently reconnects the same Connection object before the error ever reaches _plugin_invalidation_listener's except Exception: — confirmed by connection id() staying identical across kills, and in_use counts flat even on the pre-fix baseline.

The real trigger is a message-handler exception: _GlobalToggleMsg/_ModeChangeMsg in _handle_invalidation_message() have no per-case try/except (unlike _TeamBindingChangeMsg/_BindingChangeMsg), so any exception there propagates straight out of the listener's async for loop — this matches the exact seam the PR's own new regression tests target. I reproduced this with 100% real Redis I/O (real subscribe(), real PUBLISH from a second admin connection, real listen()) and only substituted _handle_invalidation_message to raise on a real delivered message.

Result — pre-fix baseline (merge-base 3fe3035e4)

[baseline] pool max_connections=5 available=1 in_use=0
[iter 0] in_use=1  handler_hits=1
[iter 1] in_use=2  handler_hits=1
[iter 3] in_use=3  handler_hits=2
[iter 5] in_use=4  handler_hits=3
[iter 7] in_use=5  handler_hits=4
[final] in_use=5 (pool at max_connections, exhausted)
RESULT: LEAK DETECTED

Result — PR head (20ed493ff)

[baseline] pool max_connections=5 available=1 in_use=0
[iter 0..7] in_use=1 flat across all 4 real handler failures
[final] in_use=1
RESULT: NO LEAK

client.connection_pool._in_use_connections grows unbounded pre-fix and stays flat post-fix under identical real-Redis fault conditions — the fix closes the leak as described in #5918.

Regression suite & blast radius

  • tests/unit/mcpgateway/plugins/test_plugin_runtime_management.py: 98 passed, 0 failed (.venv/bin/python -m pytest ... -q, exit 0).
  • mypy mcpgateway/plugins/__init__.py: no new errors in the changed lines (490–530); all pre-existing findings are unrelated to this diff.
  • Checked every other pubsub() consumer in the codebase (runtime_state.py, cache/registry_cache.py, cache/session_registry.py, transports/server_event_bus.py, services/cancellation_service.py, services/event_service.py, services/session_affinity.py) — all already use the same finally: await asyncio.wait_for(pubsub.aclose(), timeout=...) (or async with) pattern. This PR brings plugins/__init__.py in line with the established convention rather than introducing a new one.
  • No schema/migration changes, no API surface changes, import-reorder is cosmetic/unrelated.

Approving — fix verified end-to-end against a real Redis instance under the actual failure condition that triggers it, tests pass, and no regressions found in surrounding code.

Signed-off-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com>
@msureshkumar88
msureshkumar88 dismissed stale reviews from gandhipratik203 and themself via e3e32e5 July 27, 2026 14:07
Signed-off-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com>

@msureshkumar88 msureshkumar88 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGMT

@msureshkumar88
msureshkumar88 added this pull request to the merge queue Jul 27, 2026
Merged via the queue into IBM:main with commit bf9130a Jul 27, 2026
36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG]: Redis connection leak in mcpgateway/plugins/__init__.py

4 participants