Skip to content

Commit bf9130a

Browse files
aidbutlrSuresh Kumar Moharajan
andauthored
Fix redis connection leak (#5711)
* Fix redis connection leak Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Fix redis connection leak Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Reworked fix to avoid test execution hanging Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Added log message and unit tests to cover new code branches Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Updates to address comments Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Updates to address comments Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Updated detect secrets file Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * Update secrets baseline Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * updated secrets baseline Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> * chore: remove changelog entry for redis connection leak fix Signed-off-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com> * chore: fix secrets baseline line numbers after changelog edit Signed-off-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com> --------- Signed-off-by: Aidan Butler <aidbutlr@ie.ibm.com> Signed-off-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com> Co-authored-by: Suresh Kumar Moharajan <suresh.kumar.m@ibm.com>
1 parent 91f8920 commit bf9130a

3 files changed

Lines changed: 158 additions & 3 deletions

File tree

.secrets.baseline

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
"files": "(?x)( package-lock\\.json$ |Cargo\\.lock$ |uv\\.lock$ |go\\.sum$ |mcpgateway/sri_hashes\\.json$ )|^\\.secrets\\.baseline$",
44
"lines": null
55
},
6-
"generated_at": "2026-07-24T15:22:57Z",
6+
"generated_at": "2026-07-27T14:40:12Z",
77
"plugins_used": [
88
{
99
"name": "AWSKeyDetector"

mcpgateway/plugins/__init__.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import logging
2424
import random
2525
import time
26-
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
26+
from typing import Any, Callable, Literal, Optional, TYPE_CHECKING, Union
2727

2828
# Third-Party
2929
from cpex.framework import ObservabilityProvider, TenantPluginManager
@@ -102,9 +102,10 @@ class _TeamBindingChangeMsg(BaseModel):
102102

103103
def _get_invalidation_hmac_key() -> Optional[bytes]:
104104
"""Return the HMAC signing key for pub/sub messages, or None if unconfigured."""
105-
from mcpgateway.config import settings # pylint: disable=import-outside-toplevel
106105
from pydantic import SecretStr # pylint: disable=import-outside-toplevel
107106

107+
from mcpgateway.config import settings # pylint: disable=import-outside-toplevel
108+
108109
secret = settings.jwt_secret_key
109110
if not secret:
110111
return None
@@ -491,6 +492,7 @@ async def _plugin_invalidation_listener() -> None:
491492
consecutive_failures = 0
492493

493494
while True:
495+
pubsub = None
494496
try:
495497
client = await _redis()
496498
if not client:
@@ -524,6 +526,17 @@ async def _plugin_invalidation_listener() -> None:
524526
)
525527
await asyncio.sleep(delay)
526528
backoff = min(max_backoff, backoff * 2)
529+
finally:
530+
# Always release the pubsub connection back to the pool.
531+
# wait_for bounds aclose() so a stalled redis-py connection cannot
532+
# block the finally clause indefinitely.
533+
if pubsub is not None:
534+
try:
535+
await asyncio.wait_for(pubsub.aclose(), timeout=2.0)
536+
except asyncio.TimeoutError:
537+
_logger.debug("Plugin invalidation listener: pubsub aclose timed out")
538+
except Exception as exc:
539+
_logger.error("Plugin invalidation listener: pubsub aclose failed (%s)", exc)
527540

528541

529542
async def start_plugin_invalidation_listener() -> None:

tests/unit/mcpgateway/plugins/test_plugin_runtime_management.py

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1705,6 +1705,7 @@ async def _fake_handle(msg):
17051705

17061706
pubsub = MagicMock()
17071707
pubsub.subscribe = AsyncMock()
1708+
pubsub.aclose = AsyncMock()
17081709

17091710
class _Listen:
17101711
def __aiter__(self):
@@ -1724,6 +1725,147 @@ async def __anext__(self):
17241725
pubsub.subscribe.assert_awaited_once()
17251726
assert len(seen) == 1
17261727

1728+
@pytest.mark.asyncio
1729+
async def test_pubsub_aclose_called_on_cancellation(self, monkeypatch):
1730+
"""``pubsub.aclose()`` must be called when CancelledError propagates out of the listen loop.
1731+
1732+
Regression guard: if the ``finally: await pubsub.aclose()`` block is ever
1733+
removed, the pubsub connection would be leaked on every cancellation.
1734+
"""
1735+
import mcpgateway.plugins as framework
1736+
1737+
async def _fake_handle(msg):
1738+
raise asyncio.CancelledError()
1739+
1740+
monkeypatch.setattr(framework, "_handle_invalidation_message", _fake_handle)
1741+
1742+
pubsub = MagicMock()
1743+
pubsub.subscribe = AsyncMock()
1744+
pubsub.aclose = AsyncMock()
1745+
1746+
class _Listen:
1747+
def __aiter__(self):
1748+
return self
1749+
1750+
async def __anext__(self):
1751+
return {"type": "message", "data": "{}"}
1752+
1753+
pubsub.listen = MagicMock(return_value=_Listen())
1754+
1755+
client = MagicMock()
1756+
client.pubsub = MagicMock(return_value=pubsub)
1757+
monkeypatch.setattr(framework, "_redis", AsyncMock(return_value=client))
1758+
1759+
await framework._plugin_invalidation_listener()
1760+
1761+
# aclose() must have been awaited exactly once — proving the connection
1762+
# was released via the finally block before the listener returned.
1763+
pubsub.aclose.assert_awaited_once()
1764+
1765+
@pytest.mark.asyncio
1766+
async def test_pubsub_aclose_called_on_reconnect_after_exception(self, monkeypatch):
1767+
"""``pubsub.aclose()`` must be called when a non-cancellation exception escapes the listen loop.
1768+
1769+
A transient Redis error must not leak the pubsub handle — the finally
1770+
block must close it before the backoff sleep and reconnection attempt.
1771+
"""
1772+
import mcpgateway.plugins as framework
1773+
1774+
iterations: list[int] = []
1775+
1776+
async def _fake_handle(msg):
1777+
iterations.append(1)
1778+
raise RuntimeError("transient broker error")
1779+
1780+
monkeypatch.setattr(framework, "_handle_invalidation_message", _fake_handle)
1781+
1782+
pubsub = MagicMock()
1783+
pubsub.subscribe = AsyncMock()
1784+
pubsub.aclose = AsyncMock()
1785+
1786+
class _Listen:
1787+
def __aiter__(self):
1788+
return self
1789+
1790+
async def __anext__(self):
1791+
return {"type": "message", "data": "{}"}
1792+
1793+
pubsub.listen = MagicMock(return_value=_Listen())
1794+
1795+
client = MagicMock()
1796+
client.pubsub = MagicMock(return_value=pubsub)
1797+
monkeypatch.setattr(framework, "_redis", AsyncMock(return_value=client))
1798+
1799+
# Stop the reconnection loop after the first backoff sleep.
1800+
async def _fake_sleep(_delay):
1801+
raise asyncio.CancelledError()
1802+
1803+
monkeypatch.setattr(framework.asyncio, "sleep", _fake_sleep)
1804+
1805+
# CancelledError raised inside _fake_sleep propagates out of the
1806+
# except-Exception handler, through the finally block, and escapes the
1807+
# coroutine — that is the intended termination path when cancellation
1808+
# races with the backoff sleep.
1809+
try:
1810+
await framework._plugin_invalidation_listener()
1811+
except asyncio.CancelledError:
1812+
pass
1813+
1814+
# One dispatch happened, then an exception fired.
1815+
assert len(iterations) == 1
1816+
# aclose() must have been awaited once before the backoff sleep.
1817+
pubsub.aclose.assert_awaited_once()
1818+
1819+
@pytest.mark.asyncio
1820+
async def test_pubsub_aclose_timeout_guard_prevents_hang(self, monkeypatch):
1821+
"""``wait_for(aclose(), timeout=2.0)`` lets the listener exit even when aclose() hangs.
1822+
1823+
Regression guard for the timeout introduced to bound a hanging aclose():
1824+
if the ``timeout=2.0`` parameter were removed (or the wait_for dropped),
1825+
the finally block would suspend indefinitely and the listener task would
1826+
never return.
1827+
1828+
The test replaces aclose() with a coroutine that sleeps for 999 s and
1829+
asserts the entire listener finishes well inside that window.
1830+
"""
1831+
import time
1832+
1833+
import mcpgateway.plugins as framework
1834+
1835+
async def _fake_handle(msg):
1836+
raise asyncio.CancelledError()
1837+
1838+
monkeypatch.setattr(framework, "_handle_invalidation_message", _fake_handle)
1839+
1840+
async def _hanging_aclose():
1841+
await asyncio.sleep(999)
1842+
1843+
pubsub = MagicMock()
1844+
pubsub.subscribe = AsyncMock()
1845+
pubsub.aclose = _hanging_aclose # not an AsyncMock — returns a real coroutine
1846+
1847+
class _Listen:
1848+
def __aiter__(self):
1849+
return self
1850+
1851+
async def __anext__(self):
1852+
return {"type": "message", "data": "{}"}
1853+
1854+
pubsub.listen = MagicMock(return_value=_Listen())
1855+
1856+
client = MagicMock()
1857+
client.pubsub = MagicMock(return_value=pubsub)
1858+
monkeypatch.setattr(framework, "_redis", AsyncMock(return_value=client))
1859+
1860+
start = time.monotonic()
1861+
await framework._plugin_invalidation_listener()
1862+
elapsed = time.monotonic() - start
1863+
1864+
# The listener must complete well before aclose()'s 999 s sleep expires.
1865+
# A generous upper bound of 5 s covers slow CI environments while still
1866+
# being far shorter than the hang duration.
1867+
assert elapsed < 5.0, f"Listener took {elapsed:.2f}s — timeout guard may be missing"
1868+
17271869

17281870
# ---------------------------------------------------------------------------
17291871
# HMAC signing and verification tests

0 commit comments

Comments
 (0)