bitbake: asyncrpc: Abstract sockets

Rewrites the asyncrpc client and server code to make it possible to have
other transport backends that are not stream based (e.g. websockets
which are message based). The connection handling classes are now shared
between both the client and server to make it easier to implement new
transport mechanisms

(Bitbake rev: 2aaeae53696e4c2f13a169830c3b7089cbad6eca)

Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>

1 files changed, 156 insertions, 148 deletions

diff --git a/bitbake/lib/bb/asyncrpc/serv.py b/bitbake/lib/bb/asyncrpc/serv.py
index d2de4891b8..3e0d0632cb 100644
--- a/bitbake/lib/bb/asyncrpc/serv.py
+++ b/bitbake/lib/bb/asyncrpc/serv.py
@@ -12,241 +12,248 @@ import signal
 import socket
 import sys
 import multiprocessing
-from . import chunkify, DEFAULT_MAX_CHUNK
-
-
-class ClientError(Exception):
-    pass
-
-
-class ServerError(Exception):
-    pass
+from .connection import StreamConnection
+from .exceptions import ClientError, ServerError, ConnectionClosedError
 
 
 class AsyncServerConnection(object):
-    def __init__(self, reader, writer, proto_name, logger):
-        self.reader = reader
-        self.writer = writer
+    # If a handler returns this object (e.g. `return self.NO_RESPONSE`), no
+    # return message will be automatically be sent back to the client
+    NO_RESPONSE = object()
+
+    def __init__(self, socket, proto_name, logger):
+        self.socket = socket
         self.proto_name = proto_name
-        self.max_chunk = DEFAULT_MAX_CHUNK
         self.handlers = {
-            'chunk-stream': self.handle_chunk,
-            'ping': self.handle_ping,
+            "ping": self.handle_ping,
         }
         self.logger = logger
 
+    async def close(self):
+        await self.socket.close()
+
     async def process_requests(self):
         try:
-            self.addr = self.writer.get_extra_info('peername')
-            self.logger.debug('Client %r connected' % (self.addr,))
+            self.logger.info("Client %r connected" % (self.socket.address,))
 
             # Read protocol and version
-            client_protocol = await self.reader.readline()
+            client_protocol = await self.socket.recv()
             if not client_protocol:
                 return
 
-            (client_proto_name, client_proto_version) = client_protocol.decode('utf-8').rstrip().split()
+            (client_proto_name, client_proto_version) = client_protocol.split()
             if client_proto_name != self.proto_name:
-                self.logger.debug('Rejecting invalid protocol %s' % (self.proto_name))
+                self.logger.debug("Rejecting invalid protocol %s" % (self.proto_name))
                 return
 
-            self.proto_version = tuple(int(v) for v in client_proto_version.split('.'))
+            self.proto_version = tuple(int(v) for v in client_proto_version.split("."))
             if not self.validate_proto_version():
-                self.logger.debug('Rejecting invalid protocol version %s' % (client_proto_version))
+                self.logger.debug(
+                    "Rejecting invalid protocol version %s" % (client_proto_version)
+                )
                 return
 
             # Read headers. Currently, no headers are implemented, so look for
             # an empty line to signal the end of the headers
             while True:
-                line = await self.reader.readline()
-                if not line:
-                    return
-
-                line = line.decode('utf-8').rstrip()
-                if not line:
+                header = await self.socket.recv()
+                if not header:
                     break
 
             # Handle messages
             while True:
-                d = await self.read_message()
+                d = await self.socket.recv_message()
                 if d is None:
                     break
-                await self.dispatch_message(d)
-                await self.writer.drain()
-        except ClientError as e:
+                response = await self.dispatch_message(d)
+                if response is not self.NO_RESPONSE:
+                    await self.socket.send_message(response)
+
+        except ConnectionClosedError as e:
+            self.logger.info(str(e))
+        except (ClientError, ConnectionError) as e:
             self.logger.error(str(e))
         finally:
-            self.writer.close()
+            await self.close()
 
     async def dispatch_message(self, msg):
         for k in self.handlers.keys():
             if k in msg:
-                self.logger.debug('Handling %s' % k)
-                await self.handlers[k](msg[k])
-                return
+                self.logger.debug("Handling %s" % k)
+                return await self.handlers[k](msg[k])
 
         raise ClientError("Unrecognized command %r" % msg)
 
-    def write_message(self, msg):
-        for c in chunkify(json.dumps(msg), self.max_chunk):
-            self.writer.write(c.encode('utf-8'))
+    async def handle_ping(self, request):
+        return {"alive": True}
 
-    async def read_message(self):
-        l = await self.reader.readline()
-        if not l:
-            return None
 
-        try:
-            message = l.decode('utf-8')
+class StreamServer(object):
+    def __init__(self, handler, logger):
+        self.handler = handler
+        self.logger = logger
+        self.closed = False
 
-            if not message.endswith('\n'):
-                return None
+    async def handle_stream_client(self, reader, writer):
+        # writer.transport.set_write_buffer_limits(0)
+        socket = StreamConnection(reader, writer, -1)
+        if self.closed:
+            await socket.close()
+            return
+
+        await self.handler(socket)
+
+    async def stop(self):
+        self.closed = True
+
+
+class TCPStreamServer(StreamServer):
+    def __init__(self, host, port, handler, logger):
+        super().__init__(handler, logger)
+        self.host = host
+        self.port = port
+
+    def start(self, loop):
+        self.server = loop.run_until_complete(
+            asyncio.start_server(self.handle_stream_client, self.host, self.port)
+        )
+
+        for s in self.server.sockets:
+            self.logger.debug("Listening on %r" % (s.getsockname(),))
+            # Newer python does this automatically. Do it manually here for
+            # maximum compatibility
+            s.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
+            s.setsockopt(socket.SOL_TCP, socket.TCP_QUICKACK, 1)
+
+            # Enable keep alives. This prevents broken client connections
+            # from persisting on the server for long periods of time.
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+            s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30)
+            s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 15)
+            s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 4)
+
+        name = self.server.sockets[0].getsockname()
+        if self.server.sockets[0].family == socket.AF_INET6:
+            self.address = "[%s]:%d" % (name[0], name[1])
+        else:
+            self.address = "%s:%d" % (name[0], name[1])
+
+        return [self.server.wait_closed()]
+
+    async def stop(self):
+        await super().stop()
+        self.server.close()
+
+    def cleanup(self):
+        pass
 
-            return json.loads(message)
-        except (json.JSONDecodeError, UnicodeDecodeError) as e:
-            self.logger.error('Bad message from client: %r' % message)
-            raise e
 
-    async def handle_chunk(self, request):
-        lines = []
-        try:
-            while True:
-                l = await self.reader.readline()
-                l = l.rstrip(b"\n").decode("utf-8")
-                if not l:
-                    break
-                lines.append(l)
+class UnixStreamServer(StreamServer):
+    def __init__(self, path, handler, logger):
+        super().__init__(handler, logger)
+        self.path = path
 
-            msg = json.loads(''.join(lines))
-        except (json.JSONDecodeError, UnicodeDecodeError) as e:
-            self.logger.error('Bad message from client: %r' % lines)
-            raise e
+    def start(self, loop):
+        cwd = os.getcwd()
+        try:
+            # Work around path length limits in AF_UNIX
+            os.chdir(os.path.dirname(self.path))
+            self.server = loop.run_until_complete(
+                asyncio.start_unix_server(
+                    self.handle_stream_client, os.path.basename(self.path)
+                )
+            )
+        finally:
+            os.chdir(cwd)
 
-        if 'chunk-stream' in msg:
-            raise ClientError("Nested chunks are not allowed")
+        self.logger.debug("Listening on %r" % self.path)
+        self.address = "unix://%s" % os.path.abspath(self.path)
+        return [self.server.wait_closed()]
 
-        await self.dispatch_message(msg)
+    async def stop(self):
+        await super().stop()
+        self.server.close()
 
-    async def handle_ping(self, request):
-        response = {'alive': True}
-        self.write_message(response)
+    def cleanup(self):
+        os.unlink(self.path)
 
 
 class AsyncServer(object):
     def __init__(self, logger):
-        self._cleanup_socket = None
         self.logger = logger
-        self.start = None
-        self.address = None
         self.loop = None
+        self.run_tasks = []
 
     def start_tcp_server(self, host, port):
-        def start_tcp():
-            self.server = self.loop.run_until_complete(
-                asyncio.start_server(self.handle_client, host, port)
-            )
-
-            for s in self.server.sockets:
-                self.logger.debug('Listening on %r' % (s.getsockname(),))
-                # Newer python does this automatically. Do it manually here for
-                # maximum compatibility
-                s.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1)
-                s.setsockopt(socket.SOL_TCP, socket.TCP_QUICKACK, 1)
-
-                # Enable keep alives. This prevents broken client connections
-                # from persisting on the server for long periods of time.
-                s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
-                s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30)
-                s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 15)
-                s.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 4)
-
-            name = self.server.sockets[0].getsockname()
-            if self.server.sockets[0].family == socket.AF_INET6:
-                self.address = "[%s]:%d" % (name[0], name[1])
-            else:
-                self.address = "%s:%d" % (name[0], name[1])
-
-        self.start = start_tcp
+        self.server = TCPStreamServer(host, port, self._client_handler, self.logger)
 
     def start_unix_server(self, path):
-        def cleanup():
-            os.unlink(path)
-
-        def start_unix():
-            cwd = os.getcwd()
-            try:
-                # Work around path length limits in AF_UNIX
-                os.chdir(os.path.dirname(path))
-                self.server = self.loop.run_until_complete(
-                    asyncio.start_unix_server(self.handle_client, os.path.basename(path))
-                )
-            finally:
-                os.chdir(cwd)
-
-            self.logger.debug('Listening on %r' % path)
+        self.server = UnixStreamServer(path, self._client_handler, self.logger)
 
-            self._cleanup_socket = cleanup
-            self.address = "unix://%s" % os.path.abspath(path)
-
-        self.start = start_unix
-
-    @abc.abstractmethod
-    def accept_client(self, reader, writer):
-        pass
-
-    async def handle_client(self, reader, writer):
-        # writer.transport.set_write_buffer_limits(0)
+    async def _client_handler(self, socket):
         try:
-            client = self.accept_client(reader, writer)
+            client = self.accept_client(socket)
             await client.process_requests()
         except Exception as e:
             import traceback
-            self.logger.error('Error from client: %s' % str(e), exc_info=True)
+
+            self.logger.error("Error from client: %s" % str(e), exc_info=True)
             traceback.print_exc()
-            writer.close()
-        self.logger.debug('Client disconnected')
+            await socket.close()
+        self.logger.debug("Client disconnected")
 
-    def run_loop_forever(self):
-        try:
-            self.loop.run_forever()
-        except KeyboardInterrupt:
-            pass
+    @abc.abstractmethod
+    def accept_client(self, socket):
+        pass
+
+    async def stop(self):
+        self.logger.debug("Stopping server")
+        await self.server.stop()
+
+    def start(self):
+        tasks = self.server.start(self.loop)
+        self.address = self.server.address
+        return tasks
 
     def signal_handler(self):
         self.logger.debug("Got exit signal")
-        self.loop.stop()
+        self.loop.create_task(self.stop())
 
-    def _serve_forever(self):
+    def _serve_forever(self, tasks):
         try:
             self.loop.add_signal_handler(signal.SIGTERM, self.signal_handler)
+            self.loop.add_signal_handler(signal.SIGINT, self.signal_handler)
+            self.loop.add_signal_handler(signal.SIGQUIT, self.signal_handler)
             signal.pthread_sigmask(signal.SIG_UNBLOCK, [signal.SIGTERM])
 
-            self.run_loop_forever()
-            self.server.close()
+            self.loop.run_until_complete(asyncio.gather(*tasks))
 
-            self.loop.run_until_complete(self.server.wait_closed())
-            self.logger.debug('Server shutting down')
+            self.logger.debug("Server shutting down")
         finally:
-            if self._cleanup_socket is not None:
-                self._cleanup_socket()
+            self.server.cleanup()
 
     def serve_forever(self):
         """
         Serve requests in the current process
         """
+        self._create_loop()
+        tasks = self.start()
+        self._serve_forever(tasks)
+        self.loop.close()
+
+    def _create_loop(self):
         # Create loop and override any loop that may have existed in
         # a parent process.  It is possible that the usecases of
         # serve_forever might be constrained enough to allow using
         # get_event_loop here, but better safe than sorry for now.
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
-        self.start()
-        self._serve_forever()
 
     def serve_as_process(self, *, prefunc=None, args=()):
         """
         Serve requests in a child process
         """
+
         def run(queue):
             # Create loop and override any loop that may have existed
             # in a parent process.  Without doing this and instead
@@ -259,18 +266,19 @@ class AsyncServer(object):
             # more general, though, as any potential use of asyncio in
             # Cooker could create a loop that needs to replaced in this
             # new process.
-            self.loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(self.loop)
+            self._create_loop()
             try:
-                self.start()
+                self.address = None
+                tasks = self.start()
             finally:
+                # Always put the server address to wake up the parent task
                 queue.put(self.address)
                 queue.close()
 
             if prefunc is not None:
                 prefunc(self, *args)
 
-            self._serve_forever()
+            self._serve_forever(tasks)
 
             if sys.version_info >= (3, 6):
                 self.loop.run_until_complete(self.loop.shutdown_asyncgens())