moonraker: use native asyncio unix streams

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Eric Callahan 2021-12-09 11:08:56 -05:00
parent df674ae476
commit ae4c4ad678
1 changed files with 33 additions and 23 deletions

View File

@ -20,7 +20,6 @@ import signal
import confighelper import confighelper
import utils import utils
import asyncio import asyncio
from tornado import iostream
from tornado.httpclient import AsyncHTTPClient from tornado.httpclient import AsyncHTTPClient
from eventloop import EventLoop from eventloop import EventLoop
from app import MoonrakerApp from app import MoonrakerApp
@ -50,6 +49,7 @@ if TYPE_CHECKING:
INIT_TIME = .25 INIT_TIME = .25
LOG_ATTEMPT_INTERVAL = int(2. / INIT_TIME + .5) LOG_ATTEMPT_INTERVAL = int(2. / INIT_TIME + .5)
MAX_LOG_ATTEMPTS = 10 * LOG_ATTEMPT_INTERVAL MAX_LOG_ATTEMPTS = 10 * LOG_ATTEMPT_INTERVAL
UNIX_BUFFER_LIMIT = 2 * 1024 * 1024
CORE_COMPONENTS = [ CORE_COMPONENTS = [
'database', 'file_manager', 'klippy_apis', 'machine', 'database', 'file_manager', 'klippy_apis', 'machine',
@ -637,7 +637,7 @@ class Server:
try: try:
if self.klippy_connection.is_connected(): if self.klippy_connection.is_connected():
self.klippy_disconnect_evt = asyncio.Event() self.klippy_disconnect_evt = asyncio.Event()
self.klippy_connection.close() await self.klippy_connection.close()
await asyncio.wait_for( await asyncio.wait_for(
self.klippy_disconnect_evt.wait(), 2.) self.klippy_disconnect_evt.wait(), 2.)
self.klippy_disconnect_evt = None self.klippy_disconnect_evt = None
@ -696,7 +696,9 @@ class KlippyConnection:
on_close: Callable[[], None], on_close: Callable[[], None],
event_loop: EventLoop event_loop: EventLoop
) -> None: ) -> None:
self.iostream: Optional[iostream.IOStream] = None self.writer: Optional[asyncio.StreamWriter] = None
self.closed: bool = True
self.close_mutex: asyncio.Lock = asyncio.Lock()
self.on_recd = on_recd self.on_recd = on_recd
self.on_close = on_close self.on_close = on_close
self.event_loop = event_loop self.event_loop = event_loop
@ -714,24 +716,23 @@ class KlippyConnection:
self.log_no_access = False self.log_no_access = False
return False return False
self.log_no_access = True self.log_no_access = True
ksock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
kstream = iostream.IOStream(ksock)
try: try:
await kstream.connect(address) reader, writer = await asyncio.open_unix_connection(
except iostream.StreamClosedError: address, limit=UNIX_BUFFER_LIMIT)
except Exception:
return False return False
logging.info("Klippy Connection Established") logging.info("Klippy Connection Established")
self.iostream = kstream self.closed = False
self.iostream.set_close_callback(self.on_close) self.writer = writer
self.event_loop.register_callback(self._read_stream, self.iostream) self.event_loop.register_callback(self._read_stream, reader)
return True return True
async def _read_stream(self, stream: iostream.IOStream) -> None: async def _read_stream(self, reader: asyncio.StreamReader) -> None:
while not stream.closed(): while not reader.at_eof():
try: try:
data = await stream.read_until(b'\x03') data = await reader.readuntil(b'\x03')
except iostream.StreamClosedError as e: except asyncio.IncompleteReadError:
return break
except Exception: except Exception:
logging.exception("Klippy Stream Read Error") logging.exception("Klippy Stream Read Error")
continue continue
@ -741,24 +742,33 @@ class KlippyConnection:
except Exception: except Exception:
logging.exception( logging.exception(
f"Error processing Klippy Host Response: {data.decode()}") f"Error processing Klippy Host Response: {data.decode()}")
await self.close()
async def send_request(self, request: BaseRequest) -> None: async def send_request(self, request: BaseRequest) -> None:
if self.iostream is None: if self.writer is None:
request.notify(ServerError("Klippy Host not connected", 503)) request.notify(ServerError("Klippy Host not connected", 503))
return return
data = json.dumps(request.to_dict()).encode() + b"\x03" data = json.dumps(request.to_dict()).encode() + b"\x03"
try: try:
await self.iostream.write(data) self.writer.write(data)
except iostream.StreamClosedError: await self.writer.drain()
except Exception:
request.notify(ServerError("Klippy Host not connected", 503)) request.notify(ServerError("Klippy Host not connected", 503))
await self.close()
def is_connected(self) -> bool: def is_connected(self) -> bool:
return self.iostream is not None and not self.iostream.closed() return not self.closed
def close(self) -> None: async def close(self) -> None:
if self.iostream is not None and \ async with self.close_mutex:
not self.iostream.closed(): if self.writer is not None:
self.iostream.close() self.writer.close()
await self.writer.wait_closed()
self.writer = None
if not self.closed:
self.closed = True
self.on_close()
self.closing = False
# Basic WebRequest class, easily converted to dict for json encoding # Basic WebRequest class, easily converted to dict for json encoding
class BaseRequest: class BaseRequest: