@@ -430,7 +430,20 @@ def set_parser(self, parser_class):
430
430
def set_maintenance_event_pool_handler (
431
431
self , maintenance_event_pool_handler : MaintenanceEventPoolHandler
432
432
):
433
- self ._parser .set_node_moving_push_handler (maintenance_event_pool_handler )
433
+ self ._parser .set_node_moving_push_handler (
434
+ maintenance_event_pool_handler .handle_event
435
+ )
436
+
437
+ # Initialize maintenance event connection handler if it doesn't exist
438
+ if not hasattr (self , "_maintenance_event_connection_handler" ):
439
+ self ._maintenance_event_connection_handler = (
440
+ MaintenanceEventConnectionHandler (
441
+ self , maintenance_event_pool_handler .config
442
+ )
443
+ )
444
+ self ._parser .set_maintenance_push_handler (
445
+ self ._maintenance_event_connection_handler .handle_event
446
+ )
434
447
435
448
def connect (self ):
436
449
"Connects to the Redis server if not already connected"
@@ -796,10 +809,6 @@ def should_reconnect(self):
796
809
def update_current_socket_timeout (self , relax_timeout : Optional [float ] = None ):
797
810
if self ._sock :
798
811
timeout = relax_timeout if relax_timeout != - 1 else self .socket_timeout
799
- logging .debug (
800
- f"***** Connection --> Updating timeout for { self ._sock .getpeername ()} "
801
- f" to timeout { timeout } ; relax_timeout: { relax_timeout } "
802
- )
803
812
self ._sock .settimeout (timeout )
804
813
self .update_parser_buffer_timeout (timeout )
805
814
@@ -852,10 +861,6 @@ def _connect(self):
852
861
# ipv4/ipv6, but we want to set options prior to calling
853
862
# socket.connect()
854
863
err = None
855
- if self .tmp_host_address is not None :
856
- logging .debug (
857
- f"***** Connection --> Using tmp_host_address: { self .tmp_host_address } "
858
- )
859
864
host = self .tmp_host_address or self .host
860
865
861
866
for res in socket .getaddrinfo (
@@ -876,31 +881,18 @@ def _connect(self):
876
881
877
882
# set the socket_connect_timeout before we connect
878
883
if self .tmp_relax_timeout != - 1 :
879
- logging .debug (
880
- f"***** Connection connect --> Using relax_timeout: { self .tmp_relax_timeout } "
881
- )
882
884
sock .settimeout (self .tmp_relax_timeout )
883
885
else :
884
- logging .debug (
885
- f"***** Connection connect --> Using default socket_connect_timeout: { self .socket_connect_timeout } "
886
- )
887
886
sock .settimeout (self .socket_connect_timeout )
888
887
889
888
# connect
890
889
sock .connect (socket_address )
891
890
892
891
# set the socket_timeout now that we're connected
893
892
if self .tmp_relax_timeout != - 1 :
894
- logging .debug (
895
- f"***** Connection --> Using relax_timeout: { self .tmp_relax_timeout } "
896
- )
897
893
sock .settimeout (self .tmp_relax_timeout )
898
894
else :
899
- logging .debug (
900
- f"***** Connection --> Using default socket_timeout: { self .socket_timeout } "
901
- )
902
895
sock .settimeout (self .socket_timeout )
903
- logging .debug (f"Connected to { sock .getpeername ()} " )
904
896
return sock
905
897
906
898
except OSError as _ :
@@ -1599,14 +1591,10 @@ def _update_maintenance_events_configs_for_connections(
1599
1591
):
1600
1592
with self ._lock :
1601
1593
for conn in self ._available_connections :
1602
- conn .set_maintenance_events_pool_handler (
1603
- maintenance_events_pool_handler
1604
- )
1594
+ conn .set_maintenance_event_pool_handler (maintenance_events_pool_handler )
1605
1595
conn .maintenance_events_config = maintenance_events_pool_handler .config
1606
1596
for conn in self ._in_use_connections :
1607
- conn .set_maintenance_events_pool_handler (
1608
- maintenance_events_pool_handler
1609
- )
1597
+ conn .set_maintenance_event_pool_handler (maintenance_events_pool_handler )
1610
1598
conn .maintenance_events_config = maintenance_events_pool_handler .config
1611
1599
1612
1600
def reset (self ) -> None :
@@ -1748,9 +1736,6 @@ def release(self, connection: "Connection") -> None:
1748
1736
1749
1737
if self .owns_connection (connection ):
1750
1738
if connection .should_reconnect ():
1751
- logging .debug (
1752
- f"***** Pool--> disconnecting in release { connection ._sock .getpeername ()} "
1753
- )
1754
1739
connection .disconnect ()
1755
1740
self ._available_connections .append (connection )
1756
1741
self ._event_dispatcher .dispatch (
@@ -1910,21 +1895,13 @@ def update_connections_current_timeout(
1910
1895
If -1 is provided - the relax timeout is disabled.
1911
1896
:param include_available_connections: Whether to include available connections in the update.
1912
1897
"""
1913
- logging .debug (f"***** Pool --> Updating timeouts. New value: { relax_timeout } " )
1914
- start_time = time .time ()
1915
-
1916
1898
for conn in self ._in_use_connections :
1917
1899
self ._update_connection_timeout (conn , relax_timeout )
1918
1900
1919
1901
if include_free_connections :
1920
1902
for conn in self ._available_connections :
1921
1903
self ._update_connection_timeout (conn , relax_timeout )
1922
1904
1923
- execution_time_us = (time .time () - start_time ) * 1000000
1924
- logging .error (
1925
- f"###### TIMEOUTS execution time: { execution_time_us :.0f} microseconds"
1926
- )
1927
-
1928
1905
def _update_connection_for_reconnect (
1929
1906
self ,
1930
1907
connection : "Connection" ,
@@ -2014,6 +1991,8 @@ def __init__(
2014
1991
):
2015
1992
self .queue_class = queue_class
2016
1993
self .timeout = timeout
1994
+ self ._in_maintenance = False
1995
+ self ._locked = False
2017
1996
super ().__init__ (
2018
1997
connection_class = connection_class ,
2019
1998
max_connections = max_connections ,
@@ -2022,7 +2001,10 @@ def __init__(
2022
2001
2023
2002
def reset (self ):
2024
2003
# Create and fill up a thread safe queue with ``None`` values.
2025
- with self ._lock :
2004
+ try :
2005
+ if self ._in_maintenance :
2006
+ self ._lock .acquire ()
2007
+ self ._locked = True
2026
2008
self .pool = self .queue_class (self .max_connections )
2027
2009
while True :
2028
2010
try :
@@ -2033,6 +2015,13 @@ def reset(self):
2033
2015
# Keep a list of actual connection instances so that we can
2034
2016
# disconnect them later.
2035
2017
self ._connections = []
2018
+ finally :
2019
+ if self ._locked :
2020
+ try :
2021
+ self ._lock .release ()
2022
+ except Exception :
2023
+ pass
2024
+ self ._locked = False
2036
2025
2037
2026
# this must be the last operation in this method. while reset() is
2038
2027
# called when holding _fork_lock, other threads in this process
@@ -2047,7 +2036,10 @@ def reset(self):
2047
2036
2048
2037
def make_connection (self ):
2049
2038
"Make a fresh connection."
2050
- with self ._lock :
2039
+ try :
2040
+ if self ._in_maintenance :
2041
+ self ._lock .acquire ()
2042
+ self ._locked = True
2051
2043
if self .cache is not None :
2052
2044
connection = CacheProxyConnection (
2053
2045
self .connection_class (** self .connection_kwargs ),
@@ -2059,6 +2051,13 @@ def make_connection(self):
2059
2051
2060
2052
self ._connections .append (connection )
2061
2053
return connection
2054
+ finally :
2055
+ if self ._locked :
2056
+ try :
2057
+ self ._lock .release ()
2058
+ except Exception :
2059
+ pass
2060
+ self ._locked = False
2062
2061
2063
2062
@deprecated_args (
2064
2063
args_to_warn = ["*" ],
@@ -2083,7 +2082,10 @@ def get_connection(self, command_name=None, *keys, **options):
2083
2082
# Try and get a connection from the pool. If one isn't available within
2084
2083
# self.timeout then raise a ``ConnectionError``.
2085
2084
connection = None
2086
- with self ._lock :
2085
+ try :
2086
+ if self ._in_maintenance :
2087
+ self ._lock .acquire ()
2088
+ self ._locked = True
2087
2089
try :
2088
2090
connection = self .pool .get (block = True , timeout = self .timeout )
2089
2091
except Empty :
@@ -2095,6 +2097,13 @@ def get_connection(self, command_name=None, *keys, **options):
2095
2097
# a new connection to add to the pool.
2096
2098
if connection is None :
2097
2099
connection = self .make_connection ()
2100
+ finally :
2101
+ if self ._locked :
2102
+ try :
2103
+ self ._lock .release ()
2104
+ except Exception :
2105
+ pass
2106
+ self ._locked = False
2098
2107
2099
2108
try :
2100
2109
# ensure this connection is connected to Redis
@@ -2123,7 +2132,10 @@ def release(self, connection):
2123
2132
# Make sure we haven't changed process.
2124
2133
self ._checkpid ()
2125
2134
2126
- with self ._lock :
2135
+ try :
2136
+ if self ._in_maintenance :
2137
+ self ._lock .acquire ()
2138
+ self ._locked = True
2127
2139
if not self .owns_connection (connection ):
2128
2140
# pool doesn't own this connection. do not add it back
2129
2141
# to the pool. instead add a None value which is a placeholder
@@ -2133,24 +2145,39 @@ def release(self, connection):
2133
2145
self .pool .put_nowait (None )
2134
2146
return
2135
2147
if connection .should_reconnect ():
2136
- logging .debug (
2137
- f"***** Blocking Pool--> disconnecting in release { connection ._sock .getpeername ()} "
2138
- )
2139
2148
connection .disconnect ()
2140
2149
# Put the connection back into the pool.
2141
2150
try :
2151
+ print ("Releasing connection - in the pool" )
2142
2152
self .pool .put_nowait (connection )
2143
2153
except Full :
2144
2154
# perhaps the pool has been reset() after a fork? regardless,
2145
2155
# we don't want this connection
2146
2156
pass
2157
+ finally :
2158
+ if self ._locked :
2159
+ try :
2160
+ self ._lock .release ()
2161
+ except Exception :
2162
+ pass
2163
+ self ._locked = False
2147
2164
2148
2165
def disconnect (self ):
2149
2166
"Disconnects all connections in the pool."
2150
2167
self ._checkpid ()
2151
- with self ._lock :
2168
+ try :
2169
+ if self ._in_maintenance :
2170
+ self ._lock .acquire ()
2171
+ self ._locked = True
2152
2172
for connection in self ._connections :
2153
2173
connection .disconnect ()
2174
+ finally :
2175
+ if self ._locked :
2176
+ try :
2177
+ self ._lock .release ()
2178
+ except Exception :
2179
+ pass
2180
+ self ._locked = False
2154
2181
2155
2182
def update_active_connections_for_reconnect (
2156
2183
self ,
@@ -2229,3 +2256,7 @@ def _update_maintenance_events_configs_for_connections(
2229
2256
conn .maintenance_events_config = (
2230
2257
maintenance_events_pool_handler .config
2231
2258
)
2259
+
2260
+ def set_in_maintenance (self , in_maintenance : bool ):
2261
+ """Set the maintenance mode for the connection pool."""
2262
+ self ._in_maintenance = in_maintenance
0 commit comments