1029 lines
37 KiB
Python
1029 lines
37 KiB
Python
"""
|
|
Async Redis Keyspace Notifications support for redis-py.
|
|
|
|
This module provides async utilities for subscribing to and parsing Redis
|
|
keyspace notifications.
|
|
|
|
Standalone Redis Example:
|
|
>>> from redis.asyncio import Redis
|
|
>>> from redis.asyncio.keyspace_notifications import (
|
|
... AsyncKeyspaceNotifications,
|
|
... )
|
|
>>> from redis.keyspace_notifications import KeyspaceChannel, EventType
|
|
>>>
|
|
>>> async def main():
|
|
... async with Redis() as r:
|
|
... async with AsyncKeyspaceNotifications(r) as ksn:
|
|
... channel = KeyspaceChannel("user:*")
|
|
... await ksn.subscribe(channel)
|
|
... async for notification in ksn.listen():
|
|
... print(f"Key: {notification.key}, Event: {notification.event_type}")
|
|
|
|
Redis Cluster Example:
|
|
>>> from redis.asyncio.cluster import RedisCluster
|
|
>>> from redis.asyncio.keyspace_notifications import (
|
|
... AsyncClusterKeyspaceNotifications,
|
|
... )
|
|
>>> from redis.keyspace_notifications import KeyspaceChannel, EventType
|
|
>>>
|
|
>>> async def main():
|
|
... async with RedisCluster(host="localhost", port=7000) as rc:
|
|
... async with AsyncClusterKeyspaceNotifications(rc) as ksn:
|
|
... channel = KeyspaceChannel("user:*")
|
|
... await ksn.subscribe(channel)
|
|
... async for notification in ksn.listen():
|
|
... print(f"Key: {notification.key}, Event: {notification.event_type}")
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import logging
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from redis.asyncio.client import PubSub, Redis
|
|
from redis.asyncio.cluster import ClusterNode, RedisCluster, _ClusterNodePoolAdapter
|
|
from redis.exceptions import (
|
|
ConnectionError,
|
|
RedisError,
|
|
TimeoutError,
|
|
)
|
|
from redis.keyspace_notifications import (
|
|
ChannelT,
|
|
KeyeventChannel,
|
|
KeyNotification,
|
|
KeyspaceChannel,
|
|
SubkeyeventChannel,
|
|
SubkeyspaceChannel,
|
|
SubkeyspaceeventChannel,
|
|
SubkeyspaceitemChannel,
|
|
_is_pattern,
|
|
)
|
|
from redis.utils import safe_str
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Type alias for handlers that can be sync or async
|
|
AsyncHandlerT = Callable[[KeyNotification], None | Awaitable[None]]
|
|
|
|
|
|
# =============================================================================
|
|
# Async Interface for Keyspace Notifications
|
|
# =============================================================================
|
|
|
|
|
|
class AsyncKeyspaceNotificationsInterface(ABC):
|
|
"""
|
|
Async interface for keyspace notification managers.
|
|
|
|
This interface provides a consistent async API for both standalone
|
|
(AsyncKeyspaceNotifications) and cluster (AsyncClusterKeyspaceNotifications)
|
|
implementations.
|
|
"""
|
|
|
|
@abstractmethod
|
|
async def subscribe(
|
|
self,
|
|
*channels: ChannelT,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to keyspace notification channels."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def unsubscribe(self, *channels: ChannelT):
|
|
"""Unsubscribe from keyspace notification channels."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_keyspace(
|
|
self,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to keyspace notifications for specific keys."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_keyevent(
|
|
self,
|
|
event: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to keyevent notifications for specific event types."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_subkeyspace(
|
|
self,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspace notifications for specific keys."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_subkeyevent(
|
|
self,
|
|
event: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyevent notifications for specific event types."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_subkeyspaceitem(
|
|
self,
|
|
key_or_pattern: str,
|
|
subkey_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspaceitem notifications for a specific subkey."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def subscribe_subkeyspaceevent(
|
|
self,
|
|
event: str,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspaceevent notifications for an event on a key."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def get_message(
|
|
self,
|
|
ignore_subscribe_messages: bool | None = None,
|
|
timeout: float = 0.0,
|
|
) -> KeyNotification | None:
|
|
"""Get the next keyspace notification if one is available."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def listen(self) -> AsyncIterator[KeyNotification]:
|
|
"""Listen for keyspace notifications."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def aclose(self):
|
|
"""Close the notification manager and clean up resources."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def __aenter__(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def __aexit__(self, _exc_type, _exc_val, _exc_tb):
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def subscribed(self) -> bool:
|
|
"""Check if there are any active subscriptions and not closed."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def run(
|
|
self,
|
|
poll_timeout: float = 1.0,
|
|
exception_handler: Callable[
|
|
[BaseException, AsyncKeyspaceNotificationsInterface],
|
|
None | Awaitable[None],
|
|
]
|
|
| None = None,
|
|
) -> None:
|
|
"""
|
|
Run the notification loop as a coroutine.
|
|
|
|
This is the async equivalent of run_in_thread() for sync notifications.
|
|
Use asyncio.create_task() to run in the background.
|
|
|
|
The exception_handler can be either a sync or async function.
|
|
"""
|
|
pass
|
|
|
|
|
|
# =============================================================================
|
|
# Abstract Base Class for Async Keyspace Notifications
|
|
# =============================================================================
|
|
|
|
|
|
class AbstractAsyncKeyspaceNotifications(AsyncKeyspaceNotificationsInterface):
|
|
"""
|
|
Abstract base class for async keyspace notification managers.
|
|
|
|
Provides shared implementation for subscribe/unsubscribe logic.
|
|
Subclasses must implement:
|
|
- _execute_subscribe: Execute the subscribe operation
|
|
- _execute_unsubscribe: Execute the unsubscribe operation
|
|
- get_message: Get the next notification
|
|
- listen: Async generator for notifications
|
|
- aclose: Clean up resources
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
key_prefix: str | bytes | None = None,
|
|
ignore_subscribe_messages: bool = True,
|
|
):
|
|
"""
|
|
Initialize the base async keyspace notification manager.
|
|
|
|
Args:
|
|
key_prefix: Optional prefix to filter and strip from keys in notifications
|
|
ignore_subscribe_messages: If True, subscribe/unsubscribe confirmations
|
|
are not returned by get_message/listen
|
|
"""
|
|
self.key_prefix = key_prefix
|
|
self.ignore_subscribe_messages = ignore_subscribe_messages
|
|
self._closed = False
|
|
|
|
async def subscribe(
|
|
self,
|
|
*channels: ChannelT,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""
|
|
Subscribe to keyspace notification channels.
|
|
|
|
Automatically detects whether each channel is a pattern (contains
|
|
wildcards like *, ?, [) or an exact channel name and uses the
|
|
appropriate Redis subscribe command internally.
|
|
|
|
The handler can be either a sync or async function. Note that a
|
|
**sync** handler will be called directly on the event loop thread,
|
|
so it must not perform blocking I/O or long-running computation —
|
|
prefer an ``async`` handler whenever possible.
|
|
"""
|
|
# Wrap the handler to convert raw messages to KeyNotification objects
|
|
wrapped_handler: Callable | None = None
|
|
if handler is not None:
|
|
key_prefix = self.key_prefix
|
|
is_async_handler = inspect.iscoroutinefunction(handler)
|
|
|
|
if is_async_handler:
|
|
# We've verified handler is async, so the result is awaitable
|
|
async_handler = handler
|
|
|
|
async def _async_wrap_handler(message):
|
|
notification = KeyNotification.from_message(
|
|
message, key_prefix=key_prefix
|
|
)
|
|
if notification is not None:
|
|
await async_handler(notification)
|
|
|
|
wrapped_handler = _async_wrap_handler
|
|
else:
|
|
|
|
def _sync_wrap_handler(message):
|
|
notification = KeyNotification.from_message(
|
|
message, key_prefix=key_prefix
|
|
)
|
|
if notification is not None:
|
|
handler(notification)
|
|
|
|
wrapped_handler = _sync_wrap_handler
|
|
|
|
patterns = {}
|
|
exact_channels = {}
|
|
|
|
for channel in channels:
|
|
if hasattr(channel, "_channel_str"):
|
|
channel_str = str(channel)
|
|
else:
|
|
channel_str = safe_str(channel)
|
|
if _is_pattern(channel):
|
|
patterns[channel_str] = wrapped_handler
|
|
else:
|
|
exact_channels[channel_str] = wrapped_handler
|
|
|
|
# Delegate to subclass implementation first. For standalone Redis
|
|
# this raises on failure, keeping tracking state clean. For cluster
|
|
# implementations the operation is best-effort (partial failures are
|
|
# logged, not raised) so tracking state is always updated afterwards.
|
|
await self._execute_subscribe(patterns, exact_channels)
|
|
self._track_subscribe(patterns, exact_channels)
|
|
|
|
@abstractmethod
|
|
async def _execute_subscribe(
|
|
self, patterns: dict[str, Any], exact_channels: dict[str, Any]
|
|
) -> None:
|
|
"""Execute the subscribe operation."""
|
|
pass
|
|
|
|
async def unsubscribe(self, *channels: ChannelT):
|
|
"""Unsubscribe from keyspace notification channels."""
|
|
patterns = []
|
|
exact_channels = []
|
|
|
|
for channel in channels:
|
|
if hasattr(channel, "_channel_str"):
|
|
channel_str = str(channel)
|
|
else:
|
|
channel_str = safe_str(channel)
|
|
if _is_pattern(channel):
|
|
patterns.append(channel_str)
|
|
else:
|
|
exact_channels.append(channel_str)
|
|
|
|
# Delegate to subclass implementation first. For standalone Redis
|
|
# this raises on failure, keeping tracking state intact. For cluster
|
|
# implementations the operation is best-effort (partial failures are
|
|
# logged, not raised) so tracking state is always removed afterwards
|
|
# — this is intentional: the user asked to unsubscribe, so
|
|
# refresh_subscriptions should not re-subscribe these channels.
|
|
await self._execute_unsubscribe(patterns, exact_channels)
|
|
self._untrack_subscribe(patterns, exact_channels)
|
|
|
|
@abstractmethod
|
|
async def _execute_unsubscribe(
|
|
self, patterns: list[str], exact_channels: list[str]
|
|
) -> None:
|
|
"""Execute the unsubscribe operation."""
|
|
pass
|
|
|
|
def _track_subscribe(
|
|
self, patterns: dict[str, Any], exact_channels: dict[str, Any]
|
|
) -> None:
|
|
"""Track newly subscribed patterns/channels.
|
|
|
|
Override in subclasses that need to maintain their own subscription
|
|
registry (e.g. cluster implementations that must re-subscribe
|
|
new/failed-over nodes). The default is a no-op because standalone
|
|
implementations delegate tracking to the underlying PubSub object.
|
|
"""
|
|
|
|
def _untrack_subscribe(
|
|
self, patterns: list[str], exact_channels: list[str]
|
|
) -> None:
|
|
"""Remove patterns/channels from the subscription registry.
|
|
|
|
Override in subclasses that maintain their own subscription registry.
|
|
The default is a no-op.
|
|
"""
|
|
|
|
async def subscribe_keyspace(
|
|
self,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to keyspace notifications for specific keys."""
|
|
channel = KeyspaceChannel(key_or_pattern, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def subscribe_keyevent(
|
|
self,
|
|
event: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to keyevent notifications for specific event types."""
|
|
channel = KeyeventChannel(event, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def subscribe_subkeyspace(
|
|
self,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspace notifications for specific keys."""
|
|
channel = SubkeyspaceChannel(key_or_pattern, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def subscribe_subkeyevent(
|
|
self,
|
|
event: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyevent notifications for specific event types."""
|
|
channel = SubkeyeventChannel(event, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def subscribe_subkeyspaceitem(
|
|
self,
|
|
key_or_pattern: str,
|
|
subkey_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspaceitem notifications for a specific subkey."""
|
|
channel = SubkeyspaceitemChannel(key_or_pattern, subkey_or_pattern, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def subscribe_subkeyspaceevent(
|
|
self,
|
|
event: str,
|
|
key_or_pattern: str,
|
|
db: int = 0,
|
|
handler: AsyncHandlerT | None = None,
|
|
):
|
|
"""Subscribe to subkeyspaceevent notifications for an event on a key."""
|
|
channel = SubkeyspaceeventChannel(event, key_or_pattern, db=db)
|
|
await self.subscribe(channel, handler=handler)
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, _exc_type, _exc_val, _exc_tb):
|
|
await self.aclose()
|
|
return False
|
|
|
|
async def run(
|
|
self,
|
|
poll_timeout: float = 1.0,
|
|
exception_handler: Callable[
|
|
[BaseException, AsyncKeyspaceNotificationsInterface],
|
|
None | Awaitable[None],
|
|
]
|
|
| None = None,
|
|
) -> None:
|
|
"""
|
|
Run the notification loop as a coroutine.
|
|
|
|
This continuously polls for notifications and triggers handlers.
|
|
Use asyncio.create_task() to run in the background.
|
|
|
|
Args:
|
|
poll_timeout: Timeout in seconds for each get_message call.
|
|
exception_handler: Optional callback for handling exceptions.
|
|
Can be sync or async.
|
|
"""
|
|
while self.subscribed:
|
|
try:
|
|
await self.get_message(timeout=poll_timeout)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except BaseException as e:
|
|
if exception_handler is not None:
|
|
result = exception_handler(e, self)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
else:
|
|
raise
|
|
|
|
|
|
# =============================================================================
|
|
# Standalone Async Keyspace Notification Manager
|
|
# =============================================================================
|
|
|
|
|
|
class AsyncKeyspaceNotifications(AbstractAsyncKeyspaceNotifications):
|
|
"""
|
|
Manages keyspace notification subscriptions for standalone async Redis.
|
|
|
|
For standalone Redis, keyspace notifications work with a single PubSub
|
|
connection. This class wraps that connection and provides:
|
|
- Automatic pattern vs exact channel detection
|
|
- KeyNotification parsing with optional key_prefix filtering
|
|
- Convenience methods for keyspace and keyevent subscriptions
|
|
- Context manager and run() coroutine support
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_client: Redis,
|
|
key_prefix: str | bytes | None = None,
|
|
ignore_subscribe_messages: bool = True,
|
|
):
|
|
"""
|
|
Initialize the standalone async keyspace notification manager.
|
|
|
|
Note: Keyspace notifications must be enabled on the Redis server via
|
|
the ``notify-keyspace-events`` configuration option.
|
|
|
|
Args:
|
|
redis_client: An async Redis client instance
|
|
key_prefix: Optional prefix to filter and strip from keys in notifications
|
|
ignore_subscribe_messages: If True, subscribe/unsubscribe confirmations
|
|
are not returned by get_message/listen
|
|
"""
|
|
super().__init__(key_prefix, ignore_subscribe_messages)
|
|
self.redis = redis_client
|
|
# Create PubSub with ignore_subscribe_messages=False so per-call arg works
|
|
self._pubsub: PubSub = redis_client.pubsub(ignore_subscribe_messages=False)
|
|
|
|
async def _execute_subscribe(
|
|
self, patterns: dict[str, Any], exact_channels: dict[str, Any]
|
|
) -> None:
|
|
"""Execute subscribe on the single pubsub connection."""
|
|
if patterns:
|
|
await self._pubsub.psubscribe(**patterns)
|
|
if exact_channels:
|
|
await self._pubsub.subscribe(**exact_channels)
|
|
|
|
async def _execute_unsubscribe(
|
|
self, patterns: list[str], exact_channels: list[str]
|
|
) -> None:
|
|
"""Execute unsubscribe on the single pubsub connection."""
|
|
if patterns:
|
|
await self._pubsub.punsubscribe(*patterns)
|
|
if exact_channels:
|
|
await self._pubsub.unsubscribe(*exact_channels)
|
|
|
|
async def get_message(
|
|
self,
|
|
ignore_subscribe_messages: bool | None = None,
|
|
timeout: float = 0.0,
|
|
) -> KeyNotification | None:
|
|
"""
|
|
Get the next keyspace notification if one is available.
|
|
|
|
Args:
|
|
ignore_subscribe_messages: If True, skip subscribe/unsubscribe messages.
|
|
Defaults to the value set in __init__ (True).
|
|
timeout: Time to wait for a message.
|
|
|
|
Returns:
|
|
A KeyNotification if a notification is available and no handler
|
|
was registered for the channel, None otherwise.
|
|
"""
|
|
if ignore_subscribe_messages is None:
|
|
ignore_subscribe_messages = self.ignore_subscribe_messages
|
|
|
|
if self._closed:
|
|
return None
|
|
|
|
message = await self._pubsub.get_message(
|
|
ignore_subscribe_messages=ignore_subscribe_messages,
|
|
timeout=timeout,
|
|
)
|
|
|
|
if message is not None:
|
|
return KeyNotification.from_message(message, key_prefix=self.key_prefix)
|
|
|
|
return None
|
|
|
|
async def listen(self) -> AsyncIterator[KeyNotification]:
|
|
"""
|
|
Listen for keyspace notifications.
|
|
|
|
This is an async generator that yields KeyNotification objects as they arrive.
|
|
|
|
Yields:
|
|
KeyNotification objects for each keyspace/keyevent notification.
|
|
|
|
Example:
|
|
>>> async for notification in ksn.listen():
|
|
... print(f"{notification.key}: {notification.event_type}")
|
|
"""
|
|
while self.subscribed:
|
|
notification = await self.get_message(timeout=1.0)
|
|
if notification is not None:
|
|
yield notification
|
|
|
|
@property
|
|
def subscribed(self) -> bool:
|
|
"""Check if there are any active subscriptions and not closed."""
|
|
return not self._closed and self._pubsub.subscribed
|
|
|
|
async def aclose(self):
|
|
"""Close the pubsub connection and clean up resources."""
|
|
self._closed = True
|
|
try:
|
|
await self._pubsub.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# =============================================================================
|
|
# Cluster-Aware Async Keyspace Notification Manager
|
|
# =============================================================================
|
|
|
|
|
|
class AsyncClusterKeyspaceNotifications(AbstractAsyncKeyspaceNotifications):
|
|
"""
|
|
Manages keyspace notification subscriptions across all nodes in an async Redis Cluster.
|
|
|
|
In Redis Cluster, keyspace notifications are NOT broadcast between nodes.
|
|
Each node only emits notifications for keys it owns. This class automatically
|
|
subscribes to all primary nodes in the cluster and handles topology changes.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_cluster: RedisCluster,
|
|
key_prefix: str | bytes | None = None,
|
|
ignore_subscribe_messages: bool = True,
|
|
):
|
|
"""
|
|
Initialize the async cluster keyspace notification manager.
|
|
|
|
Note: Keyspace notifications must be enabled on all Redis cluster nodes via
|
|
the ``notify-keyspace-events`` configuration option.
|
|
|
|
Args:
|
|
redis_cluster: An async RedisCluster instance
|
|
key_prefix: Optional prefix to filter and strip from keys in notifications
|
|
ignore_subscribe_messages: If True, subscribe/unsubscribe confirmations
|
|
are not returned by get_message/listen
|
|
"""
|
|
super().__init__(key_prefix, ignore_subscribe_messages)
|
|
self.cluster = redis_cluster
|
|
|
|
# Canonical subscription registry: pattern/channel -> wrapped handler.
|
|
# In cluster mode there are multiple PubSub objects (one per node), so
|
|
# this is the single source of truth used to (re-)subscribe new or
|
|
# failed-over nodes.
|
|
self._subscribed_patterns: dict[str, Any] = {}
|
|
self._subscribed_channels: dict[str, Any] = {}
|
|
|
|
# Track subscriptions per node
|
|
self._node_pubsubs: dict[str, PubSub] = {}
|
|
|
|
# Lock for topology refresh operations
|
|
self._refresh_lock = asyncio.Lock()
|
|
|
|
# Current pubsub index for round-robin polling
|
|
self._poll_index = 0
|
|
|
|
@property
|
|
def subscribed(self) -> bool:
|
|
"""Check if there are any active subscriptions and not closed."""
|
|
return not self._closed and bool(
|
|
self._subscribed_patterns or self._subscribed_channels
|
|
)
|
|
|
|
def _track_subscribe(
|
|
self, patterns: dict[str, Any], exact_channels: dict[str, Any]
|
|
) -> None:
|
|
"""Track newly subscribed patterns/channels in the cluster registry."""
|
|
if patterns:
|
|
self._subscribed_patterns.update(patterns)
|
|
if exact_channels:
|
|
self._subscribed_channels.update(exact_channels)
|
|
|
|
def _untrack_subscribe(
|
|
self, patterns: list[str], exact_channels: list[str]
|
|
) -> None:
|
|
"""Remove patterns/channels from the cluster registry."""
|
|
for p in patterns:
|
|
self._subscribed_patterns.pop(p, None)
|
|
for c in exact_channels:
|
|
self._subscribed_channels.pop(c, None)
|
|
|
|
def _get_all_primary_nodes(self) -> list[ClusterNode]:
|
|
"""Get all primary nodes in the cluster."""
|
|
return self.cluster.get_primaries()
|
|
|
|
async def _ensure_node_pubsub(self, node: ClusterNode) -> PubSub:
|
|
"""Get or create a PubSub instance for a node.
|
|
|
|
Uses a :class:`_ClusterNodePoolAdapter` to borrow a connection
|
|
from the node's existing pool. When the ``PubSub`` is closed
|
|
the connection is disconnected and returned to the node,
|
|
ensuring no subscribed socket is left in the free queue.
|
|
"""
|
|
if node.name not in self._node_pubsubs:
|
|
pool_adapter = _ClusterNodePoolAdapter(node)
|
|
pubsub = PubSub(
|
|
connection_pool=pool_adapter, # type: ignore[arg-type]
|
|
ignore_subscribe_messages=False,
|
|
)
|
|
self._node_pubsubs[node.name] = pubsub
|
|
return self._node_pubsubs[node.name]
|
|
|
|
async def _cleanup_node(self, node_name: str) -> None:
|
|
"""Remove and close a node's PubSub.
|
|
|
|
``PubSub.aclose()`` disconnects the connection and releases it
|
|
back to the underlying :class:`ClusterNode` via the adapter.
|
|
"""
|
|
pubsub = self._node_pubsubs.pop(node_name, None)
|
|
if pubsub:
|
|
try:
|
|
await pubsub.aclose()
|
|
except Exception:
|
|
pass
|
|
|
|
async def _execute_subscribe(
|
|
self, patterns: dict[str, Any], exact_channels: dict[str, Any]
|
|
) -> None:
|
|
"""Execute subscribe on all cluster nodes.
|
|
|
|
Patterns and exact channels are subscribed in a single pass over
|
|
nodes so that a mid-batch node failure cannot create a
|
|
partially-caught-up replacement. If a node fails during this
|
|
call it is removed from ``_node_pubsubs`` and will be fully
|
|
re-subscribed on the next ``refresh_subscriptions`` cycle.
|
|
|
|
If a newly discovered node is encountered (not yet in
|
|
``_node_pubsubs``), it is also subscribed to all *previously*
|
|
tracked patterns/channels so it doesn't miss notifications for
|
|
subscriptions that were established before this node joined.
|
|
"""
|
|
if not patterns and not exact_channels:
|
|
return
|
|
|
|
failed_nodes: list[str] = []
|
|
for node in self._get_all_primary_nodes():
|
|
is_new_node = node.name not in self._node_pubsubs
|
|
pubsub = await self._ensure_node_pubsub(node)
|
|
try:
|
|
# If this is a brand-new node, catch it up on existing
|
|
# subscriptions before adding the new channels.
|
|
if is_new_node:
|
|
if self._subscribed_patterns:
|
|
await pubsub.psubscribe(**self._subscribed_patterns)
|
|
if self._subscribed_channels:
|
|
await pubsub.subscribe(**self._subscribed_channels)
|
|
|
|
if patterns:
|
|
await pubsub.psubscribe(**patterns)
|
|
if exact_channels:
|
|
await pubsub.subscribe(**exact_channels)
|
|
except Exception:
|
|
# Remove the broken pubsub and its connection pool
|
|
# so refresh_subscriptions can re-create both later.
|
|
await self._cleanup_node(node.name)
|
|
failed_nodes.append(node.name)
|
|
|
|
if failed_nodes:
|
|
logger.warning(
|
|
"Failed to subscribe on cluster nodes: %s. "
|
|
"These nodes will be retried on the next refresh cycle.",
|
|
", ".join(failed_nodes),
|
|
)
|
|
|
|
async def _execute_unsubscribe(
|
|
self, patterns: list[str], exact_channels: list[str]
|
|
) -> None:
|
|
"""Execute unsubscribe on all cluster nodes."""
|
|
if patterns:
|
|
await self._unsubscribe_from_all_nodes(patterns, use_punsubscribe=True)
|
|
if exact_channels:
|
|
await self._unsubscribe_from_all_nodes(
|
|
exact_channels, use_punsubscribe=False
|
|
)
|
|
|
|
async def _unsubscribe_from_all_nodes(
|
|
self, channels: list[str], use_punsubscribe: bool
|
|
):
|
|
"""Unsubscribe from patterns/channels on all nodes.
|
|
|
|
Best-effort: tries every node so that a single broken connection
|
|
does not prevent the remaining nodes from being unsubscribed.
|
|
Broken pubsubs are cleaned up; the tracking state is still removed
|
|
by the caller, so ``refresh_subscriptions`` will *not* re-subscribe
|
|
these channels on replacement nodes.
|
|
"""
|
|
failed_nodes: list[str] = []
|
|
for node_name in list(self._node_pubsubs.keys()):
|
|
pubsub = self._node_pubsubs.get(node_name)
|
|
if pubsub is None:
|
|
continue
|
|
try:
|
|
if use_punsubscribe:
|
|
await pubsub.punsubscribe(*channels)
|
|
else:
|
|
await pubsub.unsubscribe(*channels)
|
|
except Exception:
|
|
await self._cleanup_node(node_name)
|
|
failed_nodes.append(node_name)
|
|
|
|
if failed_nodes:
|
|
logger.warning(
|
|
"Failed to unsubscribe on cluster nodes: %s. "
|
|
"These nodes will be re-created on the next refresh cycle.",
|
|
", ".join(failed_nodes),
|
|
)
|
|
|
|
async def get_message(
|
|
self,
|
|
ignore_subscribe_messages: bool | None = None,
|
|
timeout: float = 0.0,
|
|
) -> KeyNotification | None:
|
|
"""
|
|
Get the next keyspace notification if one is available.
|
|
|
|
This method polls all node pubsubs in round-robin fashion until
|
|
a message is received or the timeout expires.
|
|
If a connection error occurs, subscriptions are automatically refreshed.
|
|
|
|
Args:
|
|
ignore_subscribe_messages: If True, skip subscribe/unsubscribe messages.
|
|
Defaults to the value set in __init__ (True).
|
|
timeout: Total time to wait for a message (distributed across all nodes)
|
|
|
|
Returns:
|
|
A KeyNotification if a notification is available, None otherwise.
|
|
"""
|
|
if self._closed:
|
|
return None
|
|
|
|
total_nodes = len(self._node_pubsubs)
|
|
if total_nodes == 0:
|
|
# Sleep for the requested timeout so callers that loop
|
|
# (run(), listen) don't spin the CPU when all node
|
|
# connections have been cleaned up.
|
|
if timeout > 0:
|
|
await asyncio.sleep(timeout)
|
|
return None
|
|
|
|
if ignore_subscribe_messages is None:
|
|
ignore_subscribe_messages = self.ignore_subscribe_messages
|
|
|
|
# Handle timeout=0 as a single non-blocking poll over all pubsubs
|
|
if timeout == 0.0:
|
|
return await self._poll_all_nodes_once(ignore_subscribe_messages)
|
|
|
|
# Calculate per-node timeout for each poll
|
|
per_node_timeout = min(0.1, timeout / max(total_nodes, 1))
|
|
|
|
start_time = time.monotonic()
|
|
end_time = start_time + timeout
|
|
|
|
while True:
|
|
if time.monotonic() >= end_time:
|
|
return None
|
|
|
|
pubsubs = list(self._node_pubsubs.values())
|
|
if not pubsubs:
|
|
return None
|
|
|
|
# Round-robin polling
|
|
self._poll_index = self._poll_index % len(pubsubs)
|
|
pubsub = pubsubs[self._poll_index]
|
|
self._poll_index += 1
|
|
|
|
try:
|
|
message = await pubsub.get_message(
|
|
ignore_subscribe_messages=ignore_subscribe_messages,
|
|
timeout=per_node_timeout,
|
|
)
|
|
except (ConnectionError, TimeoutError, RedisError):
|
|
await self._refresh_subscriptions_on_error()
|
|
continue
|
|
|
|
if message is not None:
|
|
notification = KeyNotification.from_message(
|
|
message, key_prefix=self.key_prefix
|
|
)
|
|
if notification is not None:
|
|
return notification
|
|
|
|
async def _poll_all_nodes_once(
|
|
self, ignore_subscribe_messages: bool
|
|
) -> KeyNotification | None:
|
|
"""
|
|
Perform a single non-blocking poll over all node pubsubs.
|
|
|
|
Returns:
|
|
A KeyNotification if one is available, None otherwise.
|
|
"""
|
|
had_error = False
|
|
for pubsub in list(self._node_pubsubs.values()):
|
|
try:
|
|
message = await pubsub.get_message(
|
|
ignore_subscribe_messages=ignore_subscribe_messages,
|
|
timeout=0.0,
|
|
)
|
|
except (ConnectionError, TimeoutError, RedisError):
|
|
# Record the error but continue polling remaining healthy
|
|
# nodes so that already-buffered notifications are not lost.
|
|
had_error = True
|
|
continue
|
|
|
|
if message is not None:
|
|
notification = KeyNotification.from_message(
|
|
message, key_prefix=self.key_prefix
|
|
)
|
|
if notification is not None:
|
|
# Refresh before returning if any node had an error,
|
|
# so the next poll cycle has fresh state.
|
|
if had_error:
|
|
await self._refresh_subscriptions_on_error()
|
|
return notification
|
|
|
|
# Refresh after polling all nodes if any had errors
|
|
if had_error:
|
|
await self._refresh_subscriptions_on_error()
|
|
return None
|
|
|
|
async def listen(self) -> AsyncIterator[KeyNotification]:
|
|
"""
|
|
Listen for keyspace notifications from all cluster nodes.
|
|
|
|
This is an async generator that yields KeyNotification objects as they arrive.
|
|
|
|
Yields:
|
|
KeyNotification objects for each keyspace/keyevent notification.
|
|
|
|
Example:
|
|
>>> async for notification in ksn.listen():
|
|
... print(f"{notification.key}: {notification.event_type}")
|
|
"""
|
|
while self.subscribed:
|
|
notification = await self.get_message(timeout=1.0)
|
|
if notification is not None:
|
|
yield notification
|
|
|
|
async def _refresh_subscriptions_on_error(self):
|
|
"""
|
|
Refresh subscriptions after a connection error.
|
|
|
|
This is called automatically when a connection error occurs during
|
|
get_message(). It checks if nodes changed before refreshing.
|
|
"""
|
|
self._poll_index = 0 # Reset round-robin index
|
|
|
|
try:
|
|
await self.refresh_subscriptions()
|
|
except Exception:
|
|
logger.warning(
|
|
"Failed to refresh cluster subscriptions, will retry on next error",
|
|
exc_info=True,
|
|
)
|
|
|
|
def _is_pubsub_connected(self, pubsub) -> bool:
|
|
"""Check if a pubsub connection is still alive."""
|
|
try:
|
|
conn = pubsub.connection
|
|
if conn is None:
|
|
return False
|
|
# Async connections use is_connected property
|
|
return conn.is_connected
|
|
except Exception:
|
|
return False
|
|
|
|
async def refresh_subscriptions(self):
|
|
"""
|
|
Refresh subscriptions after a topology change.
|
|
|
|
This method is called automatically when topology changes are detected
|
|
or when connection errors occur. You can also call it manually if needed.
|
|
|
|
This method:
|
|
1. Discovers any new primary nodes and subscribes them
|
|
2. Removes pubsubs for nodes that are no longer primaries
|
|
3. Re-creates broken pubsub connections for existing nodes
|
|
"""
|
|
async with self._refresh_lock:
|
|
current_primaries = {
|
|
node.name: node for node in self._get_all_primary_nodes()
|
|
}
|
|
|
|
# Remove pubsubs for nodes that are no longer primaries
|
|
removed_nodes = set(self._node_pubsubs.keys()) - set(
|
|
current_primaries.keys()
|
|
)
|
|
for node_name in removed_nodes:
|
|
await self._cleanup_node(node_name)
|
|
|
|
# Detect broken connections for existing nodes and remove them
|
|
# so they get re-created below
|
|
existing_nodes = set(self._node_pubsubs.keys()) & set(
|
|
current_primaries.keys()
|
|
)
|
|
for node_name in existing_nodes:
|
|
pubsub = self._node_pubsubs.get(node_name)
|
|
if pubsub and not self._is_pubsub_connected(pubsub):
|
|
# Connection is broken, remove it so it gets re-created
|
|
await self._cleanup_node(node_name)
|
|
|
|
# Subscribe new nodes (and nodes with broken connections)
|
|
# to existing patterns/channels
|
|
new_nodes = set(current_primaries.keys()) - set(self._node_pubsubs.keys())
|
|
failed_nodes: list[str] = []
|
|
for node_name in new_nodes:
|
|
node = current_primaries[node_name]
|
|
pubsub = await self._ensure_node_pubsub(node)
|
|
|
|
try:
|
|
if self._subscribed_patterns:
|
|
await pubsub.psubscribe(**self._subscribed_patterns)
|
|
if self._subscribed_channels:
|
|
await pubsub.subscribe(**self._subscribed_channels)
|
|
except Exception:
|
|
# Subscription failed - remove from dict so retry is possible
|
|
await self._cleanup_node(node_name)
|
|
failed_nodes.append(node_name)
|
|
|
|
# Raise after attempting all nodes so we don't skip any
|
|
if failed_nodes:
|
|
raise ConnectionError(
|
|
f"Failed to subscribe to cluster nodes: {', '.join(failed_nodes)}"
|
|
)
|
|
|
|
async def aclose(self):
|
|
"""Close all pubsub connections and clean up resources."""
|
|
self._closed = True
|
|
for node_name in list(self._node_pubsubs.keys()):
|
|
await self._cleanup_node(node_name)
|
|
self._subscribed_patterns.clear()
|
|
self._subscribed_channels.clear()
|