@@ -184,8 +184,8 @@ def __init__(
184184 self ._socket_pool = socket_pool
185185 # Hang onto open sockets so that we can reuse them.
186186 self ._available_sockets = set ()
187+ self ._key_by_managed_socket = {}
187188 self ._managed_socket_by_key = {}
188- self ._managed_socket_by_socket = {}
189189
190190 def _free_sockets (self , force : bool = False ) -> None :
191191 # cloning lists since items are being removed
@@ -241,7 +241,7 @@ def close_socket(self, socket: SocketType) -> None:
241241 if socket not in self ._managed_socket_by_key .values ():
242242 raise RuntimeError ("Socket not managed" )
243243 socket .close ()
244- key = self ._managed_socket_by_socket .pop (socket )
244+ key = self ._key_by_managed_socket .pop (socket )
245245 del self ._managed_socket_by_key [key ]
246246 if socket in self ._available_sockets :
247247 self ._available_sockets .remove (socket )
@@ -299,8 +299,8 @@ def get_socket(
299299 if isinstance (result , Exception ):
300300 raise RuntimeError (f"Error connecting socket: { result } " ) from result
301301
302+ self ._key_by_managed_socket [result ] = key
302303 self ._managed_socket_by_key [key ] = result
303- self ._managed_socket_by_socket [result ] = key
304304 return result
305305
306306
0 commit comments