import asyncio import functools import logging from typing import TYPE_CHECKING, Any, Awaitable, Callable from redis.asyncio.observability.recorder import ( record_connection_handoff, record_connection_relaxed_timeout, record_maint_notification_count, ) from redis.maint_notifications import ( MaintenanceNotification, MaintenanceState, MaintNotificationsConfig, NodeMovingNotification, OSSNodeMigratedNotification, OSSNodeMigratingNotification, _get_maintenance_notification_name, _get_maintenance_notification_type, _should_skip_connection_timeout_update, ) from redis.observability.attributes import get_pool_name if TYPE_CHECKING: from redis.asyncio.cluster import RedisCluster from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection logger = logging.getLogger(__name__) _ScheduledCallback = Callable[..., Awaitable[None]] def _log_task_exception(message: str, task: "asyncio.Task[Any]") -> None: """Task done-callback that surfaces a failed task's exception in the logs. Without it, an unhandled exception in a fire-and-forget task is only reported by asyncio as a noisy "Task exception was never retrieved" warning. Retrieving the exception here suppresses that warning and logs a meaningful error instead. Cancellation is expected and left unlogged. Bind ``message`` with ``functools.partial`` before passing to ``add_done_callback`` (which supplies ``task``). """ try: exc = task.exception() except asyncio.CancelledError: return if exc: logger.error(message, exc_info=exc) def add_debug_log_for_notification( connection: object, notification: str | MaintenanceNotification, ) -> None: if not logger.isEnabledFor(logging.DEBUG): return socket_address = None try: writer = getattr(connection, "_writer", None) socket_name = writer.get_extra_info("sockname") if writer else None socket_address = ( socket_name[1] if socket_name and len(socket_name) > 1 else None ) except (AttributeError, OSError, TypeError): pass resolved_ip = None try: get_resolved_ip = getattr(connection, "get_resolved_ip", None) if callable(get_resolved_ip): resolved_ip = get_resolved_ip() except (AttributeError, OSError): pass logger.debug( f"Handling maintenance notification: {notification}, " f"with connection: {connection}, connected to ip {resolved_ip}, " f"local socket port: {socket_address}", ) class AsyncMaintNotificationsPoolHandler: def __init__( self, pool: Any, config: MaintNotificationsConfig, ) -> None: self.pool = pool self.config = config self._processed_notifications: set[MaintenanceNotification] = set() self._scheduled_tasks: set[asyncio.Task[None]] = set() self._lock = asyncio.Lock() self.connection: Any | None = None def set_connection( self, connection: "AsyncMaintNotificationsAbstractConnection" ) -> None: self.connection = connection def get_handler_for_connection(self) -> "AsyncMaintNotificationsPoolHandler": # Copy all data that should be shared between connections, while each # connection gets its own handler instance and current connection state. copy = AsyncMaintNotificationsPoolHandler(self.pool, self.config) copy._processed_notifications = self._processed_notifications copy._scheduled_tasks = self._scheduled_tasks copy._lock = self._lock copy.connection = None return copy async def remove_expired_notifications(self) -> None: async with self._lock: for notification in tuple(self._processed_notifications): if notification.is_expired(): self._processed_notifications.remove(notification) async def handle_notification(self, notification: MaintenanceNotification) -> None: await self.remove_expired_notifications() if isinstance(notification, NodeMovingNotification): await self.handle_node_moving_notification(notification) else: logger.error(f"Unhandled notification type: {notification}") async def handle_node_moving_notification( self, notification: NodeMovingNotification ) -> None: if ( not self.config.proactive_reconnect and not self.config.is_relaxed_timeouts_enabled() ): return async with self._lock: if notification in self._processed_notifications: # nothing to do in the connection pool handling # the notification has already been handled or is expired # just return return if logger.isEnabledFor(logging.DEBUG): logger.debug( f"Handling node MOVING notification: {notification}, " f"with connection: {self.connection}, connected to ip " f"{self.connection.get_resolved_ip() if self.connection else None}" ) # Get the current connected address - if any # This is the address that is being moved # and we need to handle only connections # connected to the same address moving_address_src = ( self.connection.getpeername() if self.connection else None ) # The async pool owns the active/free connection collections and # asyncio.Lock is not reentrant, so the whole MOVING pool mutation # has to be one pool-owned atomic operation. The handler still owns # the notification policy and passes the already-decided inputs. await self.pool.apply_moving_notification( notification=notification, config=self.config, moving_address_src=moving_address_src, run_proactive_reconnect=( self.config.proactive_reconnect and notification.new_node_host is not None ), ) if self.config.proactive_reconnect and notification.new_node_host is None: self._schedule( notification.ttl / 2, self.run_proactive_reconnect, moving_address_src, ) self._schedule( notification.ttl, self.handle_node_moved_notification, notification, ) await record_connection_handoff( pool_name=get_pool_name(self.pool), ) self._processed_notifications.add(notification) async def run_proactive_reconnect( self, moving_address_src: str | None = None ) -> None: """ Run proactive reconnect for the pool. Active connections are marked for reconnect after they complete the current command. Inactive connections are disconnected and will be connected on next use. """ async with self._lock: # This delayed reconnect must be atomic for the same reason as the # initial MOVING mutation: a connection can move between active and # free lists while another task is acquiring or releasing it. await self.pool.run_proactive_reconnect( moving_address_src=moving_address_src, ) async def handle_node_moved_notification( self, notification: NodeMovingNotification ) -> None: """ Handle the cleanup after a node moving notification expires. """ notification_hash = hash(notification) async with self._lock: if logger.isEnabledFor(logging.DEBUG): logger.debug( f"Reverting temporary changes related to notification: {notification}, " f"with connection: {self.connection}, connected to ip " f"{self.connection.get_resolved_ip() if self.connection else None}" ) reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled() reset_host_address = self.config.proactive_reconnect # Cleanup has to reset future connection kwargs and existing # matching connections together under the pool lock. Splitting it # lets an acquire/release interleave and leaves stale MOVING state. await self.pool.cleanup_moving_notification( notification_hash=notification_hash, reset_relaxed_timeout=reset_relaxed_timeout, reset_host_address=reset_host_address, ) def _schedule( self, delay: float, callback: _ScheduledCallback, *args: Any, ) -> None: # Record the absolute deadline now so that any lag between create_task # and the task's first execution slice does not push the fire time out. deadline = asyncio.get_running_loop().time() + delay task = asyncio.create_task(self._run_after(deadline, callback, *args)) self._scheduled_tasks.add(task) task.add_done_callback(self._scheduled_tasks.discard) task.add_done_callback( functools.partial( _log_task_exception, "Error handling scheduled maintenance notification", ) ) async def _run_after( self, deadline: float, callback: _ScheduledCallback, *args: Any, ) -> None: remaining = deadline - asyncio.get_running_loop().time() if remaining > 0: await asyncio.sleep(remaining) await callback(*args) async def cancel_scheduled_tasks(self) -> None: if not self._scheduled_tasks: return tasks = tuple(self._scheduled_tasks) for task in tasks: task.cancel() await asyncio.gather(*tasks, return_exceptions=True) class AsyncMaintNotificationsConnectionHandler: def __init__( self, connection: "AsyncMaintNotificationsAbstractConnection", config: MaintNotificationsConfig, ) -> None: self.connection = connection self.config = config def _get_pool_name(self) -> str: """ Get the pool name from the connection's pool handler. Falls back to connection representation if pool is not available. """ pool_handler = getattr( self.connection, "_maint_notifications_pool_handler", None ) if pool_handler and getattr(pool_handler, "pool", None): return get_pool_name(pool_handler.pool) # Fallback for standalone connections without a pool return repr(self.connection) async def handle_notification(self, notification: MaintenanceNotification) -> None: # 1 for start, 0 for end notification type, None for unknown. notification_type = _get_maintenance_notification_type(notification) maint_notification = _get_maintenance_notification_name(notification) await record_maint_notification_count( server_address=self.connection.host, server_port=self.connection.port, network_peer_address=self.connection.host, network_peer_port=self.connection.port, maint_notification=maint_notification, ) if notification_type is None: logger.error(f"Unhandled notification type: {notification}") return if notification_type: await self.handle_maintenance_start_notification( MaintenanceState.MAINTENANCE, notification ) else: await self.handle_maintenance_completed_notification( notification=notification ) async def handle_maintenance_start_notification( self, maintenance_state: MaintenanceState, notification: MaintenanceNotification, ) -> None: add_debug_log_for_notification(self.connection, notification) if _should_skip_connection_timeout_update( self.connection.maintenance_state, self.config ): return self.connection.maintenance_state = maintenance_state self.connection.set_tmp_settings( tmp_relaxed_timeout=self.config.relaxed_timeout ) self.connection.update_current_socket_timeout(self.config.relaxed_timeout) if isinstance(notification, OSSNodeMigratingNotification): # add the notification id to the set of processed start maint notifications # this is used to skip the unrelaxing of the timeouts if we have received more than # one start notification before the the final end notification self.connection.add_maint_start_notification(notification.id) maint_notification = _get_maintenance_notification_name(notification) await record_connection_relaxed_timeout( connection_name=self._get_pool_name(), maint_notification=maint_notification, relaxed=True, ) async def handle_maintenance_completed_notification(self, **kwargs: Any) -> None: # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled if _should_skip_connection_timeout_update( self.connection.maintenance_state, self.config ): return notification = None if kwargs.get("notification"): notification = kwargs["notification"] add_debug_log_for_notification( self.connection, notification if notification else "MAINTENANCE_COMPLETED" ) self.connection.reset_tmp_settings(reset_relaxed_timeout=True) # Maintenance completed - reset the connection # timeouts by providing -1 as the relaxed timeout self.connection.update_current_socket_timeout(-1) self.connection.maintenance_state = MaintenanceState.NONE # reset the sets that keep track of received start maint # notifications and skipped end maint notifications self.connection.reset_received_notifications() if notification: maint_notification = _get_maintenance_notification_name(notification) await record_connection_relaxed_timeout( connection_name=self._get_pool_name(), maint_notification=maint_notification, relaxed=False, ) class AsyncOSSMaintNotificationsHandler: """ Cluster-wide handler for OSS (open-source) cluster maintenance notifications. Reacts to SMIGRATED (slot migration completed) push notifications and triggers topology re-initialization via nodes_manager.initialize(). Lock discipline: _lock is held across the whole pool mutation, including await initialize(), so the topology refresh and the subsequent connection marking/disconnect happen as one atomic operation — releasing the lock mid-mutation would let other tasks observe partial state. Re-entrancy is not a deadlock risk here even though asyncio.Lock is non-reentrant: initialize() may dispatch commands whose responses carry new push notifications, but handle_notification schedules the actual handling as a separate background task rather than calling into it inline, so the re-entrant arrival never tries to re-acquire the lock on this call stack. The cheap _in_progress/_processed dedup that gates that scheduling runs without the lock — those sets are only mutated from the single event loop. """ def __init__( self, cluster_client: "RedisCluster", config: MaintNotificationsConfig, ) -> None: self.cluster_client = cluster_client self.config = config self._processed_notifications: set[MaintenanceNotification] = set() self._in_progress: set[MaintenanceNotification] = set() self._lock = asyncio.Lock() self._background_tasks: set[asyncio.Task] = set() async def remove_expired_notifications(self) -> None: async with self._lock: for n in tuple(self._processed_notifications): if n.is_expired(): self._processed_notifications.remove(n) async def handle_notification(self, notification: MaintenanceNotification) -> None: # Synchronous pre-dedup BEFORE scheduling a task. # # The same SMIGRATED notification is delivered by the server on every # connection, so under load this callback fires many times for the same # notification. Without this guard we would schedule one background task # per arrival - a flood of tasks that each acquire the lock and dedup, # saturating the single event loop. Reserving the notification in # _in_progress here (and skipping if already in-progress/processed) caps # it at exactly ONE handling task per unique notification. # # These sets are only mutated from the single event loop, so this # check-and-add is race-free without holding the lock. if ( notification in self._in_progress or notification in self._processed_notifications ): return self._in_progress.add(notification) # Schedule as a background task so the parser's read path is not blocked. # This also breaks the inline call chain that would otherwise deadlock: # initialize() dispatches commands whose responses may carry more push # notifications; as a separate task it can run while this one awaits. # # If scheduling the task (or wiring its callbacks) fails, release the # in-progress reservation made above. Otherwise the notification would be # stuck in _in_progress forever - the dedup guard would skip it on every # future arrival, and _do_handle_notification's finally (which normally # clears it) never runs because the task was never started. try: task = asyncio.get_running_loop().create_task( self._do_handle_notification(notification) ) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) task.add_done_callback( functools.partial( _log_task_exception, "Error handling maintenance notification background task", ) ) except Exception: self._in_progress.discard(notification) raise async def _do_handle_notification( self, notification: MaintenanceNotification ) -> None: try: if isinstance(notification, OSSNodeMigratedNotification): await self.handle_oss_maintenance_completed_notification(notification) else: logger.error(f"Unhandled notification type: {notification}") finally: # Release the in-progress reservation. On success the notification is # also in _processed_notifications (so it won't be re-handled); on # failure it is not, allowing a later retry. self._in_progress.discard(notification) async def handle_oss_maintenance_completed_notification( self, notification: OSSNodeMigratedNotification ) -> None: await self.remove_expired_notifications() async with self._lock: # handle_notification already reserved this notification in # _in_progress and guaranteed uniqueness; the processed check here is # defensive (e.g. across handler copies sharing the same sets). if notification in self._processed_notifications: return if logger.isEnabledFor(logging.DEBUG): logger.debug(f"Handling SMIGRATED notification: {notification}") # Extract the information about the src and destination nodes that are # affected by the maintenance. nodes_to_slots_mapping structure: # { # "src_host:port": [ # {"dest_host:port": "slot_range"}, # ... # ], # ... # } additional_startup_nodes_info = [] affected_nodes = set() for ( src_address, dest_mappings, ) in notification.nodes_to_slots_mapping.items(): src_host, src_port = src_address.rsplit(":", 1) src_node = self.cluster_client.nodes_manager.get_node( host=src_host, port=int(src_port) ) if src_node is not None: affected_nodes.add(src_node) for dest_mapping in dest_mappings: for dest_address in dest_mapping.keys(): dest_host, dest_port = dest_address.rsplit(":", 1) additional_startup_nodes_info.append( (dest_host, int(dest_port)) ) # Updates the cluster slots cache with the new slots mapping # This will also update the nodes cache with the new nodes mapping await self.cluster_client.nodes_manager.initialize( additional_startup_nodes_info=additional_startup_nodes_info, ) all_nodes = set(affected_nodes) all_nodes = all_nodes.union( self.cluster_client.nodes_manager.nodes_cache.values() ) for current_node in all_nodes: handoff_recorded = False if current_node in affected_nodes: # mark for reconnect all in-use connections to the node — this # forces them to disconnect after completing their current command free_set = set(current_node._free) for conn in current_node._connections: if conn not in free_set: add_debug_log_for_notification( conn, "SMIGRATED - mark for reconnect" ) conn.mark_for_reconnect() await record_connection_handoff( pool_name=f"{current_node.host}:{current_node.port}" ) handoff_recorded = True else: if logger.isEnabledFor(logging.DEBUG): logger.debug( f"SMIGRATED: Node {current_node.name} not affected " f"by maintenance, skipping mark for reconnect" ) if ( current_node not in self.cluster_client.nodes_manager.nodes_cache.values() ): task = asyncio.get_running_loop().create_task( current_node.disconnect_free_connections() ) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) task.add_done_callback( functools.partial( _log_task_exception, "Error disconnecting free connections after " "maintenance notification", ) ) if not handoff_recorded: await record_connection_handoff( pool_name=f"{current_node.host}:{current_node.port}" ) self._processed_notifications.add(notification)