@@ -184,8 +184,8 @@ def __init__(
184
184
self ._socket_pool = socket_pool
185
185
# Hang onto open sockets so that we can reuse them.
186
186
self ._available_sockets = set ()
187
+ self ._key_by_managed_socket = {}
187
188
self ._managed_socket_by_key = {}
188
- self ._managed_socket_by_socket = {}
189
189
190
190
def _free_sockets (self , force : bool = False ) -> None :
191
191
# cloning lists since items are being removed
@@ -241,7 +241,7 @@ def close_socket(self, socket: SocketType) -> None:
241
241
if socket not in self ._managed_socket_by_key .values ():
242
242
raise RuntimeError ("Socket not managed" )
243
243
socket .close ()
244
- key = self ._managed_socket_by_socket .pop (socket )
244
+ key = self ._key_by_managed_socket .pop (socket )
245
245
del self ._managed_socket_by_key [key ]
246
246
if socket in self ._available_sockets :
247
247
self ._available_sockets .remove (socket )
@@ -299,8 +299,8 @@ def get_socket(
299
299
if isinstance (result , Exception ):
300
300
raise RuntimeError (f"Error connecting socket: { result } " ) from result
301
301
302
+ self ._key_by_managed_socket [result ] = key
302
303
self ._managed_socket_by_key [key ] = result
303
- self ._managed_socket_by_socket [result ] = key
304
304
return result
305
305
306
306
0 commit comments