moonraker: use tornado's iostream wrapper for klippy connection

This reduces the amount of code needed to handle the connection and adds more robust error handling.   This also prepares moonraker for the eventual transition of Klippy hosting the server socket.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>

moonraker/moonraker.py

@@ -15,7 +15,7 @@ import errno
 import tornado
 import tornado.netutil
 import confighelper
-from tornado import gen
+from tornado import gen, iostream
 from tornado.ioloop import IOLoop, PeriodicCallback
 from tornado.util import TimeoutError
 from tornado.locks import Event
@@ -48,11 +48,9 @@ class Server:
             socketfile, backlog=1)
         self.remove_server_sock = tornado.netutil.add_accept_handler(
             self.klippy_server_sock, self._handle_klippy_connection)
-        self.klippy_sock = None
+        self.klippy_iostream = None
-        self.is_klippy_connected = False
         self.is_klippy_ready = False
         self.moonraker_available = False
-        self.partial_data = b""
 
         # Server/IOLoop
         self.server_running = False
@@ -60,7 +58,7 @@ class Server:
         self.register_endpoint = app.register_local_handler
         self.register_static_file_handler = app.register_static_file_handler
         self.register_upload_handler = app.register_upload_handler
-        self.io_loop = IOLoop.current()
+        self.ioloop = IOLoop.current()
         self.init_cb = PeriodicCallback(self._initialize, INIT_MS)
 
         # Setup remote methods accessable to Klippy.  Note that all
@@ -138,7 +136,7 @@ class Server:
     def send_event(self, event, *args):
         events = self.events.get(event, [])
         for evt in events:
-            self.io_loop.spawn_callback(evt, *args)
+            self.ioloop.spawn_callback(evt, *args)
 
     def register_remote_method(self, method_name, cb):
         if method_name in self.remote_methods:
@@ -150,42 +148,30 @@ class Server:
 
     # ***** Klippy Connection *****
     def _handle_klippy_connection(self, conn, addr):
-        if self.is_klippy_connected:
+        if self.klippy_iostream is not None and \
+                not self.klippy_iostream.closed():
             logging.info("New Connection received while Klippy Connected")
-            self.close_client_sock()
+            self.klippy_iostream.close()
         logging.info("Klippy Connection Established")
-        self.is_klippy_connected = True
+        self.klippy_iostream = iostream.IOStream(conn)
-        conn.setblocking(0)
+        self.klippy_iostream.set_close_callback(
-        self.klippy_sock = conn
+            self._handle_stream_closed)
-        self.io_loop.add_handler(
+        self.ioloop.spawn_callback(
-            self.klippy_sock.fileno(), self._handle_klippy_data,
+            self._read_klippy_stream, self.klippy_iostream)
-            IOLoop.READ | IOLoop.ERROR)
         # begin server iniialization
         self.init_cb.start()
 
-    def _handle_klippy_data(self, fd, events):
+    async def _read_klippy_stream(self, stream):
-        if events & IOLoop.ERROR:
+        while not stream.closed():
-            self.close_client_sock()
-            return
             try:
-            data = self.klippy_sock.recv(4096)
+                data = await stream.read_until(b'\x03')
-        except socket.error as e:
+            except iostream.StreamClosedError as e:
-            # If bad file descriptor allow connection to be
-            # closed by the data check
-            if e.errno == errno.EBADF:
-                data = b''
-            else:
                 return
-        if data == b'':
+            except Exception:
-            # Socket Closed
+                logging.exception("Klippy Stream Read Error")
-            self.close_client_sock()
+                continue
-            return
-        commands = data.split(b'\x03')
-        commands[0] = self.partial_data + commands[0]
-        self.partial_data = commands.pop()
-        for cmd in commands:
             try:
-                decoded_cmd = json.loads(cmd)
+                decoded_cmd = json.loads(data[:-1])
                 method = decoded_cmd.get('method')
                 params = decoded_cmd.get('params', {})
                 cb = self.remote_methods.get(method)
@@ -198,32 +184,22 @@ class Server:
                     "Error processing Klippy Host Response: %s"
                     % (cmd.decode()))
 
-    def klippy_send(self, data):
+    def _handle_stream_closed(self):
-        # TODO: need a mutex or lock to make sure that multiple co-routines
+        self.is_klippy_ready = False
-        # do not attempt to send
+        self.moonraker_available = False
-        if not self.is_klippy_connected:
+        self.init_cb.stop()
-            return False
+        for request in self.pending_requests.values():
-        retries = 10
+            request.notify(ServerError("Klippy Disconnected", 503))
-        data = json.dumps(data).encode() + b"\x03"
+        self.pending_requests = {}
-        while data:
+        logging.info("Klippy Connection Removed")
+        self.send_event("server:klippy_state_changed", "disconnect")
+
+    async def send_klippy_request(self, request):
+        data = json.dumps(request.to_dict()).encode() + b"\x03"
         try:
-                sent = self.klippy_sock.send(data)
+            await self.klippy_iostream.write(data)
-            except socket.error as e:
+        except iostream.StreamClosedError:
-                if e.errno == errno.EBADF or e.errno == errno.EPIPE \
+            request.notify(ServerError("Klippy Host not connected", 503))
-                        or not retries:
-                    sent = 0
-                else:
-                    # XXX - Should pause for 1ms here
-                    retries -= 1
-                    continue
-            retries = 10
-            if sent > 0:
-                data = data[sent:]
-            else:
-                logging.info("Error sending client data, closing socket")
-                self.close_client_sock()
-                return False
-        return True
 
     async def _initialize(self):
         await self._request_endpoints()
@@ -308,11 +284,8 @@ class Server:
     def make_request(self, path, method, args):
         base_request = BaseRequest(path, method, args)
         self.pending_requests[base_request.id] = base_request
-        ret = self.klippy_send(base_request.to_dict())
+        self.ioloop.spawn_callback(
-        if not ret:
+            self.send_klippy_request, base_request)
-            self.pending_requests.pop(base_request.id, None)
-            base_request.notify(
-                ServerError("Klippy Host not connected", 503))
         return base_request
 
     async def _kill_server(self):
@@ -324,29 +297,14 @@ class Server:
         for plugin in self.plugins:
             if hasattr(plugin, "close"):
                 await plugin.close()
-        self.close_client_sock()
+        if self.klippy_iostream is not None and \
+                not self.klippy_iostream.closed():
+            self.klippy_iostream.close()
         self.close_server_sock()
         if self.server_running:
             self.server_running = False
             await self.moonraker_app.close()
-            self.io_loop.stop()
+            self.ioloop.stop()
-
-    def close_client_sock(self):
-        self.is_klippy_ready = False
-        self.moonraker_available = False
-        self.init_cb.stop()
-        for request in self.pending_requests.values():
-            request.notify(ServerError("Klippy Disconnected", 503))
-        self.pending_requests = {}
-        if self.is_klippy_connected:
-            self.is_klippy_connected = False
-            logging.info("Klippy Connection Removed")
-            try:
-                self.io_loop.remove_handler(self.klippy_sock.fileno())
-                self.klippy_sock.close()
-            except socket.error:
-                logging.exception("Error Closing Client Socket")
-            self.send_event("server:klippy_state_changed", "disconnect")
 
     def close_server_sock(self):
         try: