kore

commit f4cd70956b754c98a256ce918545df6d398f7167
parent 1e7ccc2adfe4964eff90c35107e856e5ccecae05
Author: Joris Vink <joris@coders.se>
Date:   Mon, 25 Feb 2019 10:35:00 +0100

Add an optional timeout to socketop.recv().

Diffstat:
examples/python-echo/src/echo.py
 | 
10
+++++++++-
include/kore/python_methods.h
 | 
1
+
src/python.c
 | 
46
+++++++++++++++++++++++++++++++++++++++++++---

3 files changed, 53 insertions(+), 4 deletions(-)
diff --git a/examples/python-echo/src/echo.py b/examples/python-echo/src/echo.py
@@ -40,13 +40,21 @@ class EchoServer:
                 kore.fatal("exception %s" % e)
 
     # Each client will run as this co-routine.
+    # In this case we pass a timeout of 1 second to the recv() call
+    # which will throw a TimeoutError exception in case the timeout
+    # is hit before data is read from the socket.
+    #
+    # This timeout argument is optional. If none is specified the call
+    # will wait until data becomes available.
     async def handle_client(self, client):
         while True:
             try:
-                data = await client.recv(1024)
+                data = await client.recv(1024, 1000)
                 if data is None:
                     break
                 await client.send(data)
+            except TimeoutError as e:
+                print("timed out reading (%s)" % e)
             except Exception as e:
                 print("client got exception %s" % e)
         client.close()
diff --git a/include/kore/python_methods.h b/include/kore/python_methods.h
@@ -210,6 +210,7 @@ struct pysocket_data {
 	struct kore_buf		buffer;
 	struct sockaddr_in	sendaddr;
 	struct pysocket		*socket;
+	struct kore_timer	*timer;
 };
 
 struct pysocket_op {
diff --git a/src/python.c b/src/python.c
@@ -50,6 +50,7 @@ static int			python_coro_run(struct python_coro *);
 static void			python_coro_wakeup(struct python_coro *);
 
 static void		pysocket_evt_handle(void *, int);
+static void		pysocket_op_timeout(void *, u_int64_t);
 static PyObject		*pysocket_op_create(struct pysocket *,
 			    int, const void *, size_t);
 
@@ -1703,12 +1704,28 @@ pysocket_sendto(struct pysocket *sock, PyObject *args)
 static PyObject *
 pysocket_recv(struct pysocket *sock, PyObject *args)
 {
-	Py_ssize_t	len;
+	Py_ssize_t		len;
+	struct pysocket_op	*op;
+	PyObject		*obj;
+	int			timeo;
 
-	if (!PyArg_ParseTuple(args, "n", &len))
+	timeo = -1;
+
+	if (!PyArg_ParseTuple(args, "n|i", &len, &timeo))
 		return (NULL);
 
-	return (pysocket_op_create(sock, PYSOCKET_TYPE_RECV, NULL, len));
+	obj = pysocket_op_create(sock, PYSOCKET_TYPE_RECV, NULL, len);
+	if (obj == NULL)
+		return (NULL);
+
+	op = (struct pysocket_op *)obj;
+
+	if (timeo != -1) {
+		op->data.timer = kore_timer_add(pysocket_op_timeout,
+		    timeo, op, KORE_TIMER_ONESHOT);
+	}
+
+	return (obj);
 }
 
 static PyObject *
@@ -1826,6 +1843,11 @@ pysocket_op_dealloc(struct pysocket_op *op)
 	    op->data.type == PYSOCKET_TYPE_SEND)
 		kore_buf_cleanup(&op->data.buffer);
 
+	if (op->data.timer != NULL) {
+		kore_timer_remove(op->data.timer);
+		op->data.timer = NULL;
+	}
+
 	op->data.coro->sockop = NULL;
 	Py_DECREF(op->data.socket);
 
@@ -1860,6 +1882,7 @@ pysocket_op_create(struct pysocket *sock, int type, const void *ptr, size_t len)
 	op->data.eof = 0;
 	op->data.self = op;
 	op->data.type = type;
+	op->data.timer = NULL;
 	op->data.socket = sock;
 	op->data.evt.flags = 0;
 	op->data.coro = coro_running;
@@ -1952,6 +1975,23 @@ pysocket_op_iternext(struct pysocket_op *op)
 	return (ret);
 }
 
+static void
+pysocket_op_timeout(void *arg, u_int64_t now)
+{
+	struct pysocket_op	*op = arg;
+
+	op->data.eof = 1;
+	op->data.timer = NULL;
+
+	op->data.coro->exception = PyExc_TimeoutError;
+	op->data.coro->exception_msg = "timeout before operation completed";
+
+	if (op->data.coro->request != NULL)
+		http_request_wakeup(op->data.coro->request);
+	else
+		python_coro_wakeup(op->data.coro);
+}
+
 static PyObject *
 pysocket_async_connect(struct pysocket_op *op)
 {