Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/maint_notifications.py: 28%

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

389 statements  

1import enum 

2import ipaddress 

3import logging 

4import re 

5import threading 

6import time 

7from abc import ABC, abstractmethod 

8from typing import ( 

9 TYPE_CHECKING, 

10 Any, 

11 Dict, 

12 List, 

13 Literal, 

14 Mapping, 

15 Optional, 

16 Union, 

17) 

18 

19from redis.observability.attributes import get_pool_name 

20from redis.observability.recorder import ( 

21 record_connection_handoff, 

22 record_connection_relaxed_timeout, 

23 record_maint_notification_count, 

24) 

25from redis.typing import Number 

26 

27if TYPE_CHECKING: 

28 from redis.cluster import MaintNotificationsAbstractRedisCluster 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33class MaintenanceState(enum.Enum): 

34 NONE = "none" 

35 MOVING = "moving" 

36 MAINTENANCE = "maintenance" 

37 

38 

39class EndpointType(enum.Enum): 

40 """Valid endpoint types used in CLIENT MAINT_NOTIFICATIONS command.""" 

41 

42 INTERNAL_IP = "internal-ip" 

43 INTERNAL_FQDN = "internal-fqdn" 

44 EXTERNAL_IP = "external-ip" 

45 EXTERNAL_FQDN = "external-fqdn" 

46 NONE = "none" 

47 

48 def __str__(self): 

49 """Return the string value of the enum.""" 

50 return self.value 

51 

52 

53if TYPE_CHECKING: 

54 from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection 

55 from redis.connection import ( 

56 MaintNotificationsAbstractConnection, 

57 MaintNotificationsAbstractConnectionPool, 

58 ) 

59 

60 

61class MaintenanceNotification(ABC): 

62 """ 

63 Base class for maintenance notifications sent through push messages by Redis server. 

64 

65 This class provides common functionality for all maintenance notifications including 

66 unique identification and TTL (Time-To-Live) functionality. 

67 

68 Attributes: 

69 id (int): Unique identifier for this notification 

70 ttl (int): Time-to-live in seconds for this notification 

71 creation_time (float): Timestamp when the notification was created/read 

72 """ 

73 

74 def __init__(self, id: int, ttl: int): 

75 """ 

76 Initialize a new MaintenanceNotification with unique ID and TTL functionality. 

77 

78 Args: 

79 id (int): Unique identifier for this notification 

80 ttl (int): Time-to-live in seconds for this notification 

81 """ 

82 self.id = id 

83 self.ttl = ttl 

84 self.creation_time = time.monotonic() 

85 self.expire_at = self.creation_time + self.ttl 

86 

87 def is_expired(self) -> bool: 

88 """ 

89 Check if this notification has expired based on its TTL 

90 and creation time. 

91 

92 Returns: 

93 bool: True if the notification has expired, False otherwise 

94 """ 

95 return time.monotonic() > (self.creation_time + self.ttl) 

96 

97 @abstractmethod 

98 def __repr__(self) -> str: 

99 """ 

100 Return a string representation of the maintenance notification. 

101 

102 This method must be implemented by all concrete subclasses. 

103 

104 Returns: 

105 str: String representation of the notification 

106 """ 

107 pass 

108 

109 @abstractmethod 

110 def __eq__(self, other) -> bool: 

111 """ 

112 Compare two maintenance notifications for equality. 

113 

114 This method must be implemented by all concrete subclasses. 

115 Notifications are typically considered equal if they have the same id 

116 and are of the same type. 

117 

118 Args: 

119 other: The other object to compare with 

120 

121 Returns: 

122 bool: True if the notifications are equal, False otherwise 

123 """ 

124 pass 

125 

126 @abstractmethod 

127 def __hash__(self) -> int: 

128 """ 

129 Return a hash value for the maintenance notification. 

130 

131 This method must be implemented by all concrete subclasses to allow 

132 instances to be used in sets and as dictionary keys. 

133 

134 Returns: 

135 int: Hash value for the notification 

136 """ 

137 pass 

138 

139 

140class NodeMovingNotification(MaintenanceNotification): 

141 """ 

142 This notification is received when a node is replaced with a new node 

143 during cluster rebalancing or maintenance operations. 

144 """ 

145 

146 def __init__( 

147 self, 

148 id: int, 

149 new_node_host: Optional[str], 

150 new_node_port: Optional[int], 

151 ttl: int, 

152 ): 

153 """ 

154 Initialize a new NodeMovingNotification. 

155 

156 Args: 

157 id (int): Unique identifier for this notification 

158 new_node_host (str): Hostname or IP address of the new replacement node 

159 new_node_port (int): Port number of the new replacement node 

160 ttl (int): Time-to-live in seconds for this notification 

161 """ 

162 super().__init__(id, ttl) 

163 self.new_node_host = new_node_host 

164 self.new_node_port = new_node_port 

165 

166 def __repr__(self) -> str: 

167 expiry_time = self.expire_at 

168 remaining = max(0, expiry_time - time.monotonic()) 

169 

170 return ( 

171 f"{self.__class__.__name__}(" 

172 f"id={self.id}, " 

173 f"new_node_host='{self.new_node_host}', " 

174 f"new_node_port={self.new_node_port}, " 

175 f"ttl={self.ttl}, " 

176 f"creation_time={self.creation_time}, " 

177 f"expires_at={expiry_time}, " 

178 f"remaining={remaining:.1f}s, " 

179 f"expired={self.is_expired()}" 

180 f")" 

181 ) 

182 

183 def __eq__(self, other) -> bool: 

184 """ 

185 Two NodeMovingNotification notifications are considered equal if they have the same 

186 id, new_node_host, and new_node_port. 

187 """ 

188 if not isinstance(other, NodeMovingNotification): 

189 return False 

190 return ( 

191 self.id == other.id 

192 and self.new_node_host == other.new_node_host 

193 and self.new_node_port == other.new_node_port 

194 ) 

195 

196 def __hash__(self) -> int: 

197 """ 

198 Return a hash value for the notification to allow 

199 instances to be used in sets and as dictionary keys. 

200 

201 Returns: 

202 int: Hash value based on notification type class name, id, 

203 new_node_host and new_node_port 

204 """ 

205 try: 

206 node_port = int(self.new_node_port) if self.new_node_port else None 

207 except ValueError: 

208 node_port = 0 

209 

210 return hash( 

211 ( 

212 self.__class__.__name__, 

213 int(self.id), 

214 str(self.new_node_host), 

215 node_port, 

216 ) 

217 ) 

218 

219 

220class NodeMigratingNotification(MaintenanceNotification): 

221 """ 

222 Notification for when a Redis cluster node is in the process of migrating slots. 

223 

224 This notification is received when a node starts migrating its slots to another node 

225 during cluster rebalancing or maintenance operations. 

226 

227 Args: 

228 id (int): Unique identifier for this notification 

229 ttl (int): Time-to-live in seconds for this notification 

230 """ 

231 

232 def __init__(self, id: int, ttl: int): 

233 super().__init__(id, ttl) 

234 

235 def __repr__(self) -> str: 

236 expiry_time = self.creation_time + self.ttl 

237 remaining = max(0, expiry_time - time.monotonic()) 

238 return ( 

239 f"{self.__class__.__name__}(" 

240 f"id={self.id}, " 

241 f"ttl={self.ttl}, " 

242 f"creation_time={self.creation_time}, " 

243 f"expires_at={expiry_time}, " 

244 f"remaining={remaining:.1f}s, " 

245 f"expired={self.is_expired()}" 

246 f")" 

247 ) 

248 

249 def __eq__(self, other) -> bool: 

250 """ 

251 Two NodeMigratingNotification notifications are considered equal if they have the same 

252 id and are of the same type. 

253 """ 

254 if not isinstance(other, NodeMigratingNotification): 

255 return False 

256 return self.id == other.id and type(self) is type(other) 

257 

258 def __hash__(self) -> int: 

259 """ 

260 Return a hash value for the notification to allow 

261 instances to be used in sets and as dictionary keys. 

262 

263 Returns: 

264 int: Hash value based on notification type and id 

265 """ 

266 return hash((self.__class__.__name__, int(self.id))) 

267 

268 

269class NodeMigratedNotification(MaintenanceNotification): 

270 """ 

271 Notification for when a Redis cluster node has completed migrating slots. 

272 

273 This notification is received when a node has finished migrating all its slots 

274 to other nodes during cluster rebalancing or maintenance operations. 

275 

276 Args: 

277 id (int): Unique identifier for this notification 

278 """ 

279 

280 DEFAULT_TTL = 5 

281 

282 def __init__(self, id: int): 

283 super().__init__(id, NodeMigratedNotification.DEFAULT_TTL) 

284 

285 def __repr__(self) -> str: 

286 expiry_time = self.creation_time + self.ttl 

287 remaining = max(0, expiry_time - time.monotonic()) 

288 return ( 

289 f"{self.__class__.__name__}(" 

290 f"id={self.id}, " 

291 f"ttl={self.ttl}, " 

292 f"creation_time={self.creation_time}, " 

293 f"expires_at={expiry_time}, " 

294 f"remaining={remaining:.1f}s, " 

295 f"expired={self.is_expired()}" 

296 f")" 

297 ) 

298 

299 def __eq__(self, other) -> bool: 

300 """ 

301 Two NodeMigratedNotification notifications are considered equal if they have the same 

302 id and are of the same type. 

303 """ 

304 if not isinstance(other, NodeMigratedNotification): 

305 return False 

306 return self.id == other.id and type(self) is type(other) 

307 

308 def __hash__(self) -> int: 

309 """ 

310 Return a hash value for the notification to allow 

311 instances to be used in sets and as dictionary keys. 

312 

313 Returns: 

314 int: Hash value based on notification type and id 

315 """ 

316 return hash((self.__class__.__name__, int(self.id))) 

317 

318 

319class NodeFailingOverNotification(MaintenanceNotification): 

320 """ 

321 Notification for when a Redis cluster node is in the process of failing over. 

322 

323 This notification is received when a node starts a failover process during 

324 cluster maintenance operations or when handling node failures. 

325 

326 Args: 

327 id (int): Unique identifier for this notification 

328 ttl (int): Time-to-live in seconds for this notification 

329 """ 

330 

331 def __init__(self, id: int, ttl: int): 

332 super().__init__(id, ttl) 

333 

334 def __repr__(self) -> str: 

335 expiry_time = self.creation_time + self.ttl 

336 remaining = max(0, expiry_time - time.monotonic()) 

337 return ( 

338 f"{self.__class__.__name__}(" 

339 f"id={self.id}, " 

340 f"ttl={self.ttl}, " 

341 f"creation_time={self.creation_time}, " 

342 f"expires_at={expiry_time}, " 

343 f"remaining={remaining:.1f}s, " 

344 f"expired={self.is_expired()}" 

345 f")" 

346 ) 

347 

348 def __eq__(self, other) -> bool: 

349 """ 

350 Two NodeFailingOverNotification notifications are considered equal if they have the same 

351 id and are of the same type. 

352 """ 

353 if not isinstance(other, NodeFailingOverNotification): 

354 return False 

355 return self.id == other.id and type(self) is type(other) 

356 

357 def __hash__(self) -> int: 

358 """ 

359 Return a hash value for the notification to allow 

360 instances to be used in sets and as dictionary keys. 

361 

362 Returns: 

363 int: Hash value based on notification type and id 

364 """ 

365 return hash((self.__class__.__name__, int(self.id))) 

366 

367 

368class NodeFailedOverNotification(MaintenanceNotification): 

369 """ 

370 Notification for when a Redis cluster node has completed a failover. 

371 

372 This notification is received when a node has finished the failover process 

373 during cluster maintenance operations or after handling node failures. 

374 

375 Args: 

376 id (int): Unique identifier for this notification 

377 """ 

378 

379 DEFAULT_TTL = 5 

380 

381 def __init__(self, id: int): 

382 super().__init__(id, NodeFailedOverNotification.DEFAULT_TTL) 

383 

384 def __repr__(self) -> str: 

385 expiry_time = self.creation_time + self.ttl 

386 remaining = max(0, expiry_time - time.monotonic()) 

387 return ( 

388 f"{self.__class__.__name__}(" 

389 f"id={self.id}, " 

390 f"ttl={self.ttl}, " 

391 f"creation_time={self.creation_time}, " 

392 f"expires_at={expiry_time}, " 

393 f"remaining={remaining:.1f}s, " 

394 f"expired={self.is_expired()}" 

395 f")" 

396 ) 

397 

398 def __eq__(self, other) -> bool: 

399 """ 

400 Two NodeFailedOverNotification notifications are considered equal if they have the same 

401 id and are of the same type. 

402 """ 

403 if not isinstance(other, NodeFailedOverNotification): 

404 return False 

405 return self.id == other.id and type(self) is type(other) 

406 

407 def __hash__(self) -> int: 

408 """ 

409 Return a hash value for the notification to allow 

410 instances to be used in sets and as dictionary keys. 

411 

412 Returns: 

413 int: Hash value based on notification type and id 

414 """ 

415 return hash((self.__class__.__name__, int(self.id))) 

416 

417 

418class OSSNodeMigratingNotification(MaintenanceNotification): 

419 """ 

420 Notification for when a Redis OSS API client is used and a node is in the process of migrating slots. 

421 

422 This notification is received when a node starts migrating its slots to another node 

423 during cluster rebalancing or maintenance operations. 

424 

425 Args: 

426 id (int): Unique identifier for this notification 

427 slots (Optional[List[int]]): List of slots being migrated 

428 """ 

429 

430 DEFAULT_TTL = 30 

431 

432 def __init__( 

433 self, 

434 id: int, 

435 slots: Optional[str] = None, 

436 ): 

437 super().__init__(id, OSSNodeMigratingNotification.DEFAULT_TTL) 

438 self.slots = slots 

439 

440 def __repr__(self) -> str: 

441 expiry_time = self.creation_time + self.ttl 

442 remaining = max(0, expiry_time - time.monotonic()) 

443 return ( 

444 f"{self.__class__.__name__}(" 

445 f"id={self.id}, " 

446 f"slots={self.slots}, " 

447 f"ttl={self.ttl}, " 

448 f"creation_time={self.creation_time}, " 

449 f"expires_at={expiry_time}, " 

450 f"remaining={remaining:.1f}s, " 

451 f"expired={self.is_expired()}" 

452 f")" 

453 ) 

454 

455 def __eq__(self, other) -> bool: 

456 """ 

457 Two OSSNodeMigratingNotification notifications are considered equal if they have the same 

458 id and are of the same type. 

459 """ 

460 if not isinstance(other, OSSNodeMigratingNotification): 

461 return False 

462 return self.id == other.id and type(self) is type(other) 

463 

464 def __hash__(self) -> int: 

465 """ 

466 Return a hash value for the notification to allow 

467 instances to be used in sets and as dictionary keys. 

468 

469 Returns: 

470 int: Hash value based on notification type and id 

471 """ 

472 return hash((self.__class__.__name__, int(self.id))) 

473 

474 

475class OSSNodeMigratedNotification(MaintenanceNotification): 

476 """ 

477 Notification for when a Redis OSS API client is used and a node has completed migrating slots. 

478 

479 This notification is received when a node has finished migrating all its slots 

480 to other nodes during cluster rebalancing or maintenance operations. 

481 

482 Args: 

483 id (int): Unique identifier for this notification 

484 nodes_to_slots_mapping (Dict[str, List[Dict[str, str]]]): Map of source node address 

485 to list of destination mappings. Each destination mapping is a dict with 

486 the destination node address as key and the slot range as value. 

487 

488 Structure example: 

489 { 

490 "127.0.0.1:6379": [ 

491 {"127.0.0.1:6380": "1-100"}, 

492 {"127.0.0.1:6381": "101-200"} 

493 ], 

494 "127.0.0.1:6382": [ 

495 {"127.0.0.1:6383": "201-300"} 

496 ] 

497 } 

498 

499 Where: 

500 - Key (str): Source node address in "host:port" format 

501 - Value (List[Dict[str, str]]): List of destination mappings where each dict 

502 contains destination node address as key and slot range as value 

503 """ 

504 

505 DEFAULT_TTL = 120 

506 

507 def __init__( 

508 self, 

509 id: int, 

510 nodes_to_slots_mapping: Dict[str, List[Dict[str, str]]], 

511 ): 

512 super().__init__(id, OSSNodeMigratedNotification.DEFAULT_TTL) 

513 self.nodes_to_slots_mapping = nodes_to_slots_mapping 

514 

515 def __repr__(self) -> str: 

516 expiry_time = self.creation_time + self.ttl 

517 remaining = max(0, expiry_time - time.monotonic()) 

518 return ( 

519 f"{self.__class__.__name__}(" 

520 f"id={self.id}, " 

521 f"nodes_to_slots_mapping={self.nodes_to_slots_mapping}, " 

522 f"ttl={self.ttl}, " 

523 f"creation_time={self.creation_time}, " 

524 f"expires_at={expiry_time}, " 

525 f"remaining={remaining:.1f}s, " 

526 f"expired={self.is_expired()}" 

527 f")" 

528 ) 

529 

530 def __eq__(self, other) -> bool: 

531 """ 

532 Two OSSNodeMigratedNotification notifications are considered equal if they have the same 

533 id and are of the same type. 

534 """ 

535 if not isinstance(other, OSSNodeMigratedNotification): 

536 return False 

537 return self.id == other.id and type(self) is type(other) 

538 

539 def __hash__(self) -> int: 

540 """ 

541 Return a hash value for the notification to allow 

542 instances to be used in sets and as dictionary keys. 

543 

544 Returns: 

545 int: Hash value based on notification type and id 

546 """ 

547 return hash((self.__class__.__name__, int(self.id))) 

548 

549 

550def _is_private_fqdn(host: str) -> bool: 

551 """ 

552 Determine if an FQDN is likely to be internal/private. 

553 

554 This uses heuristics based on RFC 952 and RFC 1123 standards: 

555 - .local domains (RFC 6762 - Multicast DNS) 

556 - .internal domains (common internal convention) 

557 - Single-label hostnames (no dots) 

558 - Common internal TLDs 

559 

560 Args: 

561 host (str): The FQDN to check 

562 

563 Returns: 

564 bool: True if the FQDN appears to be internal/private 

565 """ 

566 host_lower = host.lower().rstrip(".") 

567 

568 # Single-label hostnames (no dots) are typically internal 

569 if "." not in host_lower: 

570 return True 

571 

572 # Common internal/private domain patterns 

573 internal_patterns = [ 

574 r"\.local$", # mDNS/Bonjour domains 

575 r"\.internal$", # Common internal convention 

576 r"\.corp$", # Corporate domains 

577 r"\.lan$", # Local area network 

578 r"\.intranet$", # Intranet domains 

579 r"\.private$", # Private domains 

580 ] 

581 

582 for pattern in internal_patterns: 

583 if re.search(pattern, host_lower): 

584 return True 

585 

586 # If none of the internal patterns match, assume it's external 

587 return False 

588 

589 

590notification_types_mapping: dict[type[MaintenanceNotification], str] = { 

591 NodeMovingNotification: "MOVING", 

592 NodeMigratingNotification: "MIGRATING", 

593 NodeMigratedNotification: "MIGRATED", 

594 NodeFailingOverNotification: "FAILING_OVER", 

595 NodeFailedOverNotification: "FAILED_OVER", 

596 OSSNodeMigratingNotification: "SMIGRATING", 

597 OSSNodeMigratedNotification: "SMIGRATED", 

598} 

599 

600 

601def add_debug_log_for_notification( 

602 connection: "MaintNotificationsAbstractConnection", 

603 notification: Union[str, MaintenanceNotification], 

604): 

605 if logger.isEnabledFor(logging.DEBUG): 

606 socket_address = None 

607 try: 

608 socket_address = ( 

609 connection._sock.getsockname() if connection._sock else None 

610 ) 

611 socket_address = socket_address[1] if socket_address else None 

612 except (AttributeError, OSError): 

613 pass 

614 

615 logger.debug( 

616 f"Handling maintenance notification: {notification}, " 

617 f"with connection: {connection}, connected to ip {connection.get_resolved_ip()}, " 

618 f"local socket port: {socket_address}", 

619 ) 

620 

621 

622class MaintNotificationsConfig: 

623 """ 

624 Configuration class for maintenance notifications handling behaviour. Notifications are received through 

625 push notifications. 

626 

627 This class defines how the Redis client should react to different push notifications 

628 such as node moving, migrations, etc. in a Redis cluster. 

629 

630 """ 

631 

632 def __init__( 

633 self, 

634 enabled: Union[bool, Literal["auto"]] = "auto", 

635 proactive_reconnect: bool = True, 

636 relaxed_timeout: Optional[Number] = 10, 

637 endpoint_type: Optional[EndpointType] = None, 

638 ): 

639 """ 

640 Initialize a new MaintNotificationsConfig. 

641 

642 Args: 

643 enabled (bool | "auto"): Controls maintenance notifications handling behavior. 

644 - True: The CLIENT MAINT_NOTIFICATIONS command must succeed during connection setup, 

645 otherwise a ResponseError is raised. 

646 - "auto": The CLIENT MAINT_NOTIFICATIONS command is attempted but failures are 

647 gracefully handled - a warning is logged and normal operation continues. 

648 - False: Maintenance notifications are completely disabled. 

649 Defaults to "auto". 

650 proactive_reconnect (bool): Whether to proactively reconnect when a node is replaced. 

651 Defaults to True. 

652 relaxed_timeout (Number): The relaxed timeout to use for the connection during maintenance. 

653 If -1 is provided - the relaxed timeout is disabled. Defaults to 20. 

654 endpoint_type (Optional[EndpointType]): Override for the endpoint type to use in CLIENT MAINT_NOTIFICATIONS. 

655 If None, the endpoint type will be automatically determined based on the host and TLS configuration. 

656 Defaults to None. 

657 

658 Raises: 

659 ValueError: If endpoint_type is provided but is not a valid endpoint type. 

660 """ 

661 self.enabled = enabled 

662 self.relaxed_timeout = relaxed_timeout 

663 self.proactive_reconnect = proactive_reconnect 

664 self.endpoint_type = endpoint_type 

665 

666 def __repr__(self) -> str: 

667 return ( 

668 f"{self.__class__.__name__}(" 

669 f"enabled={self.enabled}, " 

670 f"proactive_reconnect={self.proactive_reconnect}, " 

671 f"relaxed_timeout={self.relaxed_timeout}, " 

672 f"endpoint_type={self.endpoint_type!r}" 

673 f")" 

674 ) 

675 

676 def is_relaxed_timeouts_enabled(self) -> bool: 

677 """ 

678 Check if the relaxed_timeout is enabled. The '-1' value is used to disable the relaxed_timeout. 

679 If relaxed_timeout is set to None, it will make the operation blocking 

680 and waiting until any response is received. 

681 

682 Returns: 

683 True if the relaxed_timeout is enabled, False otherwise. 

684 """ 

685 return self.relaxed_timeout != -1 

686 

687 def get_endpoint_type( 

688 self, 

689 host: str, 

690 connection: "MaintNotificationsAbstractConnection | AsyncMaintNotificationsAbstractConnection", 

691 ) -> EndpointType: 

692 """ 

693 Determine the appropriate endpoint type for CLIENT MAINT_NOTIFICATIONS command. 

694 

695 Logic: 

696 1. If endpoint_type is explicitly set, use it 

697 2. Otherwise, check the original host from connection.host: 

698 - If host is an IP address, use it directly to determine internal-ip vs external-ip 

699 - If host is an FQDN, get the resolved IP to determine internal-fqdn vs external-fqdn 

700 

701 Args: 

702 host: User provided hostname to analyze 

703 connection: The connection object to analyze for endpoint type determination 

704 

705 Returns: 

706 """ 

707 

708 # If endpoint_type is explicitly set, use it 

709 if self.endpoint_type is not None: 

710 return self.endpoint_type 

711 

712 # Check if the host is an IP address 

713 try: 

714 ip_addr = ipaddress.ip_address(host) 

715 # Host is an IP address - use it directly 

716 is_private = ip_addr.is_private 

717 return EndpointType.INTERNAL_IP if is_private else EndpointType.EXTERNAL_IP 

718 except ValueError: 

719 # Host is an FQDN - need to check resolved IP to determine internal vs external 

720 pass 

721 

722 # Host is an FQDN, get the resolved IP to determine if it's internal or external 

723 resolved_ip = connection.get_resolved_ip() 

724 

725 if resolved_ip: 

726 try: 

727 ip_addr = ipaddress.ip_address(resolved_ip) 

728 is_private = ip_addr.is_private 

729 # Use FQDN types since the original host was an FQDN 

730 return ( 

731 EndpointType.INTERNAL_FQDN 

732 if is_private 

733 else EndpointType.EXTERNAL_FQDN 

734 ) 

735 except ValueError: 

736 # This shouldn't happen since we got the IP from the socket, but fallback 

737 pass 

738 

739 # Final fallback: use heuristics on the FQDN itself 

740 is_private = _is_private_fqdn(host) 

741 return EndpointType.INTERNAL_FQDN if is_private else EndpointType.EXTERNAL_FQDN 

742 

743 

744_MAINTENANCE_START_NOTIFICATION_TYPES = ( 

745 NodeMigratingNotification, 

746 NodeFailingOverNotification, 

747 OSSNodeMigratingNotification, 

748) 

749_MAINTENANCE_COMPLETED_NOTIFICATION_TYPES = ( 

750 NodeMigratedNotification, 

751 NodeFailedOverNotification, 

752 OSSNodeMigratedNotification, 

753) 

754 

755 

756def _get_maintenance_notification_type( 

757 notification: MaintenanceNotification, 

758) -> Optional[int]: 

759 if notification.__class__ in _MAINTENANCE_START_NOTIFICATION_TYPES: 

760 return 1 

761 if notification.__class__ in _MAINTENANCE_COMPLETED_NOTIFICATION_TYPES: 

762 return 0 

763 return None 

764 

765 

766def _get_maintenance_notification_name( 

767 notification: MaintenanceNotification, 

768) -> str: 

769 return notification_types_mapping.get(notification.__class__, "") 

770 

771 

772def _should_skip_connection_timeout_update( 

773 maintenance_state: MaintenanceState, 

774 config: MaintNotificationsConfig, 

775) -> bool: 

776 return ( 

777 maintenance_state == MaintenanceState.MOVING 

778 or not config.is_relaxed_timeouts_enabled() 

779 ) 

780 

781 

782def _build_moving_connection_kwargs( 

783 notification: NodeMovingNotification, 

784 config: MaintNotificationsConfig, 

785) -> dict[str, Any]: 

786 kwargs: dict[str, Any] = { 

787 "maintenance_state": MaintenanceState.MOVING, 

788 "maintenance_notification_hash": hash(notification), 

789 } 

790 if notification.new_node_host is not None: 

791 # the host is not updated if the new node host is None 

792 # this happens when the MOVING push notification does not contain 

793 # the new node host - in this case we only update the timeouts 

794 kwargs["host"] = notification.new_node_host 

795 if config.is_relaxed_timeouts_enabled(): 

796 kwargs.update( 

797 { 

798 "socket_timeout": config.relaxed_timeout, 

799 "socket_connect_timeout": config.relaxed_timeout, 

800 } 

801 ) 

802 return kwargs 

803 

804 

805def _build_moving_cleanup_connection_kwargs( 

806 connection_kwargs: Mapping[str, Any], 

807 notification_hash: int, 

808) -> Optional[dict[str, Any]]: 

809 # if the current maintenance_notification_hash in kwargs is not matching the notification 

810 # it means there has been a new moving notification after this one 

811 # and we don't need to revert the kwargs yet 

812 if connection_kwargs.get("maintenance_notification_hash") != notification_hash: 

813 return None 

814 

815 return { 

816 "maintenance_state": MaintenanceState.NONE, 

817 "maintenance_notification_hash": None, 

818 "host": connection_kwargs.get("orig_host_address"), 

819 "socket_timeout": connection_kwargs.get("orig_socket_timeout"), 

820 "socket_connect_timeout": connection_kwargs.get("orig_socket_connect_timeout"), 

821 } 

822 

823 

824class MaintNotificationsPoolHandler: 

825 def __init__( 

826 self, 

827 pool: "MaintNotificationsAbstractConnectionPool", 

828 config: MaintNotificationsConfig, 

829 ) -> None: 

830 self.pool = pool 

831 self.config = config 

832 self._processed_notifications = set() 

833 self._lock = threading.RLock() 

834 self.connection = None 

835 

836 def set_connection(self, connection: "MaintNotificationsAbstractConnection"): 

837 self.connection = connection 

838 

839 def get_handler_for_connection(self): 

840 # Copy all data that should be shared between connections 

841 # but each connection should have its own pool handler 

842 # since each connection can be in a different state 

843 copy = MaintNotificationsPoolHandler(self.pool, self.config) 

844 copy._processed_notifications = self._processed_notifications 

845 copy._lock = self._lock 

846 copy.connection = None 

847 return copy 

848 

849 def remove_expired_notifications(self): 

850 with self._lock: 

851 for notification in tuple(self._processed_notifications): 

852 if notification.is_expired(): 

853 self._processed_notifications.remove(notification) 

854 

855 def handle_notification(self, notification: MaintenanceNotification): 

856 self.remove_expired_notifications() 

857 

858 if isinstance(notification, NodeMovingNotification): 

859 return self.handle_node_moving_notification(notification) 

860 else: 

861 logger.error(f"Unhandled notification type: {notification}") 

862 

863 def handle_node_moving_notification(self, notification: NodeMovingNotification): 

864 if ( 

865 not self.config.proactive_reconnect 

866 and not self.config.is_relaxed_timeouts_enabled() 

867 ): 

868 return 

869 

870 with self._lock: 

871 if notification in self._processed_notifications: 

872 # nothing to do in the connection pool handling 

873 # the notification has already been handled or is expired 

874 # just return 

875 return 

876 

877 with self.pool._lock: 

878 if logger.isEnabledFor(logging.DEBUG): 

879 logger.debug( 

880 f"Handling node MOVING notification: {notification}, " 

881 f"with connection: {self.connection}, connected to ip " 

882 f"{self.connection.get_resolved_ip() if self.connection else None}" 

883 ) 

884 # Get the current connected address - if any 

885 # This is the address that is being moved 

886 # and we need to handle only connections 

887 # connected to the same address 

888 moving_address_src = ( 

889 self.connection.getpeername() if self.connection else None 

890 ) 

891 

892 if getattr(self.pool, "set_in_maintenance", False): 

893 # Set pool in maintenance mode - executed only if 

894 # BlockingConnectionPool is used 

895 self.pool.set_in_maintenance(True) 

896 

897 # Update maintenance state, timeout and optionally host address 

898 # connection settings for matching connections 

899 self.pool.update_connections_settings( 

900 state=MaintenanceState.MOVING, 

901 maintenance_notification_hash=hash(notification), 

902 relaxed_timeout=self.config.relaxed_timeout, 

903 host_address=notification.new_node_host, 

904 matching_address=moving_address_src, 

905 matching_pattern="connected_address", 

906 update_notification_hash=True, 

907 include_free_connections=True, 

908 ) 

909 

910 if self.config.proactive_reconnect: 

911 if notification.new_node_host is not None: 

912 self.run_proactive_reconnect(moving_address_src) 

913 else: 

914 threading.Timer( 

915 notification.ttl / 2, 

916 self.run_proactive_reconnect, 

917 args=(moving_address_src,), 

918 ).start() 

919 

920 # Update config for new connections: 

921 # Set state to MOVING 

922 # update host 

923 # if relax timeouts are enabled - update timeouts 

924 self.pool.update_connection_kwargs( 

925 **_build_moving_connection_kwargs(notification, self.config) 

926 ) 

927 

928 if getattr(self.pool, "set_in_maintenance", False): 

929 self.pool.set_in_maintenance(False) 

930 

931 threading.Timer( 

932 notification.ttl, 

933 self.handle_node_moved_notification, 

934 args=(notification,), 

935 ).start() 

936 

937 record_connection_handoff( 

938 pool_name=get_pool_name(self.pool), 

939 ) 

940 

941 self._processed_notifications.add(notification) 

942 

943 def run_proactive_reconnect(self, moving_address_src: Optional[str] = None): 

944 """ 

945 Run proactive reconnect for the pool. 

946 Active connections are marked for reconnect after they complete the current command. 

947 Inactive connections are disconnected and will be connected on next use. 

948 """ 

949 with self._lock: 

950 with self.pool._lock: 

951 # take care for the active connections in the pool 

952 # mark them for reconnect after they complete the current command 

953 self.pool.update_active_connections_for_reconnect( 

954 moving_address_src=moving_address_src, 

955 ) 

956 # take care for the inactive connections in the pool 

957 # delete them and create new ones 

958 self.pool.disconnect_free_connections( 

959 moving_address_src=moving_address_src, 

960 ) 

961 

962 def handle_node_moved_notification(self, notification: NodeMovingNotification): 

963 """ 

964 Handle the cleanup after a node moving notification expires. 

965 """ 

966 notification_hash = hash(notification) 

967 

968 with self._lock: 

969 if logger.isEnabledFor(logging.DEBUG): 

970 logger.debug( 

971 f"Reverting temporary changes related to notification: {notification}, " 

972 f"with connection: {self.connection}, connected to ip " 

973 f"{self.connection.get_resolved_ip() if self.connection else None}" 

974 ) 

975 kwargs = _build_moving_cleanup_connection_kwargs( 

976 self.pool.connection_kwargs, notification_hash 

977 ) 

978 if kwargs is not None: 

979 self.pool.update_connection_kwargs(**kwargs) 

980 

981 with self.pool._lock: 

982 reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled() 

983 reset_host_address = self.config.proactive_reconnect 

984 

985 self.pool.update_connections_settings( 

986 relaxed_timeout=-1, 

987 state=MaintenanceState.NONE, 

988 maintenance_notification_hash=None, 

989 matching_notification_hash=notification_hash, 

990 matching_pattern="notification_hash", 

991 update_notification_hash=True, 

992 reset_relaxed_timeout=reset_relaxed_timeout, 

993 reset_host_address=reset_host_address, 

994 include_free_connections=True, 

995 ) 

996 

997 

998class MaintNotificationsConnectionHandler: 

999 # 1 = "starting maintenance" notifications, 0 = "completed maintenance" notifications 

1000 _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = { 

1001 NodeMigratingNotification: 1, 

1002 NodeFailingOverNotification: 1, 

1003 OSSNodeMigratingNotification: 1, 

1004 NodeMigratedNotification: 0, 

1005 NodeFailedOverNotification: 0, 

1006 OSSNodeMigratedNotification: 0, 

1007 } 

1008 

1009 def __init__( 

1010 self, 

1011 connection: "MaintNotificationsAbstractConnection", 

1012 config: MaintNotificationsConfig, 

1013 ) -> None: 

1014 self.connection = connection 

1015 self.config = config 

1016 

1017 def _get_pool_name(self) -> str: 

1018 """ 

1019 Get the pool name from the connection's pool handler. 

1020 Falls back to connection representation if pool is not available. 

1021 """ 

1022 pool_handler = getattr( 

1023 self.connection, "_maint_notifications_pool_handler", None 

1024 ) 

1025 if pool_handler and getattr(pool_handler, "pool", None): 

1026 return get_pool_name(pool_handler.pool) 

1027 # Fallback for standalone connections without a pool 

1028 return repr(self.connection) 

1029 

1030 def handle_notification(self, notification: MaintenanceNotification): 

1031 # get the notification type by checking its class 

1032 # 1 for start, 0 for end notification type, None for unknown 

1033 notification_type = _get_maintenance_notification_type(notification) 

1034 maint_notification = _get_maintenance_notification_name(notification) 

1035 

1036 record_maint_notification_count( 

1037 server_address=self.connection.host, 

1038 server_port=self.connection.port, 

1039 network_peer_address=self.connection.host, 

1040 network_peer_port=self.connection.port, 

1041 maint_notification=maint_notification, 

1042 ) 

1043 

1044 if notification_type is None: 

1045 logger.error(f"Unhandled notification type: {notification}") 

1046 return 

1047 

1048 if notification_type: 

1049 self.handle_maintenance_start_notification( 

1050 MaintenanceState.MAINTENANCE, notification 

1051 ) 

1052 else: 

1053 self.handle_maintenance_completed_notification(notification=notification) 

1054 

1055 def handle_maintenance_start_notification( 

1056 self, maintenance_state: MaintenanceState, notification: MaintenanceNotification 

1057 ): 

1058 add_debug_log_for_notification(self.connection, notification) 

1059 

1060 if _should_skip_connection_timeout_update( 

1061 self.connection.maintenance_state, self.config 

1062 ): 

1063 return 

1064 

1065 self.connection.maintenance_state = maintenance_state 

1066 self.connection.set_tmp_settings( 

1067 tmp_relaxed_timeout=self.config.relaxed_timeout 

1068 ) 

1069 # extend the timeout for all created connections 

1070 self.connection.update_current_socket_timeout(self.config.relaxed_timeout) 

1071 if isinstance(notification, OSSNodeMigratingNotification): 

1072 # add the notification id to the set of processed start maint notifications 

1073 # this is used to skip the unrelaxing of the timeouts if we have received more than 

1074 # one start notification before the the final end notification 

1075 self.connection.add_maint_start_notification(notification.id) 

1076 

1077 maint_notification = _get_maintenance_notification_name(notification) 

1078 record_connection_relaxed_timeout( 

1079 connection_name=self._get_pool_name(), 

1080 maint_notification=maint_notification, 

1081 relaxed=True, 

1082 ) 

1083 

1084 def handle_maintenance_completed_notification(self, **kwargs: Any) -> None: 

1085 # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled 

1086 if _should_skip_connection_timeout_update( 

1087 self.connection.maintenance_state, self.config 

1088 ): 

1089 return 

1090 

1091 notification = None 

1092 if kwargs.get("notification"): 

1093 notification = kwargs["notification"] 

1094 add_debug_log_for_notification( 

1095 self.connection, notification if notification else "MAINTENANCE_COMPLETED" 

1096 ) 

1097 self.connection.reset_tmp_settings(reset_relaxed_timeout=True) 

1098 # Maintenance completed - reset the connection 

1099 # timeouts by providing -1 as the relaxed timeout 

1100 self.connection.update_current_socket_timeout(-1) 

1101 self.connection.maintenance_state = MaintenanceState.NONE 

1102 # reset the sets that keep track of received start maint 

1103 # notifications and skipped end maint notifications 

1104 self.connection.reset_received_notifications() 

1105 

1106 if notification: 

1107 maint_notification = _get_maintenance_notification_name(notification) 

1108 record_connection_relaxed_timeout( 

1109 connection_name=self._get_pool_name(), 

1110 maint_notification=maint_notification, 

1111 relaxed=False, 

1112 ) 

1113 

1114 

1115class OSSMaintNotificationsHandler: 

1116 def __init__( 

1117 self, 

1118 cluster_client: "MaintNotificationsAbstractRedisCluster", 

1119 config: MaintNotificationsConfig, 

1120 ) -> None: 

1121 self.cluster_client = cluster_client 

1122 self.config = config 

1123 self._processed_notifications = set() 

1124 self._in_progress = set() 

1125 self._lock = threading.RLock() 

1126 

1127 def remove_expired_notifications(self): 

1128 with self._lock: 

1129 for notification in tuple(self._processed_notifications): 

1130 if notification.is_expired(): 

1131 self._processed_notifications.remove(notification) 

1132 

1133 def handle_notification(self, notification: MaintenanceNotification): 

1134 if isinstance(notification, OSSNodeMigratedNotification): 

1135 self.handle_oss_maintenance_completed_notification(notification) 

1136 else: 

1137 logger.error(f"Unhandled notification type: {notification}") 

1138 

1139 def handle_oss_maintenance_completed_notification( 

1140 self, notification: OSSNodeMigratedNotification 

1141 ): 

1142 self.remove_expired_notifications() 

1143 

1144 with self._lock: 

1145 if ( 

1146 notification in self._in_progress 

1147 or notification in self._processed_notifications 

1148 ): 

1149 # we are already handling this notification or it has already been processed 

1150 # we should skip in_progress notification since when we reinitialize the cluster 

1151 # we execute a CLUSTER SLOTS command that can use a different connection 

1152 # that has also has the notification and we don't want to 

1153 # process the same notification twice 

1154 return 

1155 

1156 if logger.isEnabledFor(logging.DEBUG): 

1157 logger.debug(f"Handling SMIGRATED notification: {notification}") 

1158 self._in_progress.add(notification) 

1159 

1160 try: 

1161 # Extract the information about the src and destination nodes that are affected 

1162 # by the maintenance. nodes_to_slots_mapping structure: 

1163 # { 

1164 # "src_host:port": [ 

1165 # {"dest_host:port": "slot_range"}, 

1166 # ... 

1167 # ], 

1168 # ... 

1169 # } 

1170 additional_startup_nodes_info = [] 

1171 affected_nodes = set() 

1172 for ( 

1173 src_address, 

1174 dest_mappings, 

1175 ) in notification.nodes_to_slots_mapping.items(): 

1176 src_host, src_port = src_address.rsplit(":", 1) 

1177 src_node = self.cluster_client.nodes_manager.get_node( 

1178 host=src_host, port=int(src_port) 

1179 ) 

1180 if src_node is not None: 

1181 affected_nodes.add(src_node) 

1182 

1183 for dest_mapping in dest_mappings: 

1184 for dest_address in dest_mapping.keys(): 

1185 dest_host, dest_port = dest_address.rsplit(":", 1) 

1186 additional_startup_nodes_info.append( 

1187 (dest_host, int(dest_port)) 

1188 ) 

1189 

1190 # Updates the cluster slots cache with the new slots mapping 

1191 # This will also update the nodes cache with the new nodes mapping 

1192 self.cluster_client.nodes_manager.initialize( 

1193 disconnect_startup_nodes_pools=False, 

1194 additional_startup_nodes_info=additional_startup_nodes_info, 

1195 ) 

1196 

1197 all_nodes = set(affected_nodes) 

1198 all_nodes = all_nodes.union( 

1199 self.cluster_client.nodes_manager.nodes_cache.values() 

1200 ) 

1201 

1202 for current_node in all_nodes: 

1203 if current_node.redis_connection is None: 

1204 continue 

1205 with current_node.redis_connection.connection_pool._lock: 

1206 handoff_recorded = False 

1207 if current_node in affected_nodes: 

1208 # mark for reconnect all in use connections to the node - this will force them to 

1209 # disconnect after they complete their current commands 

1210 # Some of them might be used by sub sub and we don't know which ones - so we disconnect 

1211 # all in flight connections after they are done with current command execution 

1212 for conn in current_node.redis_connection.connection_pool._get_in_use_connections(): 

1213 add_debug_log_for_notification( 

1214 conn, "SMIGRATED - mark for reconnect" 

1215 ) 

1216 conn.mark_for_reconnect() 

1217 

1218 record_connection_handoff( 

1219 pool_name=get_pool_name( 

1220 current_node.redis_connection.connection_pool 

1221 ) 

1222 ) 

1223 handoff_recorded = True 

1224 else: 

1225 if logger.isEnabledFor(logging.DEBUG): 

1226 logger.debug( 

1227 f"SMIGRATED: Node {current_node.name} not affected by maintenance, " 

1228 f"skipping mark for reconnect" 

1229 ) 

1230 

1231 if ( 

1232 current_node 

1233 not in self.cluster_client.nodes_manager.nodes_cache.values() 

1234 ): 

1235 # disconnect all free connections to the node - this node will be dropped 

1236 # from the cluster, so we don't need to revert the timeouts 

1237 for conn in current_node.redis_connection.connection_pool._get_free_connections(): 

1238 conn.disconnect() 

1239 

1240 # Only record handoff if not already recorded for this node 

1241 if not handoff_recorded: 

1242 record_connection_handoff( 

1243 pool_name=get_pool_name( 

1244 current_node.redis_connection.connection_pool 

1245 ) 

1246 ) 

1247 

1248 # mark the notification as processed 

1249 self._processed_notifications.add(notification) 

1250 finally: 

1251 # Release the in-progress reservation. On success the notification 

1252 # is also in _processed_notifications (so it won't be re-handled); 

1253 # on failure it is not, allowing a later redelivery to retry. 

1254 self._in_progress.discard(notification)