feat: small speed up to the aio message reader (#273)

This commit is contained in:
J. Nick Koston 2023-12-03 16:49:49 -10:00 committed by GitHub
parent 210323953c
commit 8ee18a1355
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 37 additions and 25 deletions

View File

@ -3,3 +3,11 @@
import cython
from .._private.unmarshaller cimport Unmarshaller
cpdef _message_reader(
Unmarshaller unmarshaller,
object process,
object finalize,
bint negotiate_unix_fd
)

View File

@ -1,11 +1,39 @@
import logging
import socket
from functools import partial
from typing import Callable, Optional
from .._private.unmarshaller import Unmarshaller
from ..message import Message
def _message_reader(
unmarshaller: Unmarshaller,
process: Callable[[Message], None],
finalize: Callable[[Optional[Exception]], None],
negotiate_unix_fd: bool,
) -> None:
"""Reads messages from the unmarshaller and passes them to the process function."""
try:
while True:
message = unmarshaller._unmarshall()
if message is None:
return
try:
process(message)
except Exception as e:
logging.error("Unexpected error processing message: %s", exc_info=True)
# If we are not negotiating unix fds, we can stop reading as soon as we have
# the buffer is empty as asyncio will call us again when there is more data.
if (
not negotiate_unix_fd
and not unmarshaller._has_another_message_in_buffer()
):
return
except Exception as e:
finalize(e)
def build_message_reader(
sock: Optional[socket.socket],
process: Callable[[Message], None],
@ -14,28 +42,4 @@ def build_message_reader(
) -> Callable[[], None]:
"""Build a callable that reads messages from the unmarshaller and passes them to the process function."""
unmarshaller = Unmarshaller(None, sock, negotiate_unix_fd)
def _message_reader() -> None:
"""Reads messages from the unmarshaller and passes them to the process function."""
try:
while True:
message = unmarshaller._unmarshall()
if message is None:
return
try:
process(message)
except Exception as e:
logging.error(
"Unexpected error processing message: %s", exc_info=True
)
# If we are not negotiating unix fds, we can stop reading as soon as we have
# the buffer is empty as asyncio will call us again when there is more data.
if (
not negotiate_unix_fd
and not unmarshaller._has_another_message_in_buffer()
):
return
except Exception as e:
finalize(e)
return _message_reader
return partial(_message_reader, unmarshaller, process, finalize, negotiate_unix_fd)