From 8569450a85c81044f1cb44b743a07ece60577cbb Mon Sep 17 00:00:00 2001 From: yueyoum Date: Tue, 28 Jul 2015 03:18:09 +0800 Subject: [PATCH] I add this feature for using uwsgi websocket with gevent more easy and convenient. Example: tests/websockets_chat_2.py Add: add a new API `uwsgi.input_object()` for get `struct wsgi_req *` in python. Modify: `uwsgi.websocket_send/send_binary/recv/recv_nb` now has a kwarg: `uwsgi_input`. If not provide the `uwsgi_input` keyword argument, The four function's behavior not changed. When provide the `uwsgi_input` keyword argument, The four function will get `wsgi_req` from the argument. In this way, uwsgi websocket functions can be used in other greenlets. --- plugins/python/uwsgi_pymodule.c | 164 ++++++++++++++++++++++++++++---- plugins/python/uwsgi_python.h | 15 +++ tests/websockets_chat_2.py | 145 ++++++++++++++++++++++++++++ 3 files changed, 308 insertions(+), 16 deletions(-) create mode 100644 tests/websockets_chat_2.py diff --git a/plugins/python/uwsgi_pymodule.c b/plugins/python/uwsgi_pymodule.c index 54273199..5895a5e3 100644 --- a/plugins/python/uwsgi_pymodule.c +++ b/plugins/python/uwsgi_pymodule.c @@ -4,6 +4,52 @@ extern struct uwsgi_server uwsgi; extern struct uwsgi_python up; extern struct uwsgi_plugin python_plugin; +PyTypeObject uwsgi_InputIdentityType = { + PyObject_HEAD_INIT(NULL) + 0, + "uwsgi.InputIdentityType", + sizeof(uwsgi_InputIdentity), + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + Py_TPFLAGS_DEFAULT, + "", +}; + +struct wsgi_request *py_current_wsgi_req_from_input(PyObject *input) { + uwsgi_InputIdentity *identity = (uwsgi_InputIdentity *)input; + if (identity->wsgi_req == NULL || !(identity->wsgi_req->fd)) { + return NULL; + } + + //verify + if(uwsgi.mywid != identity->mywid || + uwsgi.workers[uwsgi.mywid].cores[identity->wsgi_req->async_id].requests != identity->requests) { + return NULL; + } + + int in_request = uwsgi.workers[uwsgi.mywid].cores[identity->wsgi_req->async_id].in_request; + if (!in_request) { + return NULL; + } + + return identity->wsgi_req; +} + + static PyObject *py_uwsgi_add_var(PyObject * self, PyObject * args) { char *key = NULL; Py_ssize_t keylen = 0; @@ -998,6 +1044,22 @@ PyObject *py_uwsgi_connection_fd(PyObject * self, PyObject * args) { return PyInt_FromLong(wsgi_req->fd); } +PyObject *py_uwsgi_input_object(PyObject * self, PyObject * args) { + struct wsgi_request *wsgi_req = py_current_wsgi_req(); + + uwsgi_InputIdentity *input = PyObject_New(uwsgi_InputIdentity, &uwsgi_InputIdentityType); + if(input == NULL) + { + return NULL; + } + + input->mywid = uwsgi.mywid; + input->requests = uwsgi.workers[uwsgi.mywid].cores[wsgi_req->async_id].requests; + input->wsgi_req = wsgi_req; + return (PyObject *)input; +} + + PyObject *py_uwsgi_websocket_handshake(PyObject * self, PyObject * args) { char *key = NULL; Py_ssize_t key_len = 0; @@ -1026,15 +1088,29 @@ PyObject *py_uwsgi_websocket_handshake(PyObject * self, PyObject * args) { return Py_None; } -PyObject *py_uwsgi_websocket_send(PyObject * self, PyObject * args) { + +PyObject *py_uwsgi_websocket_send(PyObject * self, PyObject * args, PyObject * kwargs) { char *message = NULL; Py_ssize_t message_len = 0; - if (!PyArg_ParseTuple(args, "s#:websocket_send", &message, &message_len)) { + PyObject *uwsgi_input = NULL; + + static char *kwlist[] = {"message", "uwsgi_input", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s#|O:websocket_send", kwlist, &message, &message_len, &uwsgi_input)) { return NULL; } - struct wsgi_request *wsgi_req = py_current_wsgi_req(); + struct wsgi_request *wsgi_req; + if(uwsgi_input == NULL) { + wsgi_req = py_current_wsgi_req(); + } else { + wsgi_req = py_current_wsgi_req_from_input(uwsgi_input); + } + + if(wsgi_req == NULL) { + return PyErr_Format(PyExc_IOError, "unable to send websocket message"); + } UWSGI_RELEASE_GIL int ret = uwsgi_websocket_send(wsgi_req, message, message_len); @@ -1046,15 +1122,28 @@ PyObject *py_uwsgi_websocket_send(PyObject * self, PyObject * args) { return Py_None; } -PyObject *py_uwsgi_websocket_send_binary(PyObject * self, PyObject * args) { +PyObject *py_uwsgi_websocket_send_binary(PyObject * self, PyObject * args, PyObject * kwargs) { char *message = NULL; Py_ssize_t message_len = 0; - if (!PyArg_ParseTuple(args, "s#:websocket_send_binary", &message, &message_len)) { + PyObject *uwsgi_input = NULL; + + static char *kwlist[] = {"message", "uwsgi_input", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s#|O:websocket_send_binary", kwlist, &message, &message_len, &uwsgi_input)) { return NULL; } - struct wsgi_request *wsgi_req = py_current_wsgi_req(); + struct wsgi_request *wsgi_req; + if(uwsgi_input == NULL) { + wsgi_req = py_current_wsgi_req(); + } else { + wsgi_req = py_current_wsgi_req_from_input(uwsgi_input); + } + + if(wsgi_req == NULL) { + return PyErr_Format(PyExc_IOError, "unable to send websocket binary message"); + } UWSGI_RELEASE_GIL int ret = uwsgi_websocket_send_binary(wsgi_req, message, message_len); @@ -1103,11 +1192,29 @@ PyObject *py_uwsgi_chunked_read_nb(PyObject * self, PyObject * args) { -PyObject *py_uwsgi_websocket_recv(PyObject * self, PyObject * args) { - struct wsgi_request *wsgi_req = py_current_wsgi_req(); - UWSGI_RELEASE_GIL +PyObject *py_uwsgi_websocket_recv(PyObject * self, PyObject * args, PyObject * kwargs) { + PyObject *uwsgi_input = NULL; + + static char *kwlist[] = {"uwsgi_input", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O:websocket_recv", kwlist, &uwsgi_input)) { + return NULL; + } + + struct wsgi_request *wsgi_req; + if(uwsgi_input == NULL) { + wsgi_req = py_current_wsgi_req(); + } else { + wsgi_req = py_current_wsgi_req_from_input(uwsgi_input); + } + + if(wsgi_req == NULL) { + return PyErr_Format(PyExc_IOError, "unable to receive websocket message"); + } + + UWSGI_RELEASE_GIL struct uwsgi_buffer *ub = uwsgi_websocket_recv(wsgi_req); - UWSGI_GET_GIL + UWSGI_GET_GIL if (!ub) { return PyErr_Format(PyExc_IOError, "unable to receive websocket message"); } @@ -1117,8 +1224,26 @@ PyObject *py_uwsgi_websocket_recv(PyObject * self, PyObject * args) { return ret; } -PyObject *py_uwsgi_websocket_recv_nb(PyObject * self, PyObject * args) { - struct wsgi_request *wsgi_req = py_current_wsgi_req(); +PyObject *py_uwsgi_websocket_recv_nb(PyObject * self, PyObject * args, PyObject * kwargs) { + PyObject *uwsgi_input = NULL; + + static char *kwlist[] = {"uwsgi_input", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O:websocket_recv_nb", kwlist, &uwsgi_input)) { + return NULL; + } + + struct wsgi_request *wsgi_req; + if(uwsgi_input == NULL) { + wsgi_req = py_current_wsgi_req(); + } else { + wsgi_req = py_current_wsgi_req_from_input(uwsgi_input); + } + + if(wsgi_req == NULL) { + return PyErr_Format(PyExc_IOError, "unable to receive websocket message"); + } + UWSGI_RELEASE_GIL struct uwsgi_buffer *ub = uwsgi_websocket_recv_nb(wsgi_req); UWSGI_GET_GIL @@ -2599,10 +2724,12 @@ static PyMethodDef uwsgi_advanced_methods[] = { {"set_user_harakiri", py_uwsgi_set_user_harakiri, METH_VARARGS, ""}, - {"websocket_recv", py_uwsgi_websocket_recv, METH_VARARGS, ""}, - {"websocket_recv_nb", py_uwsgi_websocket_recv_nb, METH_VARARGS, ""}, - {"websocket_send", py_uwsgi_websocket_send, METH_VARARGS, ""}, - {"websocket_send_binary", py_uwsgi_websocket_send_binary, METH_VARARGS, ""}, + {"input_object", py_uwsgi_input_object, METH_VARARGS, ""}, + + {"websocket_recv", (PyCFunction)py_uwsgi_websocket_recv, METH_VARARGS|METH_KEYWORDS, ""}, + {"websocket_recv_nb", (PyCFunction)py_uwsgi_websocket_recv_nb, METH_VARARGS|METH_KEYWORDS, ""}, + {"websocket_send", (PyCFunction)py_uwsgi_websocket_send, METH_VARARGS|METH_KEYWORDS|METH_KEYWORDS, ""}, + {"websocket_send_binary", (PyCFunction)py_uwsgi_websocket_send_binary, METH_VARARGS|METH_KEYWORDS, ""}, {"websocket_handshake", py_uwsgi_websocket_handshake, METH_VARARGS, ""}, {"chunked_read", py_uwsgi_chunked_read, METH_VARARGS, ""}, @@ -3455,6 +3582,11 @@ void init_uwsgi_module_advanced(PyObject * current_uwsgi_module) { Py_DECREF(func); } + uwsgi_InputIdentityType.tp_new = PyType_GenericNew; + if(PyType_Ready(&uwsgi_InputIdentityType) < 0) { + uwsgi_log("UWSGIInput not ready\n"); + exit(1); + } } void init_uwsgi_module_cache(PyObject * current_uwsgi_module) { diff --git a/plugins/python/uwsgi_python.h b/plugins/python/uwsgi_python.h index 430fde60..84db9739 100644 --- a/plugins/python/uwsgi_python.h +++ b/plugins/python/uwsgi_python.h @@ -100,6 +100,19 @@ typedef struct uwsgi_Input { struct wsgi_request *wsgi_req; } uwsgi_Input; + +// this struct is used for: +// get this object in python, +// and pass back to uwsgi. +// uwsgi verify whether this is a valid `struct wsgi_request *` pointer +typedef struct uwsgi_InputIdentity { + PyObject_HEAD + int mywid; + uint64_t requests; + struct wsgi_request *wsgi_req; +} uwsgi_InputIdentity; + + struct uwsgi_python { char *home; @@ -307,6 +320,8 @@ void uwsgi_python_set_thread_name(int); return PyErr_Format(PyExc_SystemError, "you can call uwsgi api function only from the main callable");\ } +struct wsgi_request *py_current_wsgi_req_from_input(PyObject *); + #define uwsgi_pyexit {PyErr_Print();exit(1);} #ifdef __linux__ diff --git a/tests/websockets_chat_2.py b/tests/websockets_chat_2.py new file mode 100644 index 00000000..e946fdb2 --- /dev/null +++ b/tests/websockets_chat_2.py @@ -0,0 +1,145 @@ +#!./uwsgi --http-socket :9090 --http-raw-body --gevent 100 --module tests.websockets_chat_2 +import uwsgi +import time + +import gevent +from gevent.queue import Queue + +class ClientManager(object): + clients = set() + + @classmethod + def add(cls, client): + cls.clients.add(client) + + @classmethod + def remove(cls, client): + cls.clients.remove(client) + + @classmethod + def count(cls): + return len(cls.clients) + + @classmethod + def broadcast(cls, data): + data = "{0} {1}".format(time.time(), data) + def do_broadcast(): + for c in cls.clients: + c.send(data) + + gevent.spawn(do_broadcast) + + +class Client(object): + def __init__(self): + self.uwsgi_input = None + self.send_queue = Queue() + self.jobs = [] + + + def _recv_job(self): + while True: + data = uwsgi.websocket_recv(uwsgi_input=self.uwsgi_input) + self.on_data(data) + + def _send_job(self): + while True: + data = self.send_queue.get() + uwsgi.websocket_send(data, uwsgi_input=self.uwsgi_input) + + def _exit(self, *args): + for j in self.jobs: + j.unlink(self._exit) + + gevent.killall(self.jobs) + ClientManager.remove(self) + self.on_exit() + + + def on_data(self, data): + print "GOT: {0}".format(data) + ClientManager.broadcast(data) + + + def on_exit(self): + print "bye bye..." + + + def send(self, data): + self.send_queue.put(data) + + + def start(self): + uwsgi.websocket_handshake() + self.uwsgi_input = uwsgi.input_object() + + ClientManager.add(self) + + self.jobs.extend([ + gevent.spawn(self._recv_job), + gevent.spawn(self._send_job), + ]) + + for j in self.jobs: + j.link(self._exit) + + gevent.joinall(self.jobs) + + + + +def application(env, sr): + + ws_scheme = 'ws' + if 'HTTPS' in env or env['wsgi.url_scheme'] == 'https': + ws_scheme = 'wss' + + if env['PATH_INFO'] == '/': + sr('200 OK', [('Content-Type', 'text/html')]) + return """ + + + + + +

WebSocket

+ + +
+
+ + + """ % (ws_scheme, env['HTTP_HOST']) + elif env['PATH_INFO'] == '/favicon.ico': + return "" + elif env['PATH_INFO'] == '/foobar/': + print "websockets..." + client = Client() + client.start() + + return "" +