Speed up reading of command responses

kitty/kittens.c

@@ -7,6 +7,92 @@
 
 #include "data-types.h"
 
+#define CMD_BUF_SZ 2048
+
+
+static inline bool
+append_buf(char buf[CMD_BUF_SZ], size_t *pos, PyObject *ans) {
+    if (*pos) {
+        PyObject *bytes = PyBytes_FromStringAndSize(buf, *pos);
+        if (!bytes) { PyErr_NoMemory(); return false; }
+        int ret = PyList_Append(ans, bytes);
+        Py_CLEAR(bytes);
+        if (ret != 0) return false;
+        *pos = 0;
+    }
+    return true;
+}
+
+static inline bool
+add_char(char buf[CMD_BUF_SZ], size_t *pos, char ch, PyObject *ans) {
+    if (*pos >= CMD_BUF_SZ) {
+        if (!append_buf(buf, pos, ans)) return false;
+    }
+    buf[*pos] = ch;
+    *pos += 1;
+    return true;
+}
+
+static inline bool
+read_response(int fd, double timeout, char buf[CMD_BUF_SZ], size_t *pos, PyObject *ans) {
+    enum ReadState {START, STARTING_ESC, P, AT, K, I, T, T2, Y, HYPHEN, C, M, BODY, TRAILING_ESC};
+    enum ReadState state = START;
+    char ch;
+    double end_time = monotonic() + timeout;
+    while(monotonic() <= end_time) {
+        ssize_t len = read(fd, &ch, 1);
+        if (len == 0) continue;
+        if (len < 0) {
+            if (errno == EINTR || errno == EAGAIN) continue;
+            PyErr_SetFromErrno(PyExc_OSError);
+            return false;
+            continue;
+        }
+        end_time = monotonic() + timeout;
+        switch(state) {
+            case START:
+                if (ch == 0x1b) state = STARTING_ESC;
+                break;
+#define CASE(curr, q, next) case curr: state = ch == q ? next : START; break;
+            CASE(STARTING_ESC, 'P', P);
+            CASE(P, '@', AT);
+            CASE(AT, 'k', K);
+            CASE(K, 'i', I);
+            CASE(I, 't', T);
+            CASE(T, 't', T2);
+            CASE(T2, 'y', Y);
+            CASE(Y, '-', HYPHEN);
+            CASE(HYPHEN, 'c', C);
+            CASE(C, 'm', M);
+            CASE(M, 'd', BODY);
+            case BODY:
+                if (ch == 0x1b) { state = TRAILING_ESC; }
+                else {
+                    if (!add_char(buf, pos, ch, ans)) return false;
+                }
+                break;
+            case TRAILING_ESC:
+                if (ch == '\\') return append_buf(buf, pos, ans);
+                state = BODY;
+                break;
+        }
+    }
+    PyErr_SetString(PyExc_TimeoutError, "Timed out while waiting to read cmd response");
+    return false;
+}
+
+static PyObject*
+read_command_response(PyObject *self UNUSED, PyObject *args) {
+    double timeout;
+    int fd;
+    PyObject *ans;
+    if (!PyArg_ParseTuple(args, "idO!", &fd, &timeout, &PyList_Type, &ans)) return NULL;
+    static char buf[CMD_BUF_SZ];
+    size_t pos = 0;
+    if (!read_response(fd, timeout, buf, &pos, ans)) return NULL;
+    Py_RETURN_NONE;
+}
+
 static PyObject*
 parse_input_from_terminal(PyObject *self UNUSED, PyObject *args) {
     enum State { NORMAL, ESC, CSI, ST, ESC_ST };
@@ -105,6 +191,7 @@ parse_input_from_terminal(PyObject *self UNUSED, PyObject *args) {
 
 static PyMethodDef module_methods[] = {
     METHODB(parse_input_from_terminal, METH_VARARGS),
+    METHODB(read_command_response, METH_VARARGS),
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

kitty/remote_control.py

@@ -11,6 +11,7 @@ from functools import partial
 from .cli import emph, parse_args
 from .cmds import cmap, parse_subcommand_cli
 from .constants import appname, version
+from .fast_data_types import read_command_response
 from .utils import TTYIO, parse_address_spec
 
 
@@ -73,13 +74,23 @@ class SocketIO:
                     out.flush()
         self.socket.shutdown(socket.SHUT_WR)
 
-    def recv(self, more_needed, timeout):
-        # We dont bother with more_needed since the server will close the
-        # connection after transmission
+    def recv(self, timeout):
+        dcs = re.compile(br'\x1bP@kitty-cmd([^\x1b]+)\x1b\\')
         self.socket.settimeout(timeout)
         with self.socket.makefile('rb') as src:
             data = src.read()
-        more_needed(data)
+        m = dcs.search(data)
+        if m is None:
+            raise TimeoutError('Timed out while waiting to read cmd response')
+        return m.group(1)
+
+
+class RCIO(TTYIO):
+
+    def recv(self, timeout):
+        ans = []
+        read_command_response(self.tty_fd, timeout, ans)
+        return b''.join(ans)
 
 
 def do_io(to, send, no_response):
@@ -93,28 +104,14 @@ def do_io(to, send, no_response):
                 yield encode_send(send)
         send_data = send_generator()
 
-    io = SocketIO(to) if to else TTYIO()
+    io = SocketIO(to) if to else RCIO()
     with io:
         io.send(send_data)
         if no_response:
             return {'ok': True}
+        received = io.recv(timeout=10)
 
-        dcs = re.compile(br'\x1bP@kitty-cmd([^\x1b]+)\x1b\\')
-
-        received = b''
-        match = None
-
-        def more_needed(data):
-            nonlocal received, match
-            received += data
-            match = dcs.search(received)
-            return match is None
-
-        io.recv(more_needed, timeout=10)
-
-    if match is None:
-        raise SystemExit('Failed to receive response from ' + appname)
-    response = json.loads(match.group(1).decode('ascii'))
+    response = json.loads(received.decode('ascii'))
     return response