fix(prefect-redis): make socket timeouts reachable and publisher/reconnect safe

src/integrations/prefect-redis/prefect_redis/client.py

@@ -8,11 +8,14 @@
 from redis.asyncio import Redis
 from typing_extensions import Self, TypeAlias
 
+from prefect.logging import get_logger
 from prefect.settings.base import (
     PrefectBaseSettings,
     build_settings_config,  # type: ignore[reportPrivateUsage]
 )
 
+logger = get_logger(__name__)
+
 _UNSET: Any = object()
 
 
@@ -57,17 +60,19 @@ class RedisMessagingSettings(PrefectBaseSettings):
         description="Whether to use SSL for the Redis connection",
     )
     socket_timeout: Optional[float] = Field(
-        default=None,
+        default=60.0,
         description=(
-            "Timeout in seconds for socket read operations. "
-            "None means no timeout (preserves pre-redis-py-8 behavior)."
+            "Timeout in seconds for socket read operations. None means no "
+            "timeout. Consumers with a longer PREFECT_REDIS_MESSAGING_CONSUMER_BLOCK "
+            "extend this automatically for their own client, so idle blocking "
+            "reads (e.g. XREADGROUP) don't time out spuriously."
         ),
     )
     socket_connect_timeout: Optional[float] = Field(
-        default=None,
+        default=10.0,
         description=(
             "Timeout in seconds for socket connect operations. "
-            "None means no timeout (preserves pre-redis-py-8 behavior)."
+            "None means no timeout."
         ),
     )
     protocol: int = Field(
@@ -174,10 +179,31 @@ def close_all_cached_connections() -> None:
         loop.run_until_complete(client.aclose())
 
 
-async def clear_cached_clients() -> None:
-    """Clear cached Redis clients and the messaging URL lookup."""
+async def clear_cached_clients(client: Union[Redis, None] = None) -> None:
+    """Close and clear cached Redis clients and the messaging URL lookup.
+
+    Without closing, connections stuck in-use from a hung operation are never
+    released and the pool eventually exhausts its connection cap.
+
+    Args:
+        client: If given, only this cached client is closed and evicted;
+            other cached clients (e.g. for other endpoints) are left running.
+            If omitted, every current-loop client is closed.
+    """
     _get_redis_messaging_url.cache_clear()
-    _client_cache.clear()
+
+    current_loop = _running_loop()
+    for key, cached_client in list(_client_cache.items()):
+        if client is not None and cached_client is not client:
+            continue
+        _, _, _, loop = key
+        if loop is not None and loop is not current_loop:
+            continue  # can't await a client bound to another loop
+        _client_cache.pop(key, None)
+        try:
+            await cached_client.aclose()
+        except Exception:
+            logger.exception("Error closing cached Redis client")
 
 
 @cached

src/integrations/prefect-redis/prefect_redis/messaging.py

@@ -5,7 +5,7 @@
 import socket
 import time
 import uuid
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 from datetime import timedelta
 from functools import partial
 from types import TracebackType
@@ -40,6 +40,7 @@
 from prefect.server.utilities.messaging import Publisher as _Publisher
 from prefect.settings.base import PrefectBaseSettings, build_settings_config
 from prefect_redis.client import (
+    RedisMessagingSettings,
     clear_cached_clients,
     cluster_key_prefix,
     get_async_redis_client,
@@ -271,14 +272,23 @@ def __init__(
     async def __aenter__(self) -> Self:
         self._client = get_async_redis_client()
         self._batch: list[RedisStreamsMessage] = []
+        self._claimed_count = 0
+        self._flush_lock = asyncio.Lock()
 
         if self.publish_every is not None:
             interval = self.publish_every.total_seconds()
 
             async def _publish_periodically() -> None:
                 while True:
                     await asyncio.sleep(interval)
-                    await asyncio.shield(self._publish_current_batch())
+                    try:
+                        await asyncio.shield(self._publish_current_batch())
+                    except asyncio.CancelledError:
+                        raise
+                    except Exception:
+                        logger.exception(
+                            "Error publishing batch for topic %s", self.topic
+                        )
 
             self._periodic_task = asyncio.create_task(_publish_periodically())
 
@@ -296,6 +306,8 @@ async def __aexit__(
         try:
             if self._periodic_task:
                 self._periodic_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await self._periodic_task
             await self._publish_current_batch()
         except Exception:
             if self.deduplicate_by:
@@ -311,31 +323,48 @@ async def publish_data(self, data: bytes, attributes: dict[str, Any]):
             await asyncio.shield(self._publish_current_batch())
 
     async def _publish_current_batch(self) -> None:
-        if not self._batch:
-            return
+        async with self._flush_lock:
+            if not self._batch:
+                return
 
-        if self.deduplicate_by:
-            to_publish = await self.cache.without_duplicates(
-                self.deduplicate_by, self._batch
-            )
-        else:
-            to_publish = list(self._batch)
+            if self.deduplicate_by:
+                claimed = self._batch[: self._claimed_count]
+                fresh = self._batch[self._claimed_count :]
+                to_publish = claimed + (
+                    await self.cache.without_duplicates(self.deduplicate_by, fresh)
+                    if fresh
+                    else []
+                )
+            else:
+                to_publish = list(self._batch)
 
-        self._batch.clear()
+            self._batch.clear()
+            self._claimed_count = 0
 
-        try:
-            for message in to_publish:
-                await self._client.xadd(
-                    self.stream,
-                    {
-                        "data": message.data,
-                        "attributes": orjson.dumps(message.attributes),
-                    },
-                )
-        except Exception:
-            if self.deduplicate_by:
-                await self.cache.forget_duplicates(self.deduplicate_by, to_publish)
-            raise
+            published = 0
+            try:
+                for message in to_publish:
+                    await self._client.xadd(
+                        self.stream,
+                        {
+                            "data": message.data,
+                            "attributes": orjson.dumps(message.attributes),
+                        },
+                    )
+                    published += 1
+            except Exception:
+                unsent = to_publish[published:]
+                self._batch[:0] = unsent
+                if self.deduplicate_by and unsent:
+                    try:
+                        await self.cache.forget_duplicates(self.deduplicate_by, unsent)
+                    except Exception:
+                        self._claimed_count = len(unsent)
+                        logger.exception(
+                            "Error clearing deduplication markers for topic %s",
+                            self.topic,
+                        )
+                raise
 
 
 class Consumer(_Consumer):
@@ -397,6 +426,16 @@ def __init__(
         self._read_batch_size: Optional[int] = read_batch_size
         self.use_consumer_group = use_consumer_group
 
+    def _get_redis_client(self) -> Redis:
+        """Blocking reads hold the socket for `self.block`, so a socket timeout
+        shorter than that fires before the read can return."""
+        socket_timeout = RedisMessagingSettings().socket_timeout
+        if socket_timeout is None or socket_timeout > self.block.total_seconds():
+            return get_async_redis_client()
+        return get_async_redis_client(
+            socket_timeout=self.block.total_seconds() + socket_timeout
+        )
+
     async def _ensure_stream_and_group(self, redis_client: Redis) -> None:
         """Ensure the stream and consumer group exist."""
         try:
@@ -513,10 +552,11 @@ async def run(self, handler: MessageHandler) -> None:
         attempt = 0
         base_delay = 1.0
         max_delay = 60.0
+        redis_client: Optional[Redis] = None
 
         while True:  # Outer loop for connection resilience
             try:
-                redis_client: Redis = get_async_redis_client()
+                redis_client = self._get_redis_client()
 
                 if not self.use_consumer_group:
                     await self._run_without_consumer_group(handler, redis_client)
@@ -603,8 +643,9 @@ async def run(self, handler: MessageHandler) -> None:
                     f"reconnecting in {delay:.1f}s (attempt {attempt + 1}): {e}"
                 )
 
-                # Clear cached clients to force fresh connections
-                await clear_cached_clients()
+                # Retire only the client that failed; other endpoints'
+                # cached clients are unaffected.
+                await clear_cached_clients(client=redis_client)
 
                 await asyncio.sleep(delay)
                 attempt += 1