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.
This commit is contained in:
yueyoum
2015-07-28 03:18:09 +08:00
parent ca673dfc5f
commit 8569450a85
3 changed files with 308 additions and 16 deletions
+148 -16
View File
@@ -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) {
+15
View File
@@ -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__