From 5e61737ab6ddb1b38afe54f7bd621f7bf081c8cf Mon Sep 17 00:00:00 2001 From: maximeboudier Date: Thu, 23 Apr 2026 11:00:47 +0200 Subject: [PATCH] fix(redis): cache GCP IAM token to prevent async event loop blocking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem GCPIAMCredentialProvider.get_credentials() calls _generate_gcp_iam_access_token on every Redis connection establishment. This function performs synchronous HTTP and gRPC calls (google-auth + google-cloud-iam) which block Python's asyncio event loop while running. Under concurrent load (e.g. connection pool warm-up, parallel health checks), multiple connections are established simultaneously, each triggering an independent blocking IAM token refresh. These refreshes serialise behind each other inside the single-threaded event loop, causing individual Redis spans to take 20-25 seconds instead of milliseconds. Observed in production via Datadog APM: a single INCRBYFLOAT Redis span took 25.6 seconds (90% of a 28.4s trace), with GCP metadata + GenerateAccessToken gRPC calls visible inside the span. This cascaded into aiohttp SocketTimeoutError on upstream LLM API calls — not because the upstream was slow, but because the event loop was frozen and the 30-second sock_read timer fired on a connection that was never given CPU time. ## Fix Add a module-level token cache (dict keyed by service account, value is (token, expiry_monotonic)). _get_cached_gcp_iam_token() returns the cached token on cache hit (no I/O), and refreshes only when expired using double-checked locking so only one thread performs the network round-trip. GCP IAM tokens are valid for 1 hour; the cache TTL is set to 55 minutes (_GCP_IAM_TOKEN_TTL_SECONDS = 3300) to refresh safely before expiry. The cache is shared across all GCPIAMCredentialProvider instances for the same service account, so N concurrent Redis connections on the same pod share a single token and avoid N concurrent blocking refreshes. get_credentials_async() already used asyncio.to_thread (non-blocking), and is updated to call _get_cached_gcp_iam_token so it also benefits from caching. ## Tests - Updated existing test that expected a fresh token on every call to reflect the new caching behaviour. - Added tests for: cache hit (no redundant I/O), cache expiry and refresh, and cache sharing across multiple provider instances. - Added autouse fixture to clear the module-level cache between tests. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- litellm/_redis_credential_provider.py | 59 ++++++++++++++++++- tests/test_litellm/test_redis.py | 85 ++++++++++++++++++++++++--- 2 files changed, 132 insertions(+), 12 deletions(-) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 495d2a879bd..7161ce2dd13 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,8 +1,19 @@ import asyncio -from typing import Tuple +import threading +import time +from typing import Dict, Optional, Tuple from redis.credentials import CredentialProvider # type: ignore[attr-defined] +# GCP IAM tokens are valid for 1 hour. Cache for 55 minutes to refresh before expiry. +_GCP_IAM_TOKEN_TTL_SECONDS = 3300 + +# Module-level cache shared across all GCPIAMCredentialProvider instances for the +# same service account, so multiple Redis connections on the same pod share one token. +# Keyed by service_account → (token, expiry_monotonic_timestamp). +_token_cache: Dict[str, Tuple[str, float]] = {} +_token_cache_lock = threading.Lock() + def _generate_gcp_iam_access_token(service_account: str) -> str: """ @@ -31,23 +42,65 @@ def _generate_gcp_iam_access_token(service_account: str) -> str: return str(response.access_token) +def _get_cached_gcp_iam_token(service_account: str) -> str: + """ + Return a cached GCP IAM token, refreshing only when expired. + + Uses a module-level cache shared across all GCPIAMCredentialProvider + instances for the same service account. The threading.Lock ensures only + one thread performs the network round-trip on expiry; all others wait + briefly and read the fresh token (double-checked locking pattern). + + This avoids N concurrent blocking IAM refreshes when N Redis connections + are established simultaneously (e.g. during health checks or pool warm-up), + which would otherwise serialise inside Python's async event loop and cause + cascading request latency. + """ + cached = _token_cache.get(service_account) + if cached is not None: + token, expiry = cached + if time.monotonic() < expiry: + return token + + with _token_cache_lock: + # Re-check inside the lock: another thread may have refreshed already. + cached = _token_cache.get(service_account) + if cached is not None: + token, expiry = cached + if time.monotonic() < expiry: + return token + + token = _generate_gcp_iam_access_token(service_account) + _token_cache[service_account] = ( + token, + time.monotonic() + _GCP_IAM_TOKEN_TTL_SECONDS, + ) + return token + + class GCPIAMCredentialProvider(CredentialProvider): """ redis.credentials.CredentialProvider implementation that generates a fresh GCP IAM token on every new connection. This fixes the 1-hour token expiry issue for async Redis cluster clients, which previously generated the token once at startup and cached it as a static password. + + Tokens are cached at module level per service account for + _GCP_IAM_TOKEN_TTL_SECONDS (55 min) so that repeated connection + establishments — e.g. during connection pool warm-up or health checks — + do not each trigger a synchronous network round-trip that would block + Python's async event loop and cause cascading request latency. """ def __init__(self, gcp_service_account: str) -> None: self._gcp_service_account = gcp_service_account def get_credentials(self) -> Tuple[str]: - token = _generate_gcp_iam_access_token(self._gcp_service_account) + token = _get_cached_gcp_iam_token(self._gcp_service_account) return (token,) async def get_credentials_async(self) -> Tuple[str]: token = await asyncio.to_thread( - _generate_gcp_iam_access_token, self._gcp_service_account + _get_cached_gcp_iam_token, self._gcp_service_account ) return (token,) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 54503647ce8..2483469db23 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -13,7 +13,18 @@ from litellm._redis import ( get_redis_connection_pool, get_redis_url_from_environment, ) -from litellm._redis_credential_provider import GCPIAMCredentialProvider +from litellm._redis_credential_provider import ( + GCPIAMCredentialProvider, + _token_cache, +) + + +@pytest.fixture(autouse=True) +def clear_gcp_iam_token_cache(): + """Reset the module-level GCP IAM token cache between tests.""" + _token_cache.clear() + yield + _token_cache.clear() def test_get_redis_url_from_environment_single_url(monkeypatch): @@ -202,7 +213,7 @@ def test_get_redis_async_client_without_connection_pool(): def test_gcp_iam_credential_provider_get_credentials(): - """GCPIAMCredentialProvider.get_credentials() returns a fresh token tuple on every call.""" + """GCPIAMCredentialProvider.get_credentials() returns a token tuple.""" service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" with patch( @@ -216,20 +227,76 @@ def test_gcp_iam_credential_provider_get_credentials(): mock_gen.assert_called_once_with(service_account) -def test_gcp_iam_credential_provider_regenerates_token_on_each_call(): - """Each call to get_credentials() generates a new token (no caching).""" +def test_gcp_iam_credential_provider_caches_token(): + """ + Repeated calls to get_credentials() reuse the cached token and only call + _generate_gcp_iam_access_token once, avoiding redundant blocking I/O. + """ service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" - tokens = ["tok-1", "tok-2", "tok-3"] with patch( "litellm._redis_credential_provider._generate_gcp_iam_access_token", - side_effect=tokens, + return_value="tok-cached", ) as mock_gen: provider = GCPIAMCredentialProvider(service_account) - results = [provider.get_credentials() for _ in range(3)] + results = [provider.get_credentials() for _ in range(5)] - assert results == [("tok-1",), ("tok-2",), ("tok-3",)] - assert mock_gen.call_count == 3 + assert all(r == ("tok-cached",) for r in results) + # Token must be fetched exactly once regardless of how many connections are established + mock_gen.assert_called_once_with(service_account) + + +def test_gcp_iam_credential_provider_refreshes_on_expiry(): + """ + get_credentials() fetches a new token after the cached one expires, + ensuring connections always authenticate with a valid token. + """ + import time + + import litellm._redis_credential_provider as cred_module + + service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" + + with patch( + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + side_effect=["tok-1", "tok-2"], + ) as mock_gen: + provider = GCPIAMCredentialProvider(service_account) + + # First call — populates cache + assert provider.get_credentials() == ("tok-1",) + + # Artificially expire the cached token + cred_module._token_cache[service_account] = ("tok-1", time.monotonic() - 1) + + # Second call — cache miss, must refresh + assert provider.get_credentials() == ("tok-2",) + + assert mock_gen.call_count == 2 + + +def test_gcp_iam_credential_provider_cache_shared_across_instances(): + """ + Multiple GCPIAMCredentialProvider instances for the same service account + share one cached token so concurrent Redis connections don't each trigger + a blocking IAM round-trip. + """ + service_account = ( + "projects/-/serviceAccounts/shared@project.iam.gserviceaccount.com" + ) + + with patch( + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + return_value="tok-shared", + ) as mock_gen: + p1 = GCPIAMCredentialProvider(service_account) + p2 = GCPIAMCredentialProvider(service_account) + + assert p1.get_credentials() == ("tok-shared",) + assert p2.get_credentials() == ("tok-shared",) + + # Only one network call despite two provider instances + mock_gen.assert_called_once() def test_get_redis_async_client_gcp_cluster_uses_credential_provider():