3 files changed, 88 insertions, 88 deletions

diff --git a/bitbake/lib/bb/asyncrpc/__init__.py b/bitbake/lib/bb/asyncrpc/__init__.py
index 639e1607f8..a4371643d7 100644
--- a/bitbake/lib/bb/asyncrpc/__init__.py
+++ b/bitbake/lib/bb/asyncrpc/__init__.py
@@ -5,7 +5,7 @@
 #
 
 
-from .client import AsyncClient, Client, ClientPool
+from .client import AsyncClient, Client
 from .serv import AsyncServer, AsyncServerConnection
 from .connection import DEFAULT_MAX_CHUNK
 from .exceptions import (
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
diff --git a/bitbake/lib/bb/asyncrpc/serv.py b/bitbake/lib/bb/asyncrpc/serv.py
index a66117acad..667217c5c1 100644
--- a/bitbake/lib/bb/asyncrpc/serv.py
+++ b/bitbake/lib/bb/asyncrpc/serv.py
@@ -138,14 +138,20 @@ class StreamServer(object):
 
 
 class TCPStreamServer(StreamServer):
-    def __init__(self, host, port, handler, logger):
+    def __init__(self, host, port, handler, logger, *, reuseport=False):
         super().__init__(handler, logger)
         self.host = host
         self.port = port
+        self.reuseport = reuseport
 
     def start(self, loop):
         self.server = loop.run_until_complete(
-            asyncio.start_server(self.handle_stream_client, self.host, self.port)
+            asyncio.start_server(
+                self.handle_stream_client,
+                self.host,
+                self.port,
+                reuse_port=self.reuseport,
+            )
         )
 
         for s in self.server.sockets:
@@ -209,11 +215,12 @@ class UnixStreamServer(StreamServer):
 
 
 class WebsocketsServer(object):
-    def __init__(self, host, port, handler, logger):
+    def __init__(self, host, port, handler, logger, *, reuseport=False):
         self.host = host
         self.port = port
         self.handler = handler
         self.logger = logger
+        self.reuseport = reuseport
 
     def start(self, loop):
         import websockets.server
@@ -224,6 +231,7 @@ class WebsocketsServer(object):
                 self.host,
                 self.port,
                 ping_interval=None,
+                reuse_port=self.reuseport,
             )
         )
 
@@ -262,14 +270,26 @@ class AsyncServer(object):
         self.loop = None
         self.run_tasks = []
 
-    def start_tcp_server(self, host, port):
-        self.server = TCPStreamServer(host, port, self._client_handler, self.logger)
+    def start_tcp_server(self, host, port, *, reuseport=False):
+        self.server = TCPStreamServer(
+            host,
+            port,
+            self._client_handler,
+            self.logger,
+            reuseport=reuseport,
+        )
 
     def start_unix_server(self, path):
         self.server = UnixStreamServer(path, self._client_handler, self.logger)
 
-    def start_websocket_server(self, host, port):
-        self.server = WebsocketsServer(host, port, self._client_handler, self.logger)
+    def start_websocket_server(self, host, port, reuseport=False):
+        self.server = WebsocketsServer(
+            host,
+            port,
+            self._client_handler,
+            self.logger,
+            reuseport=reuseport,
+        )
 
     async def _client_handler(self, socket):
         address = socket.address
@@ -368,8 +388,7 @@ class AsyncServer(object):
 
             self._serve_forever(tasks)
 
-            if sys.version_info >= (3, 6):
-                self.loop.run_until_complete(self.loop.shutdown_asyncgens())
+            self.loop.run_until_complete(self.loop.shutdown_asyncgens())
             self.loop.close()
 
         queue = multiprocessing.Queue()