Fixed client.wait_connected() concurrency issue (#171)

francis-clairicia · Nov 26, 2023 · 169ebe1 · 169ebe1
1 parent 012c629
commit 169ebe1
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 16 deletions.
diff --git a/src/easynetwork/api_async/client/tcp.py b/src/easynetwork/api_async/client/tcp.py
@@ -488,23 +488,23 @@ async def __ensure_connected(self) -> AsyncStreamEndpoint[_SentPacketT, _Receive
                 endpoint_and_proxy = await socket_connector.get()
             self.__socket_connector = None
             if endpoint_and_proxy is None:
-                raise ClientClosedError("Client is closing, or is already closed")
-            transport, self.__socket_proxy = endpoint_and_proxy
-            del endpoint_and_proxy
-            self.__endpoint = AsyncStreamEndpoint(transport, self.__protocol, max_recv_size=self.max_recv_size)
+                raise self.__closed()
+            if self.__endpoint is None:
+                transport, self.__socket_proxy = endpoint_and_proxy
+                self.__endpoint = AsyncStreamEndpoint(transport, self.__protocol, max_recv_size=self.max_recv_size)
 
         if self.__endpoint.is_closing():
-            raise ClientClosedError("Client is closing, or is already closed")
+            raise self.__closed()
         return self.__endpoint
 
     def __get_endpoint_sync(self) -> AsyncStreamEndpoint[_SentPacketT, _ReceivedPacketT]:
         if self.__endpoint is None:
             if self.__socket_connector is not None:
                 raise _utils.error_from_errno(_errno.ENOTSOCK)
             else:
-                raise ClientClosedError("Client is closing, or is already closed")
+                raise self.__closed()
         if self.__endpoint.is_closing():
-            raise ClientClosedError("Client is closing, or is already closed")
+            raise self.__closed()
         return self.__endpoint
 
     @contextlib.contextmanager
@@ -515,13 +515,17 @@ def __convert_socket_error(self) -> Iterator[None]:
             raise self.__abort() from exc
         except OSError as exc:
             if exc.errno in constants.CLOSED_SOCKET_ERRNOS:
-                raise ClientClosedError("Client is closing, or is already closed")
+                raise self.__closed()
             raise
 
     @staticmethod
     def __abort() -> OSError:
         return _utils.error_from_errno(_errno.ECONNABORTED)
 
+    @staticmethod
+    def __closed() -> ClientClosedError:
+        return ClientClosedError("Client is closing, or is already closed")
+
     @property
     @final
     def socket(self) -> SocketProxy:

diff --git a/src/easynetwork/api_async/client/udp.py b/src/easynetwork/api_async/client/udp.py
@@ -331,23 +331,23 @@ async def __ensure_connected(self) -> AsyncDatagramEndpoint[_SentPacketT, _Recei
                 endpoint_and_proxy = await socket_connector.get()
             self.__socket_connector = None
             if endpoint_and_proxy is None:
-                raise ClientClosedError("Client is closing, or is already closed")
-            transport, self.__socket_proxy = endpoint_and_proxy
-            del endpoint_and_proxy
-            self.__endpoint = AsyncDatagramEndpoint(transport, self.__protocol)
+                raise self.__closed()
+            if self.__endpoint is None:
+                transport, self.__socket_proxy = endpoint_and_proxy
+                self.__endpoint = AsyncDatagramEndpoint(transport, self.__protocol)
 
         if self.__endpoint.is_closing():
-            raise ClientClosedError("Client is closing, or is already closed")
+            raise self.__closed()
         return self.__endpoint
 
     def __get_endpoint_sync(self) -> AsyncDatagramEndpoint[_SentPacketT, _ReceivedPacketT]:
         if self.__endpoint is None:
             if self.__socket_connector is not None:
                 raise _utils.error_from_errno(_errno.ENOTSOCK)
             else:
-                raise ClientClosedError("Client is closing, or is already closed")
+                raise self.__closed()
         if self.__endpoint.is_closing():
-            raise ClientClosedError("Client is closing, or is already closed")
+            raise self.__closed()
         return self.__endpoint
 
     @contextlib.contextmanager
@@ -356,9 +356,13 @@ def __convert_socket_error(self) -> Iterator[None]:
             yield
         except OSError as exc:
             if exc.errno in constants.CLOSED_SOCKET_ERRNOS:
-                raise ClientClosedError("Client is closing, or is already closed")
+                raise self.__closed()
             raise
 
+    @staticmethod
+    def __closed() -> ClientClosedError:
+        return ClientClosedError("Client is closing, or is already closed")
+
     @property
     @final
     def socket(self) -> SocketProxy: