poky/bitbake/lib/bb/asyncrpc/serv.py

#
# Copyright BitBake Contributors
#
# SPDX-License-Identifier: GPL-2.0-only
#

import abc
import asyncio
import json
import os
import signal
import socket
import sys
import multiprocessing
from .connection import StreamConnection
from .exceptions import ClientError, ServerError, ConnectionClosedError


class AsyncServerConnection(object):
    # 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.handlers = {
            "ping": self.handle_ping,
        }
        self.logger = logger

    async def close(self):
        await self.socket.close()

    async def process_requests(self):
        try:
            self.logger.info("Client %r connected" % (self.socket.address,))

            # Read protocol and version
            client_protocol = await self.socket.recv()
            if not client_protocol:
                return

            (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))
                return

            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)
                )
                return

            # Read headers. Currently, no headers are implemented, so look for
            # an empty line to signal the end of the headers
            while True:
                header = await self.socket.recv()
                if not header:
                    break

            # Handle messages
            while True:
                d = await self.socket.recv_message()
                if d is None:
                    break
                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:
            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)
                return await self.handlers[k](msg[k])

        raise ClientError("Unrecognized command %r" % msg)

    async def handle_ping(self, request):
        return {"alive": True}


class StreamServer(object):
    def __init__(self, handler, logger):
        self.handler = handler
        self.logger = logger
        self.closed = False

    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


class UnixStreamServer(StreamServer):
    def __init__(self, path, handler, logger):
        super().__init__(handler, logger)
        self.path = path

    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)

        self.logger.debug("Listening on %r" % self.path)
        self.address = "unix://%s" % os.path.abspath(self.path)
        return [self.server.wait_closed()]

    async def stop(self):
        await super().stop()
        self.server.close()

    def cleanup(self):
        os.unlink(self.path)


class AsyncServer(object):
    def __init__(self, logger):
        self.logger = logger
        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_unix_server(self, path):
        self.server = UnixStreamServer(path, self._client_handler, self.logger)

    async def _client_handler(self, socket):
        try:
            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)
            traceback.print_exc()
            await socket.close()
        self.logger.debug("Client disconnected")

    @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.create_task(self.stop())

    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.loop.run_until_complete(asyncio.gather(*tasks))

            self.logger.debug("Server shutting down")
        finally:
            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)

    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
            # using get_event_loop, at the very minimum the hashserv
            # unit tests will hang when running the second test.
            # This happens since get_event_loop in the spawned server
            # process for the second testcase ends up with the loop
            # from the hashserv client created in the unit test process
            # when running the first testcase.  The problem is somewhat
            # more general, though, as any potential use of asyncio in
            # Cooker could create a loop that needs to replaced in this
            # new process.
            self._create_loop()
            try:
                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(tasks)

            if sys.version_info >= (3, 6):
                self.loop.run_until_complete(self.loop.shutdown_asyncgens())
            self.loop.close()

        queue = multiprocessing.Queue()

        # Temporarily block SIGTERM. The server process will inherit this
        # block which will ensure it doesn't receive the SIGTERM until the
        # handler is ready for it
        mask = signal.pthread_sigmask(signal.SIG_BLOCK, [signal.SIGTERM])
        try:
            self.process = multiprocessing.Process(target=run, args=(queue,))
            self.process.start()

            self.address = queue.get()
            queue.close()
            queue.join_thread()

            return self.process
        finally:
            signal.pthread_sigmask(signal.SIG_SETMASK, mask)