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.event import (
12 MaintenanceCompletedEvent,
13 MaintenanceEvent,
14 MaintenanceStartedEvent,
15)
16from redis.maint_notifications import (
17 MaintenanceNotification,
18 MaintenanceState,
19 MaintNotificationsConfig,
20 NodeMovingNotification,
21 OSSNodeMigratedNotification,
22 OSSNodeMigratingNotification,
23 _dispatch_maintenance_event,
24 _get_maintenance_notification_name,
25 _get_maintenance_notification_type,
26 _should_skip_connection_timeout_update,
27)
28from redis.observability.attributes import get_pool_name
30if TYPE_CHECKING:
31 from redis.asyncio.cluster import RedisCluster
32 from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection
34logger = logging.getLogger(__name__)
36_ScheduledCallback = Callable[..., Awaitable[None]]
39def _log_task_exception(message: str, task: "asyncio.Task[Any]") -> None:
40 """Task done-callback that surfaces a failed task's exception in the logs.
42 Without it, an unhandled exception in a fire-and-forget task is only
43 reported by asyncio as a noisy "Task exception was never retrieved"
44 warning. Retrieving the exception here suppresses that warning and logs a
45 meaningful error instead. Cancellation is expected and left unlogged.
47 Bind ``message`` with ``functools.partial`` before passing to
48 ``add_done_callback`` (which supplies ``task``).
49 """
50 try:
51 exc = task.exception()
52 except asyncio.CancelledError:
53 return
54 if exc:
55 logger.error(message, exc_info=exc)
58def add_debug_log_for_notification(
59 connection: object,
60 notification: str | MaintenanceNotification,
61) -> None:
62 if not logger.isEnabledFor(logging.DEBUG):
63 return
65 details = "no connection"
66 try:
67 extract = getattr(connection, "extract_connection_details", None)
68 if callable(extract):
69 details = extract()
70 except (AttributeError, OSError, TypeError):
71 pass
73 logger.debug(
74 f"Handling maintenance notification: {notification}, "
75 f"with connection: {connection}, {details}",
76 )
79class AsyncMaintNotificationsPoolHandler:
80 def __init__(
81 self,
82 pool: Any,
83 config: MaintNotificationsConfig,
84 ) -> None:
85 self.pool = pool
86 self.config = config
87 self._processed_notifications: set[MaintenanceNotification] = set()
88 self._scheduled_tasks: set[asyncio.Task[None]] = set()
89 # The handoffs whose relaxation the scheduled cleanup is still to end
90 self._pending_handoffs: set[NodeMovingNotification] = set()
91 self._lock = asyncio.Lock()
92 self.connection: Any | None = None
94 def set_connection(
95 self, connection: "AsyncMaintNotificationsAbstractConnection"
96 ) -> None:
97 self.connection = connection
99 def get_handler_for_connection(self) -> "AsyncMaintNotificationsPoolHandler":
100 # Copy all data that should be shared between connections, while each
101 # connection gets its own handler instance and current connection state.
102 copy = AsyncMaintNotificationsPoolHandler(self.pool, self.config)
103 copy._processed_notifications = self._processed_notifications
104 copy._scheduled_tasks = self._scheduled_tasks
105 copy._pending_handoffs = self._pending_handoffs
106 copy._lock = self._lock
107 copy.connection = None
108 return copy
110 async def remove_expired_notifications(self) -> None:
111 async with self._lock:
112 for notification in tuple(self._processed_notifications):
113 if notification.is_expired():
114 self._processed_notifications.remove(notification)
116 async def handle_notification(self, notification: MaintenanceNotification) -> None:
117 await self.remove_expired_notifications()
119 if isinstance(notification, NodeMovingNotification):
120 await self.handle_node_moving_notification(notification)
121 else:
122 logger.error(f"Unhandled notification type: {notification}")
124 async def handle_node_moving_notification(
125 self, notification: NodeMovingNotification
126 ) -> None:
127 if (
128 not self.config.proactive_reconnect
129 and not self.config.is_relaxed_timeouts_enabled()
130 ):
131 return
133 async with self._lock:
134 if notification in self._processed_notifications:
135 # nothing to do in the connection pool handling
136 # the notification has already been handled or is expired
137 # just return
138 return
139 if logger.isEnabledFor(logging.DEBUG):
140 logger.debug(
141 f"Handling node MOVING notification: {notification}, "
142 f"with connection: {self.connection}, connected to ip "
143 f"{self.connection.get_resolved_ip() if self.connection else None}"
144 )
145 # Get the current connected address - if any
146 # This is the address that is being moved
147 # and we need to handle only connections
148 # connected to the same address
149 moving_address_src = (
150 self.connection.getpeername() if self.connection else None
151 )
153 # Reported before the pool is relaxed: taking over a connection
154 # under maintenance completes that connection's pair, and a
155 # listener must see the handoff's pair open by then rather than
156 # no relaxation at all.
157 self._dispatch_handoff_event(MaintenanceStartedEvent, notification)
159 try:
160 # The async pool owns the active/free connection collections
161 # and asyncio.Lock is not reentrant, so the whole MOVING pool
162 # mutation has to be one pool-owned atomic operation. The
163 # handler still owns the notification policy and passes the
164 # already-decided inputs.
165 await self.pool.apply_moving_notification(
166 notification=notification,
167 config=self.config,
168 moving_address_src=moving_address_src,
169 run_proactive_reconnect=(
170 self.config.proactive_reconnect
171 and notification.new_node_host is not None
172 ),
173 )
175 if (
176 self.config.proactive_reconnect
177 and notification.new_node_host is None
178 ):
179 self._schedule(
180 notification.ttl / 2,
181 self.run_proactive_reconnect,
182 moving_address_src,
183 )
185 self._schedule(
186 notification.ttl,
187 self.handle_node_moved_notification,
188 notification,
189 )
190 if self.config.is_relaxed_timeouts_enabled():
191 self._pending_handoffs.add(notification)
192 except BaseException:
193 # The relaxation reported above was not applied, or has no
194 # cleanup to end it: close its pair
195 self._dispatch_handoff_event(MaintenanceCompletedEvent, notification)
196 raise
198 await record_connection_handoff(
199 pool_name=get_pool_name(self.pool),
200 )
202 self._processed_notifications.add(notification)
204 def _dispatch_handoff_event(
205 self,
206 event_type: type[MaintenanceEvent],
207 notification: NodeMovingNotification,
208 ) -> None:
209 """
210 Report a handoff's relaxation. The events report relaxed timeouts; a
211 handoff handled only for the proactive reconnect relaxes nothing.
212 """
213 if self.config.is_relaxed_timeouts_enabled():
214 _dispatch_maintenance_event(
215 self.config,
216 event_type(
217 connection_pool=self.pool,
218 connection=self.connection,
219 state=MaintenanceState.MOVING,
220 notification=notification,
221 config=self.config,
222 ),
223 )
225 async def run_proactive_reconnect(
226 self, moving_address_src: str | None = None
227 ) -> None:
228 """
229 Run proactive reconnect for the pool.
230 Active connections are marked for reconnect after they complete the current command.
231 Inactive connections are disconnected and will be connected on next use.
232 """
233 async with self._lock:
234 # This delayed reconnect must be atomic for the same reason as the
235 # initial MOVING mutation: a connection can move between active and
236 # free lists while another task is acquiring or releasing it.
237 await self.pool.run_proactive_reconnect(
238 moving_address_src=moving_address_src,
239 )
241 async def handle_node_moved_notification(
242 self, notification: NodeMovingNotification
243 ) -> None:
244 """
245 Handle the cleanup after a node moving notification expires.
246 """
247 notification_hash = hash(notification)
249 async with self._lock:
250 self._pending_handoffs.discard(notification)
251 if logger.isEnabledFor(logging.DEBUG):
252 logger.debug(
253 f"Reverting temporary changes related to notification: {notification}, "
254 f"with connection: {self.connection}, connected to ip "
255 f"{self.connection.get_resolved_ip() if self.connection else None}"
256 )
257 reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled()
258 reset_host_address = self.config.proactive_reconnect
260 # Cleanup has to reset future connection kwargs and existing
261 # matching connections together under the pool lock. Splitting it
262 # lets an acquire/release interleave and leaves stale MOVING state.
263 reverted = await self.pool.cleanup_moving_notification(
264 notification_hash=notification_hash,
265 reset_relaxed_timeout=reset_relaxed_timeout,
266 reset_host_address=reset_host_address,
267 )
269 # A newer MOVING notification has superseded this one when the
270 # kwargs were not reverted: the pool stays relaxed until it expires.
271 if reverted:
272 self._dispatch_handoff_event(MaintenanceCompletedEvent, notification)
274 def _schedule(
275 self,
276 delay: float,
277 callback: _ScheduledCallback,
278 *args: Any,
279 ) -> None:
280 # Record the absolute deadline now so that any lag between create_task
281 # and the task's first execution slice does not push the fire time out.
282 deadline = asyncio.get_running_loop().time() + delay
283 task = asyncio.create_task(self._run_after(deadline, callback, *args))
284 self._scheduled_tasks.add(task)
285 task.add_done_callback(self._scheduled_tasks.discard)
286 task.add_done_callback(
287 functools.partial(
288 _log_task_exception,
289 "Error handling scheduled maintenance notification",
290 )
291 )
293 async def _run_after(
294 self,
295 deadline: float,
296 callback: _ScheduledCallback,
297 *args: Any,
298 ) -> None:
299 remaining = deadline - asyncio.get_running_loop().time()
300 if remaining > 0:
301 await asyncio.sleep(remaining)
302 await callback(*args)
304 async def cancel_scheduled_tasks(self) -> None:
305 if not self._scheduled_tasks:
306 return
307 # The pool is closing, which ends the relaxation the cancelled cleanups
308 # would have ended: report it so their start events are paired
309 pending_handoffs = tuple(self._pending_handoffs)
310 self._pending_handoffs.clear()
311 tasks = tuple(self._scheduled_tasks)
312 for task in tasks:
313 task.cancel()
314 await asyncio.gather(*tasks, return_exceptions=True)
315 for notification in pending_handoffs:
316 self._dispatch_handoff_event(MaintenanceCompletedEvent, notification)
319class AsyncMaintNotificationsConnectionHandler:
320 def __init__(
321 self,
322 connection: "AsyncMaintNotificationsAbstractConnection",
323 config: MaintNotificationsConfig,
324 ) -> None:
325 self.connection = connection
326 self.config = config
328 def _get_pool(self) -> Any | None:
329 """
330 Get the pool from the connection's pool handler, or None for
331 standalone connections without a pool.
332 """
333 pool_handler = getattr(
334 self.connection, "_maint_notifications_pool_handler", None
335 )
336 if pool_handler:
337 return getattr(pool_handler, "pool", None)
338 return None
340 def _get_pool_name(self) -> str:
341 """
342 Get the pool name from the connection's pool handler.
343 Falls back to connection representation if pool is not available.
344 """
345 pool = self._get_pool()
346 if pool:
347 return get_pool_name(pool)
348 # Fallback for standalone connections without a pool
349 return repr(self.connection)
351 async def handle_notification(self, notification: MaintenanceNotification) -> None:
352 # 1 for start, 0 for end notification type, None for unknown.
353 notification_type = _get_maintenance_notification_type(notification)
354 maint_notification = _get_maintenance_notification_name(notification)
356 await record_maint_notification_count(
357 server_address=self.connection.host,
358 server_port=self.connection.port,
359 network_peer_address=self.connection.host,
360 network_peer_port=self.connection.port,
361 maint_notification=maint_notification,
362 )
364 if notification_type is None:
365 logger.error(f"Unhandled notification type: {notification}")
366 return
368 if notification_type:
369 await self.handle_maintenance_start_notification(
370 MaintenanceState.MAINTENANCE, notification
371 )
372 else:
373 await self.handle_maintenance_completed_notification(
374 notification=notification
375 )
377 async def handle_maintenance_start_notification(
378 self,
379 maintenance_state: MaintenanceState,
380 notification: MaintenanceNotification,
381 ) -> None:
382 add_debug_log_for_notification(self.connection, notification)
384 if _should_skip_connection_timeout_update(
385 self.connection.maintenance_state, self.config
386 ):
387 return
389 self.connection.maintenance_state = maintenance_state
390 self.connection.set_tmp_settings(
391 tmp_relaxed_timeout=self.config.relaxed_timeout
392 )
393 self.connection.update_current_socket_timeout(self.config.relaxed_timeout)
394 if isinstance(notification, OSSNodeMigratingNotification):
395 # add the notification id to the set of processed start maint notifications
396 # this is used to skip the unrelaxing of the timeouts if we have received more than
397 # one start notification before the final end notification
398 self.connection.add_maint_start_notification(notification.id)
400 maint_notification = _get_maintenance_notification_name(notification)
401 await record_connection_relaxed_timeout(
402 connection_name=self._get_pool_name(),
403 maint_notification=maint_notification,
404 relaxed=True,
405 )
407 _dispatch_maintenance_event(
408 self.config,
409 MaintenanceStartedEvent(
410 connection_pool=self._get_pool(),
411 connection=self.connection,
412 state=maintenance_state,
413 notification=notification,
414 config=self.config,
415 ),
416 )
418 async def handle_maintenance_completed_notification(self, **kwargs: Any) -> None:
419 # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled
420 if _should_skip_connection_timeout_update(
421 self.connection.maintenance_state, self.config
422 ):
423 return
425 # A completion can reach a connection whose start went to another one,
426 # or none at all; only a connection that was relaxed has a start event
427 # to pair the completion with
428 was_relaxed = self.connection.maintenance_state == MaintenanceState.MAINTENANCE
430 notification = None
431 if kwargs.get("notification"):
432 notification = kwargs["notification"]
433 add_debug_log_for_notification(
434 self.connection, notification if notification else "MAINTENANCE_COMPLETED"
435 )
436 self.connection.reset_tmp_settings(reset_relaxed_timeout=True)
437 # Maintenance completed - reset the connection
438 # timeouts by providing -1 as the relaxed timeout
439 self.connection.update_current_socket_timeout(-1)
440 self.connection.maintenance_state = MaintenanceState.NONE
441 # reset the sets that keep track of received start maint
442 # notifications and skipped end maint notifications
443 self.connection.reset_received_notifications()
445 if notification:
446 maint_notification = _get_maintenance_notification_name(notification)
447 await record_connection_relaxed_timeout(
448 connection_name=self._get_pool_name(),
449 maint_notification=maint_notification,
450 relaxed=False,
451 )
453 if was_relaxed:
454 _dispatch_maintenance_event(
455 self.config,
456 MaintenanceCompletedEvent(
457 connection_pool=self._get_pool(),
458 connection=self.connection,
459 state=MaintenanceState.MAINTENANCE,
460 notification=notification,
461 config=self.config,
462 ),
463 )
466class AsyncOSSMaintNotificationsHandler:
467 """
468 Cluster-wide handler for OSS (open-source) cluster maintenance notifications.
470 Reacts to SMIGRATED (slot migration completed) push notifications and
471 triggers topology re-initialization via nodes_manager.initialize().
473 Lock discipline: _lock is held across the whole pool mutation, including
474 await initialize(), so the topology refresh and the subsequent connection
475 marking/disconnect happen as one atomic operation — releasing the lock
476 mid-mutation would let other tasks observe partial state.
478 Re-entrancy is not a deadlock risk here even though asyncio.Lock is
479 non-reentrant: initialize() may dispatch commands whose responses carry new
480 push notifications, but handle_notification schedules the actual handling as
481 a separate background task rather than calling into it inline, so the
482 re-entrant arrival never tries to re-acquire the lock on this call stack.
483 The cheap _in_progress/_processed dedup that gates that scheduling runs
484 without the lock — those sets are only mutated from the single event loop.
486 For the same reason the sync handler's lock-ordering rule (acquire
487 NodesManager._initialization_lock before the handler's _lock) does not apply
488 here: because the handling never runs inline on the call stack that is inside
489 initialize(), no task ever holds NodesManager._initialize_lock while waiting
490 for this _lock, so the two can be acquired in this order safely. The
491 divergence from the sync handler is deliberate, not drift.
492 """
494 def __init__(
495 self,
496 cluster_client: "RedisCluster",
497 config: MaintNotificationsConfig,
498 ) -> None:
499 self.cluster_client = cluster_client
500 self.config = config
501 self._processed_notifications: set[MaintenanceNotification] = set()
502 self._in_progress: set[MaintenanceNotification] = set()
503 self._lock = asyncio.Lock()
504 self._background_tasks: set[asyncio.Task] = set()
506 async def remove_expired_notifications(self) -> None:
507 async with self._lock:
508 for n in tuple(self._processed_notifications):
509 if n.is_expired():
510 self._processed_notifications.remove(n)
512 async def handle_notification(self, notification: MaintenanceNotification) -> None:
513 # Synchronous pre-dedup BEFORE scheduling a task.
514 #
515 # The same SMIGRATED notification is delivered by the server on every
516 # connection, so under load this callback fires many times for the same
517 # notification. Without this guard we would schedule one background task
518 # per arrival - a flood of tasks that each acquire the lock and dedup,
519 # saturating the single event loop. Reserving the notification in
520 # _in_progress here (and skipping if already in-progress/processed) caps
521 # it at exactly ONE handling task per unique notification.
522 #
523 # These sets are only mutated from the single event loop, so this
524 # check-and-add is race-free without holding the lock.
525 if (
526 notification in self._in_progress
527 or notification in self._processed_notifications
528 ):
529 return
530 self._in_progress.add(notification)
532 # Schedule as a background task so the parser's read path is not blocked.
533 # This also breaks the inline call chain that would otherwise deadlock:
534 # initialize() dispatches commands whose responses may carry more push
535 # notifications; as a separate task it can run while this one awaits.
536 #
537 # If scheduling the task (or wiring its callbacks) fails, release the
538 # in-progress reservation made above. Otherwise the notification would be
539 # stuck in _in_progress forever - the dedup guard would skip it on every
540 # future arrival, and _do_handle_notification's finally (which normally
541 # clears it) never runs because the task was never started.
542 try:
543 task = asyncio.get_running_loop().create_task(
544 self._do_handle_notification(notification)
545 )
546 self._background_tasks.add(task)
547 task.add_done_callback(self._background_tasks.discard)
548 task.add_done_callback(
549 functools.partial(
550 _log_task_exception,
551 "Error handling maintenance notification background task",
552 )
553 )
554 except Exception:
555 self._in_progress.discard(notification)
556 raise
558 async def _do_handle_notification(
559 self, notification: MaintenanceNotification
560 ) -> None:
561 try:
562 if isinstance(notification, OSSNodeMigratedNotification):
563 await self.handle_oss_maintenance_completed_notification(notification)
564 else:
565 logger.error(f"Unhandled notification type: {notification}")
566 finally:
567 # Release the in-progress reservation. On success the notification is
568 # also in _processed_notifications (so it won't be re-handled); on
569 # failure it is not, allowing a later retry.
570 self._in_progress.discard(notification)
572 async def handle_oss_maintenance_completed_notification(
573 self, notification: OSSNodeMigratedNotification
574 ) -> None:
575 await self.remove_expired_notifications()
577 async with self._lock:
578 # handle_notification already reserved this notification in
579 # _in_progress and guaranteed uniqueness; the processed check here is
580 # defensive (e.g. across handler copies sharing the same sets).
581 if notification in self._processed_notifications:
582 return
583 if logger.isEnabledFor(logging.DEBUG):
584 logger.debug(f"Handling SMIGRATED notification: {notification}")
586 # Extract the information about the src and destination nodes that are
587 # affected by the maintenance. nodes_to_slots_mapping structure:
588 # {
589 # "src_host:port": [
590 # {"dest_host:port": "slot_range"},
591 # ...
592 # ],
593 # ...
594 # }
595 additional_startup_nodes_info = []
596 affected_nodes = set()
597 for (
598 src_address,
599 dest_mappings,
600 ) in notification.nodes_to_slots_mapping.items():
601 src_host, src_port = src_address.rsplit(":", 1)
602 src_node = self.cluster_client.nodes_manager.get_node(
603 host=src_host, port=int(src_port)
604 )
605 if src_node is not None:
606 affected_nodes.add(src_node)
607 for dest_mapping in dest_mappings:
608 for dest_address in dest_mapping.keys():
609 dest_host, dest_port = dest_address.rsplit(":", 1)
610 additional_startup_nodes_info.append(
611 (dest_host, int(dest_port))
612 )
613 # Updates the cluster slots cache with the new slots mapping
614 # This will also update the nodes cache with the new nodes mapping
615 await self.cluster_client.nodes_manager.initialize(
616 additional_startup_nodes_info=additional_startup_nodes_info,
617 )
619 all_nodes = set(affected_nodes)
620 all_nodes = all_nodes.union(
621 self.cluster_client.nodes_manager.nodes_cache.values()
622 )
623 for current_node in all_nodes:
624 handoff_recorded = False
625 if current_node in affected_nodes:
626 # mark for reconnect all in-use connections to the node — this
627 # forces them to disconnect after completing their current command
628 free_set = set(current_node._free)
629 for conn in current_node._connections:
630 if conn not in free_set:
631 add_debug_log_for_notification(
632 conn, "SMIGRATED - mark for reconnect"
633 )
634 conn.mark_for_reconnect()
635 await record_connection_handoff(
636 pool_name=f"{current_node.host}:{current_node.port}"
637 )
638 handoff_recorded = True
639 else:
640 if logger.isEnabledFor(logging.DEBUG):
641 logger.debug(
642 f"SMIGRATED: Node {current_node.name} not affected "
643 f"by maintenance, skipping mark for reconnect"
644 )
645 if (
646 current_node
647 not in self.cluster_client.nodes_manager.nodes_cache.values()
648 ):
649 task = asyncio.get_running_loop().create_task(
650 current_node.disconnect_free_connections()
651 )
652 self._background_tasks.add(task)
653 task.add_done_callback(self._background_tasks.discard)
654 task.add_done_callback(
655 functools.partial(
656 _log_task_exception,
657 "Error disconnecting free connections after "
658 "maintenance notification",
659 )
660 )
661 if not handoff_recorded:
662 await record_connection_handoff(
663 pool_name=f"{current_node.host}:{current_node.port}"
664 )
666 self._processed_notifications.add(notification)