🏡 index : quine-global-labs/mirror-cockpit-project-cockpit.git

# SPDX-License-Identifier: LGPL-2.1-or-later

class CockpitBufferedProtocol(CockpitBaseProtocol, asyncio.BufferedProtocol):
    """A hopefully-reasonably-efficient implementation of the Cockpit frame protocol

    We try to operate on a principle of finding a good balance between not
    making too many small reads, and not making too many large copies.
    Perfection is impossible here: in order to avoid all copies, we would need
    to know the size of the frames in advance, and we can't know that without
    doing too many small reads.

    Most incoming messages to the bridge are high tens to low hundreds of
    bytes, and they often come grouped together.  We should be able to manage
    the vast majority by reading in medium-sized chunks and processing multiple
    frames in a single iteration.  If we have a partial frame in a chunk, we
    leave it for the next iteration (which is the only place we do copies).

    A larger buffer allows processing more frames in a single read, and also
    makes it less likely that we'll need to adjust the buffer for larger
    frames, or copy data when we get flooded with lots of request.  It also
    means that in the event that we *do* need to adjust the buffer or copy
    data, we'll end up copying more.  4096 is picked out of thin air.

    We handle large frame sizes by adjusting the buffer size for the next read,
    so that we don't have to copy (most of) the content of large frames.
    """
    min_buffer_size = None
    max_buffer_size = None
    buffer = None
    in_buffer = 0

    def __init__(self, buffer_size=4096):
        self.min_buffer_size = buffer_size
        self.max_buffer_size = buffer_size * 4
        self.buffer = memoryview(bytearray(self.min_buffer_size))
        self.in_buffer = 0

    def get_buffer(self, sizehint=-1):
        return self.buffer[self.in_buffer:]

    def buffer_updated(self, nbytes):
        self.in_buffer += nbytes

        # This is all the data we have so far
        view = self.buffer[:self.in_buffer]

        # Try to consume some frames
        consumed = 0
        while (result := self.consume_one_frame(view)) > 0:
            view = view[result:]
            consumed += result

        if not consumed and len(view) - result < len(self.buffer):
            # This is the easy case: we didn't consume anything, and there's
            # enough room in the buffer for the next packet.
            return

        if not view and len(self.buffer) <= self.max_buffer_size:
            # There is no unconsumed data.  That also implies that there is no
            # size hint, because where would it have come from?  If we're happy
            # with the buffer size, then we can just reuse it.
            self.in_buffer = 0
            return

        # Otherwise, the general case.  We need to allocate a new buffer, and
        # possibly copy the leftover data into it.
        new_buffer_size = max(len(view) - result, self.min_buffer_size)
        new_buffer = bytearray(new_buffer_size)
        new_buffer[:len(view)] = view
        self.buffer = memoryview(new_buffer)
        self.in_buffer = len(view)

    def eof_received(self):
        raise ValueError('eof received')


class BufferedProtocolFeeder:
    def __init__(self, fd, protocol, *args, **kwargs):
        self.fd = fd
        self.fileio = io.FileIO(self.fd, closefd=False)
        self.listener = protocol(*args, **kwargs)
        self.loop = asyncio.get_event_loop()
        self.loop.add_reader(self.fd, self.ready)

    # From BufferedProtocol docs:
    #  start -> CM [-> GB [-> BU?]]* [-> ER?] -> CL -> end
    def ready(self):
        buffer = self.listener.get_buffer()
        n_bytes = self.fileio.readinto(buffer)
        if n_bytes:
            self.listener.buffer_updated(n_bytes)
        else:
            self.listener.eof_received()
            self.loop.remove_reader(self.fd)