1 files changed, 337 insertions, 0 deletions

diff --git a/bitbake/lib/bb/asyncrpc/client.py b/bitbake/lib/bb/asyncrpc/client.py
new file mode 100644
index 0000000000..b49de99313
--- /dev/null
+++ b/bitbake/lib/bb/asyncrpc/client.py
@@ -0,0 +1,337 @@
+#
+# Copyright BitBake Contributors
+#
+# SPDX-License-Identifier: GPL-2.0-only
+#
+
+import abc
+import asyncio
+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__(
+        self,
+        proto_name,
+        proto_version,
+        logger,
+        timeout=30,
+        server_headers=False,
+        headers={},
+    ):
+        self.socket = None
+        self.max_chunk = DEFAULT_MAX_CHUNK
+        self.proto_name = proto_name
+        self.proto_version = proto_version
+        self.logger = logger
+        self.timeout = timeout
+        self.needs_server_headers = server_headers
+        self.server_headers = {}
+        self.headers = headers
+
+    async def connect_tcp(self, address, port):
+        async def connect_sock():
+            reader, writer = await asyncio.open_connection(address, port)
+            return StreamConnection(reader, writer, self.timeout, self.max_chunk)
+
+        self._connect_sock = connect_sock
+
+    async def connect_unix(self, path):
+        async def connect_sock():
+            # AF_UNIX has path length issues so chdir here to workaround
+            cwd = os.getcwd()
+            try:
+                os.chdir(os.path.dirname(path))
+                # The socket must be opened synchronously so that CWD doesn't get
+                # changed out from underneath us so we pass as a sock into asyncio
+                sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, 0)
+                sock.connect(os.path.basename(path))
+            finally:
+                os.chdir(cwd)
+            reader, writer = await asyncio.open_unix_connection(sock=sock)
+            return StreamConnection(reader, writer, self.timeout, self.max_chunk)
+
+        self._connect_sock = connect_sock
+
+    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)
+            return WebsocketConnection(websocket, self.timeout)
+
+        self._connect_sock = connect_sock
+
+    async def setup_connection(self):
+        # Send headers
+        await self.socket.send("%s %s" % (self.proto_name, self.proto_version))
+        await self.socket.send(
+            "needs-headers: %s" % ("true" if self.needs_server_headers else "false")
+        )
+        for k, v in self.headers.items():
+            await self.socket.send("%s: %s" % (k, v))
+
+        # End of headers
+        await self.socket.send("")
+
+        self.server_headers = {}
+        if self.needs_server_headers:
+            while True:
+                line = await self.socket.recv()
+                if not line:
+                    # End headers
+                    break
+                tag, value = line.split(":", 1)
+                self.server_headers[tag.lower()] = value.strip()
+
+    async def get_header(self, tag, default):
+        await self.connect()
+        return self.server_headers.get(tag, default)
+
+    async def connect(self):
+        if self.socket is None:
+            self.socket = await self._connect_sock()
+            await self.setup_connection()
+
+    async def disconnect(self):
+        if self.socket is not None:
+            await self.socket.close()
+            self.socket = None
+
+    async def close(self):
+        await self.disconnect()
+
+    async def _send_wrapper(self, proc):
+        count = 0
+        while True:
+            try:
+                await self.connect()
+                return await proc()
+            except (
+                OSError,
+                ConnectionError,
+                ConnectionClosedError,
+                json.JSONDecodeError,
+                UnicodeDecodeError,
+            ) as e:
+                self.logger.warning("Error talking to server: %s" % e)
+                if count >= 3:
+                    if not isinstance(e, ConnectionError):
+                        raise ConnectionError(str(e))
+                    raise e
+                await self.close()
+                count += 1
+
+    def check_invoke_error(self, msg):
+        if isinstance(msg, dict) and "invoke-error" in msg:
+            raise InvokeError(msg["invoke-error"]["message"])
+
+    async def invoke(self, msg):
+        async def proc():
+            await self.socket.send_message(msg)
+            return await self.socket.recv_message()
+
+        result = await self._send_wrapper(proc)
+        self.check_invoke_error(result)
+        return result
+
+    async def ping(self):
+        return await self.invoke({"ping": {}})
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.close()
+
+
+class Client(object):
+    def __init__(self):
+        self.client = self._get_async_client()
+        self.loop = asyncio.new_event_loop()
+
+        # Override any pre-existing loop.
+        # Without this, the PR server export selftest triggers a hang
+        # when running with Python 3.7.  The drawback is that there is
+        # potential for issues if the PR and hash equiv (or some new)
+        # clients need to both be instantiated in the same process.
+        # This should be revisited if/when Python 3.9 becomes the
+        # minimum required version for BitBake, as it seems not
+        # required (but harmless) with it.
+        asyncio.set_event_loop(self.loop)
+
+        self._add_methods("connect_tcp", "ping")
+
+    @abc.abstractmethod
+    def _get_async_client(self):
+        pass
+
+    def _get_downcall_wrapper(self, downcall):
+        def wrapper(*args, **kwargs):
+            return self.loop.run_until_complete(downcall(*args, **kwargs))
+
+        return wrapper
+
+    def _add_methods(self, *methods):
+        for m in methods:
+            downcall = getattr(self.client, m)
+            setattr(self, m, self._get_downcall_wrapper(downcall))
+
+    def connect_unix(self, path):
+        self.loop.run_until_complete(self.client.connect_unix(path))
+        self.loop.run_until_complete(self.client.connect())
+
+    @property
+    def max_chunk(self):
+        return self.client.max_chunk
+
+    @max_chunk.setter
+    def max_chunk(self, value):
+        self.client.max_chunk = value
+
+    def disconnect(self):
+        self.loop.run_until_complete(self.client.close())
+
+    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
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False