1 files changed, 59 insertions, 78 deletions

diff --git a/bitbake/lib/bb/asyncrpc/client.py b/bitbake/lib/bb/asyncrpc/client.py
index 29a5ab76aa..17b72033b9 100644
--- a/bitbake/lib/bb/asyncrpc/client.py
+++ b/bitbake/lib/bb/asyncrpc/client.py
@@ -10,11 +10,41 @@ import json
 import os
 import socket
 import sys
+import re
 import contextlib
 from threading import Thread
 from .connection import StreamConnection, WebsocketConnection, DEFAULT_MAX_CHUNK
 from .exceptions import ConnectionClosedError, InvokeError
 
+UNIX_PREFIX = "unix://"
+WS_PREFIX = "ws://"
+WSS_PREFIX = "wss://"
+
+ADDR_TYPE_UNIX = 0
+ADDR_TYPE_TCP = 1
+ADDR_TYPE_WS = 2
+
+WEBSOCKETS_MIN_VERSION = (9, 1)
+# Need websockets 10 with python 3.10+
+if sys.version_info >= (3, 10, 0):
+    WEBSOCKETS_MIN_VERSION = (10, 0)
+
+
+def parse_address(addr):
+    if addr.startswith(UNIX_PREFIX):
+        return (ADDR_TYPE_UNIX, (addr[len(UNIX_PREFIX) :],))
+    elif addr.startswith(WS_PREFIX) or addr.startswith(WSS_PREFIX):
+        return (ADDR_TYPE_WS, (addr,))
+    else:
+        m = re.match(r"\[(?P<host>[^\]]*)\]:(?P<port>\d+)$", addr)
+        if m is not None:
+            host = m.group("host")
+            port = m.group("port")
+        else:
+            host, port = addr.split(":")
+
+        return (ADDR_TYPE_TCP, (host, int(port)))
+
 
 class AsyncClient(object):
     def __init__(
@@ -63,8 +93,35 @@ class AsyncClient(object):
     async def connect_websocket(self, uri):
         import websockets
 
+        try:
+            version = tuple(
+                int(v)
+                for v in websockets.__version__.split(".")[
+                    0 : len(WEBSOCKETS_MIN_VERSION)
+                ]
+            )
+        except ValueError:
+            raise ImportError(
+                f"Unable to parse websockets version '{websockets.__version__}'"
+            )
+
+        if version < WEBSOCKETS_MIN_VERSION:
+            min_ver_str = ".".join(str(v) for v in WEBSOCKETS_MIN_VERSION)
+            raise ImportError(
+                f"Websockets version {websockets.__version__} is less than minimum required version {min_ver_str}"
+            )
+
         async def connect_sock():
-            websocket = await websockets.connect(uri, ping_interval=None)
+            try:
+                websocket = await websockets.connect(
+                    uri,
+                    ping_interval=None,
+                    open_timeout=self.timeout,
+                )
+            except asyncio.exceptions.TimeoutError:
+                raise ConnectionError("Timeout while connecting to websocket")
+            except (OSError, websockets.InvalidHandshake, websockets.InvalidURI) as exc:
+                raise ConnectionError(f"Could not connect to websocket: {exc}") from exc
             return WebsocketConnection(websocket, self.timeout)
 
         self._connect_sock = connect_sock
@@ -202,85 +259,9 @@ class Client(object):
     def close(self):
         if self.loop:
             self.loop.run_until_complete(self.client.close())
-            if sys.version_info >= (3, 6):
-                self.loop.run_until_complete(self.loop.shutdown_asyncgens())
-            self.loop.close()
-        self.loop = None
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
-        return False
-
-
-class ClientPool(object):
-    def __init__(self, max_clients):
-        self.avail_clients = []
-        self.num_clients = 0
-        self.max_clients = max_clients
-        self.loop = None
-        self.client_condition = None
-
-    @abc.abstractmethod
-    async def _new_client(self):
-        raise NotImplementedError("Must be implemented in derived class")
-
-    def close(self):
-        if self.client_condition:
-            self.client_condition = None
-
-        if self.loop:
-            self.loop.run_until_complete(self.__close_clients())
             self.loop.run_until_complete(self.loop.shutdown_asyncgens())
             self.loop.close()
-            self.loop = None
-
-    def run_tasks(self, tasks):
-        if not self.loop:
-            self.loop = asyncio.new_event_loop()
-
-        thread = Thread(target=self.__thread_main, args=(tasks,))
-        thread.start()
-        thread.join()
-
-    @contextlib.asynccontextmanager
-    async def get_client(self):
-        async with self.client_condition:
-            if self.avail_clients:
-                client = self.avail_clients.pop()
-            elif self.num_clients < self.max_clients:
-                self.num_clients += 1
-                client = await self._new_client()
-            else:
-                while not self.avail_clients:
-                    await self.client_condition.wait()
-                client = self.avail_clients.pop()
-
-        try:
-            yield client
-        finally:
-            async with self.client_condition:
-                self.avail_clients.append(client)
-                self.client_condition.notify()
-
-    def __thread_main(self, tasks):
-        async def process_task(task):
-            async with self.get_client() as client:
-                await task(client)
-
-        asyncio.set_event_loop(self.loop)
-        if not self.client_condition:
-            self.client_condition = asyncio.Condition()
-        tasks = [process_task(t) for t in tasks]
-        self.loop.run_until_complete(asyncio.gather(*tasks))
-
-    async def __close_clients(self):
-        for c in self.avail_clients:
-            await c.close()
-        self.avail_clients = []
-        self.num_clients = 0
+        self.loop = None
 
     def __enter__(self):
         return self