Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import asyncio
2import functools
3import logging
4from typing import TYPE_CHECKING, Any, Awaitable, Callable
6from redis.asyncio.observability.recorder import (
7 record_connection_handoff,
8 record_connection_relaxed_timeout,
9 record_maint_notification_count,
10)
11from redis.maint_notifications import (
12 MaintenanceNotification,
13 MaintenanceState,
14 MaintNotificationsConfig,
15 NodeMovingNotification,
16 OSSNodeMigratedNotification,
17 OSSNodeMigratingNotification,
18 _get_maintenance_notification_name,
19 _get_maintenance_notification_type,
20 _should_skip_connection_timeout_update,
21)
22from redis.observability.attributes import get_pool_name
24if TYPE_CHECKING:
25 from redis.asyncio.cluster import RedisCluster
26 from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection
28logger = logging.getLogger(__name__)
30_ScheduledCallback = Callable[..., Awaitable[None]]
33def _log_task_exception(message: str, task: "asyncio.Task[Any]") -> None:
34 """Task done-callback that surfaces a failed task's exception in the logs.
36 Without it, an unhandled exception in a fire-and-forget task is only
37 reported by asyncio as a noisy "Task exception was never retrieved"
38 warning. Retrieving the exception here suppresses that warning and logs a
39 meaningful error instead. Cancellation is expected and left unlogged.
41 Bind ``message`` with ``functools.partial`` before passing to
42 ``add_done_callback`` (which supplies ``task``).
43 """
44 try:
45 exc = task.exception()
46 except asyncio.CancelledError:
47 return
48 if exc:
49 logger.error(message, exc_info=exc)
52def add_debug_log_for_notification(
53 connection: object,
54 notification: str | MaintenanceNotification,
55) -> None:
56 if not logger.isEnabledFor(logging.DEBUG):
57 return
59 socket_address = None
60 try:
61 writer = getattr(connection, "_writer", None)
62 socket_name = writer.get_extra_info("sockname") if writer else None
63 socket_address = (
64 socket_name[1] if socket_name and len(socket_name) > 1 else None
65 )
66 except (AttributeError, OSError, TypeError):
67 pass
69 resolved_ip = None
70 try:
71 get_resolved_ip = getattr(connection, "get_resolved_ip", None)
72 if callable(get_resolved_ip):
73 resolved_ip = get_resolved_ip()
74 except (AttributeError, OSError):
75 pass
77 logger.debug(
78 f"Handling maintenance notification: {notification}, "
79 f"with connection: {connection}, connected to ip {resolved_ip}, "
80 f"local socket port: {socket_address}",
81 )
84class AsyncMaintNotificationsPoolHandler:
85 def __init__(
86 self,
87 pool: Any,
88 config: MaintNotificationsConfig,
89 ) -> None:
90 self.pool = pool
91 self.config = config
92 self._processed_notifications: set[MaintenanceNotification] = set()
93 self._scheduled_tasks: set[asyncio.Task[None]] = set()
94 self._lock = asyncio.Lock()
95 self.connection: Any | None = None
97 def set_connection(
98 self, connection: "AsyncMaintNotificationsAbstractConnection"
99 ) -> None:
100 self.connection = connection
102 def get_handler_for_connection(self) -> "AsyncMaintNotificationsPoolHandler":
103 # Copy all data that should be shared between connections, while each
104 # connection gets its own handler instance and current connection state.
105 copy = AsyncMaintNotificationsPoolHandler(self.pool, self.config)
106 copy._processed_notifications = self._processed_notifications
107 copy._scheduled_tasks = self._scheduled_tasks
108 copy._lock = self._lock
109 copy.connection = None
110 return copy
112 async def remove_expired_notifications(self) -> None:
113 async with self._lock:
114 for notification in tuple(self._processed_notifications):
115 if notification.is_expired():
116 self._processed_notifications.remove(notification)
118 async def handle_notification(self, notification: MaintenanceNotification) -> None:
119 await self.remove_expired_notifications()
121 if isinstance(notification, NodeMovingNotification):
122 await self.handle_node_moving_notification(notification)
123 else:
124 logger.error(f"Unhandled notification type: {notification}")
126 async def handle_node_moving_notification(
127 self, notification: NodeMovingNotification
128 ) -> None:
129 if (
130 not self.config.proactive_reconnect
131 and not self.config.is_relaxed_timeouts_enabled()
132 ):
133 return
135 async with self._lock:
136 if notification in self._processed_notifications:
137 # nothing to do in the connection pool handling
138 # the notification has already been handled or is expired
139 # just return
140 return
141 if logger.isEnabledFor(logging.DEBUG):
142 logger.debug(
143 f"Handling node MOVING notification: {notification}, "
144 f"with connection: {self.connection}, connected to ip "
145 f"{self.connection.get_resolved_ip() if self.connection else None}"
146 )
147 # Get the current connected address - if any
148 # This is the address that is being moved
149 # and we need to handle only connections
150 # connected to the same address
151 moving_address_src = (
152 self.connection.getpeername() if self.connection else None
153 )
155 # The async pool owns the active/free connection collections and
156 # asyncio.Lock is not reentrant, so the whole MOVING pool mutation
157 # has to be one pool-owned atomic operation. The handler still owns
158 # the notification policy and passes the already-decided inputs.
159 await self.pool.apply_moving_notification(
160 notification=notification,
161 config=self.config,
162 moving_address_src=moving_address_src,
163 run_proactive_reconnect=(
164 self.config.proactive_reconnect
165 and notification.new_node_host is not None
166 ),
167 )
169 if self.config.proactive_reconnect and notification.new_node_host is None:
170 self._schedule(
171 notification.ttl / 2,
172 self.run_proactive_reconnect,
173 moving_address_src,
174 )
176 self._schedule(
177 notification.ttl,
178 self.handle_node_moved_notification,
179 notification,
180 )
182 await record_connection_handoff(
183 pool_name=get_pool_name(self.pool),
184 )
186 self._processed_notifications.add(notification)
188 async def run_proactive_reconnect(
189 self, moving_address_src: str | None = None
190 ) -> None:
191 """
192 Run proactive reconnect for the pool.
193 Active connections are marked for reconnect after they complete the current command.
194 Inactive connections are disconnected and will be connected on next use.
195 """
196 async with self._lock:
197 # This delayed reconnect must be atomic for the same reason as the
198 # initial MOVING mutation: a connection can move between active and
199 # free lists while another task is acquiring or releasing it.
200 await self.pool.run_proactive_reconnect(
201 moving_address_src=moving_address_src,
202 )
204 async def handle_node_moved_notification(
205 self, notification: NodeMovingNotification
206 ) -> None:
207 """
208 Handle the cleanup after a node moving notification expires.
209 """
210 notification_hash = hash(notification)
212 async with self._lock:
213 if logger.isEnabledFor(logging.DEBUG):
214 logger.debug(
215 f"Reverting temporary changes related to notification: {notification}, "
216 f"with connection: {self.connection}, connected to ip "
217 f"{self.connection.get_resolved_ip() if self.connection else None}"
218 )
219 reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled()
220 reset_host_address = self.config.proactive_reconnect
222 # Cleanup has to reset future connection kwargs and existing
223 # matching connections together under the pool lock. Splitting it
224 # lets an acquire/release interleave and leaves stale MOVING state.
225 await self.pool.cleanup_moving_notification(
226 notification_hash=notification_hash,
227 reset_relaxed_timeout=reset_relaxed_timeout,
228 reset_host_address=reset_host_address,
229 )
231 def _schedule(
232 self,
233 delay: float,
234 callback: _ScheduledCallback,
235 *args: Any,
236 ) -> None:
237 # Record the absolute deadline now so that any lag between create_task
238 # and the task's first execution slice does not push the fire time out.
239 deadline = asyncio.get_running_loop().time() + delay
240 task = asyncio.create_task(self._run_after(deadline, callback, *args))
241 self._scheduled_tasks.add(task)
242 task.add_done_callback(self._scheduled_tasks.discard)
243 task.add_done_callback(
244 functools.partial(
245 _log_task_exception,
246 "Error handling scheduled maintenance notification",
247 )
248 )
250 async def _run_after(
251 self,
252 deadline: float,
253 callback: _ScheduledCallback,
254 *args: Any,
255 ) -> None:
256 remaining = deadline - asyncio.get_running_loop().time()
257 if remaining > 0:
258 await asyncio.sleep(remaining)
259 await callback(*args)
261 async def cancel_scheduled_tasks(self) -> None:
262 if not self._scheduled_tasks:
263 return
264 tasks = tuple(self._scheduled_tasks)
265 for task in tasks:
266 task.cancel()
267 await asyncio.gather(*tasks, return_exceptions=True)
270class AsyncMaintNotificationsConnectionHandler:
271 def __init__(
272 self,
273 connection: "AsyncMaintNotificationsAbstractConnection",
274 config: MaintNotificationsConfig,
275 ) -> None:
276 self.connection = connection
277 self.config = config
279 def _get_pool_name(self) -> str:
280 """
281 Get the pool name from the connection's pool handler.
282 Falls back to connection representation if pool is not available.
283 """
284 pool_handler = getattr(
285 self.connection, "_maint_notifications_pool_handler", None
286 )
287 if pool_handler and getattr(pool_handler, "pool", None):
288 return get_pool_name(pool_handler.pool)
289 # Fallback for standalone connections without a pool
290 return repr(self.connection)
292 async def handle_notification(self, notification: MaintenanceNotification) -> None:
293 # 1 for start, 0 for end notification type, None for unknown.
294 notification_type = _get_maintenance_notification_type(notification)
295 maint_notification = _get_maintenance_notification_name(notification)
297 await record_maint_notification_count(
298 server_address=self.connection.host,
299 server_port=self.connection.port,
300 network_peer_address=self.connection.host,
301 network_peer_port=self.connection.port,
302 maint_notification=maint_notification,
303 )
305 if notification_type is None:
306 logger.error(f"Unhandled notification type: {notification}")
307 return
309 if notification_type:
310 await self.handle_maintenance_start_notification(
311 MaintenanceState.MAINTENANCE, notification
312 )
313 else:
314 await self.handle_maintenance_completed_notification(
315 notification=notification
316 )
318 async def handle_maintenance_start_notification(
319 self,
320 maintenance_state: MaintenanceState,
321 notification: MaintenanceNotification,
322 ) -> None:
323 add_debug_log_for_notification(self.connection, notification)
325 if _should_skip_connection_timeout_update(
326 self.connection.maintenance_state, self.config
327 ):
328 return
330 self.connection.maintenance_state = maintenance_state
331 self.connection.set_tmp_settings(
332 tmp_relaxed_timeout=self.config.relaxed_timeout
333 )
334 self.connection.update_current_socket_timeout(self.config.relaxed_timeout)
335 if isinstance(notification, OSSNodeMigratingNotification):
336 # add the notification id to the set of processed start maint notifications
337 # this is used to skip the unrelaxing of the timeouts if we have received more than
338 # one start notification before the the final end notification
339 self.connection.add_maint_start_notification(notification.id)
341 maint_notification = _get_maintenance_notification_name(notification)
342 await record_connection_relaxed_timeout(
343 connection_name=self._get_pool_name(),
344 maint_notification=maint_notification,
345 relaxed=True,
346 )
348 async def handle_maintenance_completed_notification(self, **kwargs: Any) -> None:
349 # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled
350 if _should_skip_connection_timeout_update(
351 self.connection.maintenance_state, self.config
352 ):
353 return
355 notification = None
356 if kwargs.get("notification"):
357 notification = kwargs["notification"]
358 add_debug_log_for_notification(
359 self.connection, notification if notification else "MAINTENANCE_COMPLETED"
360 )
361 self.connection.reset_tmp_settings(reset_relaxed_timeout=True)
362 # Maintenance completed - reset the connection
363 # timeouts by providing -1 as the relaxed timeout
364 self.connection.update_current_socket_timeout(-1)
365 self.connection.maintenance_state = MaintenanceState.NONE
366 # reset the sets that keep track of received start maint
367 # notifications and skipped end maint notifications
368 self.connection.reset_received_notifications()
370 if notification:
371 maint_notification = _get_maintenance_notification_name(notification)
372 await record_connection_relaxed_timeout(
373 connection_name=self._get_pool_name(),
374 maint_notification=maint_notification,
375 relaxed=False,
376 )
379class AsyncOSSMaintNotificationsHandler:
380 """
381 Cluster-wide handler for OSS (open-source) cluster maintenance notifications.
383 Reacts to SMIGRATED (slot migration completed) push notifications and
384 triggers topology re-initialization via nodes_manager.initialize().
386 Lock discipline: _lock is held across the whole pool mutation, including
387 await initialize(), so the topology refresh and the subsequent connection
388 marking/disconnect happen as one atomic operation — releasing the lock
389 mid-mutation would let other tasks observe partial state.
391 Re-entrancy is not a deadlock risk here even though asyncio.Lock is
392 non-reentrant: initialize() may dispatch commands whose responses carry new
393 push notifications, but handle_notification schedules the actual handling as
394 a separate background task rather than calling into it inline, so the
395 re-entrant arrival never tries to re-acquire the lock on this call stack.
396 The cheap _in_progress/_processed dedup that gates that scheduling runs
397 without the lock — those sets are only mutated from the single event loop.
398 """
400 def __init__(
401 self,
402 cluster_client: "RedisCluster",
403 config: MaintNotificationsConfig,
404 ) -> None:
405 self.cluster_client = cluster_client
406 self.config = config
407 self._processed_notifications: set[MaintenanceNotification] = set()
408 self._in_progress: set[MaintenanceNotification] = set()
409 self._lock = asyncio.Lock()
410 self._background_tasks: set[asyncio.Task] = set()
412 async def remove_expired_notifications(self) -> None:
413 async with self._lock:
414 for n in tuple(self._processed_notifications):
415 if n.is_expired():
416 self._processed_notifications.remove(n)
418 async def handle_notification(self, notification: MaintenanceNotification) -> None:
419 # Synchronous pre-dedup BEFORE scheduling a task.
420 #
421 # The same SMIGRATED notification is delivered by the server on every
422 # connection, so under load this callback fires many times for the same
423 # notification. Without this guard we would schedule one background task
424 # per arrival - a flood of tasks that each acquire the lock and dedup,
425 # saturating the single event loop. Reserving the notification in
426 # _in_progress here (and skipping if already in-progress/processed) caps
427 # it at exactly ONE handling task per unique notification.
428 #
429 # These sets are only mutated from the single event loop, so this
430 # check-and-add is race-free without holding the lock.
431 if (
432 notification in self._in_progress
433 or notification in self._processed_notifications
434 ):
435 return
436 self._in_progress.add(notification)
438 # Schedule as a background task so the parser's read path is not blocked.
439 # This also breaks the inline call chain that would otherwise deadlock:
440 # initialize() dispatches commands whose responses may carry more push
441 # notifications; as a separate task it can run while this one awaits.
442 #
443 # If scheduling the task (or wiring its callbacks) fails, release the
444 # in-progress reservation made above. Otherwise the notification would be
445 # stuck in _in_progress forever - the dedup guard would skip it on every
446 # future arrival, and _do_handle_notification's finally (which normally
447 # clears it) never runs because the task was never started.
448 try:
449 task = asyncio.get_running_loop().create_task(
450 self._do_handle_notification(notification)
451 )
452 self._background_tasks.add(task)
453 task.add_done_callback(self._background_tasks.discard)
454 task.add_done_callback(
455 functools.partial(
456 _log_task_exception,
457 "Error handling maintenance notification background task",
458 )
459 )
460 except Exception:
461 self._in_progress.discard(notification)
462 raise
464 async def _do_handle_notification(
465 self, notification: MaintenanceNotification
466 ) -> None:
467 try:
468 if isinstance(notification, OSSNodeMigratedNotification):
469 await self.handle_oss_maintenance_completed_notification(notification)
470 else:
471 logger.error(f"Unhandled notification type: {notification}")
472 finally:
473 # Release the in-progress reservation. On success the notification is
474 # also in _processed_notifications (so it won't be re-handled); on
475 # failure it is not, allowing a later retry.
476 self._in_progress.discard(notification)
478 async def handle_oss_maintenance_completed_notification(
479 self, notification: OSSNodeMigratedNotification
480 ) -> None:
481 await self.remove_expired_notifications()
483 async with self._lock:
484 # handle_notification already reserved this notification in
485 # _in_progress and guaranteed uniqueness; the processed check here is
486 # defensive (e.g. across handler copies sharing the same sets).
487 if notification in self._processed_notifications:
488 return
489 if logger.isEnabledFor(logging.DEBUG):
490 logger.debug(f"Handling SMIGRATED notification: {notification}")
492 # Extract the information about the src and destination nodes that are
493 # affected by the maintenance. nodes_to_slots_mapping structure:
494 # {
495 # "src_host:port": [
496 # {"dest_host:port": "slot_range"},
497 # ...
498 # ],
499 # ...
500 # }
501 additional_startup_nodes_info = []
502 affected_nodes = set()
503 for (
504 src_address,
505 dest_mappings,
506 ) in notification.nodes_to_slots_mapping.items():
507 src_host, src_port = src_address.rsplit(":", 1)
508 src_node = self.cluster_client.nodes_manager.get_node(
509 host=src_host, port=int(src_port)
510 )
511 if src_node is not None:
512 affected_nodes.add(src_node)
513 for dest_mapping in dest_mappings:
514 for dest_address in dest_mapping.keys():
515 dest_host, dest_port = dest_address.rsplit(":", 1)
516 additional_startup_nodes_info.append(
517 (dest_host, int(dest_port))
518 )
519 # Updates the cluster slots cache with the new slots mapping
520 # This will also update the nodes cache with the new nodes mapping
521 await self.cluster_client.nodes_manager.initialize(
522 additional_startup_nodes_info=additional_startup_nodes_info,
523 )
525 all_nodes = set(affected_nodes)
526 all_nodes = all_nodes.union(
527 self.cluster_client.nodes_manager.nodes_cache.values()
528 )
529 for current_node in all_nodes:
530 handoff_recorded = False
531 if current_node in affected_nodes:
532 # mark for reconnect all in-use connections to the node — this
533 # forces them to disconnect after completing their current command
534 free_set = set(current_node._free)
535 for conn in current_node._connections:
536 if conn not in free_set:
537 add_debug_log_for_notification(
538 conn, "SMIGRATED - mark for reconnect"
539 )
540 conn.mark_for_reconnect()
541 await record_connection_handoff(
542 pool_name=f"{current_node.host}:{current_node.port}"
543 )
544 handoff_recorded = True
545 else:
546 if logger.isEnabledFor(logging.DEBUG):
547 logger.debug(
548 f"SMIGRATED: Node {current_node.name} not affected "
549 f"by maintenance, skipping mark for reconnect"
550 )
551 if (
552 current_node
553 not in self.cluster_client.nodes_manager.nodes_cache.values()
554 ):
555 task = asyncio.get_running_loop().create_task(
556 current_node.disconnect_free_connections()
557 )
558 self._background_tasks.add(task)
559 task.add_done_callback(self._background_tasks.discard)
560 task.add_done_callback(
561 functools.partial(
562 _log_task_exception,
563 "Error disconnecting free connections after "
564 "maintenance notification",
565 )
566 )
567 if not handoff_recorded:
568 await record_connection_handoff(
569 pool_name=f"{current_node.host}:{current_node.port}"
570 )
572 self._processed_notifications.add(notification)