Add kore.sendmsg(object, worker=None) to the python api.

This allows you to send Python objects that can be run through pickle
to other worker processes.

If your application implements koreapp.onmsg() you will be able to receive
these objects.
This commit is contained in:
Joris Vink 2019-10-16 12:05:27 +02:00
parent ffb9fd2fc5
commit cdc3347120
7 changed files with 153 additions and 10 deletions

View File

@ -418,6 +418,8 @@ struct kore_alog_header {
u_int16_t loglen;
} __attribute__((packed));
#define KORE_WORKER_MAX UCHAR_MAX
struct kore_worker {
u_int8_t id;
u_int8_t cpu;
@ -512,6 +514,7 @@ struct kore_timer {
#define KORE_MSG_CERTIFICATE_REQ 8
#define KORE_MSG_CRL 9
#define KORE_MSG_ACCEPT_AVAILABLE 10
#define KORE_PYTHON_SEND_OBJ 11
/* Predefined message targets. */
#define KORE_MSG_PARENT 1000
@ -743,6 +746,7 @@ void kore_websocket_broadcast(struct connection *,
void kore_msg_init(void);
void kore_msg_worker_init(void);
void kore_msg_parent_init(void);
void kore_msg_unregister(u_int8_t);
void kore_msg_parent_add(struct kore_worker *);
void kore_msg_parent_remove(struct kore_worker *);
void kore_msg_send(u_int16_t, u_int8_t, const void *, size_t);

View File

@ -55,7 +55,8 @@ static PyObject *python_kore_task_create(PyObject *, PyObject *);
static PyObject *python_kore_socket_wrap(PyObject *, PyObject *);
static PyObject *python_kore_domain(PyObject *, PyObject *, PyObject *);
static PyObject *python_kore_gather(PyObject *, PyObject *, PyObject *);
static PyObject *python_kore_sendobj(PyObject *, PyObject *,
PyObject *);
static PyObject *python_kore_server(PyObject *, PyObject *,
PyObject *);
@ -100,6 +101,7 @@ static struct PyMethodDef pykore_methods[] = {
METHOD("server", python_kore_server, METH_VARARGS | METH_KEYWORDS),
METHOD("gather", python_kore_gather, METH_VARARGS | METH_KEYWORDS),
METHOD("domain", python_kore_domain, METH_VARARGS | METH_KEYWORDS),
METHOD("sendobj", python_kore_sendobj, METH_VARARGS | METH_KEYWORDS),
METHOD("websocket_broadcast", python_websocket_broadcast, METH_VARARGS),
#if defined(KORE_USE_PGSQL)
METHOD("dbsetup", python_kore_pgsql_register, METH_VARARGS),

View File

@ -1488,7 +1488,7 @@ configure_workers(char *option)
{
int err;
worker_count = kore_strtonum(option, 10, 1, 255, &err);
worker_count = kore_strtonum(option, 10, 1, KORE_WORKER_MAX, &err);
if (err != KORE_RESULT_OK) {
printf("%s is not a valid worker number\n", option);
return (KORE_RESULT_ERROR);

View File

@ -155,7 +155,9 @@ kore_keymgr_run(void)
kore_seccomp_filter("keymgr", filter_keymgr,
KORE_FILTER_LEN(filter_keymgr));
#endif
#if defined(KORE_USE_PYTHON)
kore_msg_unregister(KORE_PYTHON_SEND_OBJ);
#endif
kore_worker_privdrop(keymgr_runas_user, keymgr_root_path);
if (rand_file != NULL) {

View File

@ -243,6 +243,7 @@ main(int argc, char *argv[])
kore_platform_init();
kore_log_init();
kore_msg_init();
#if !defined(KORE_NO_HTTP)
http_parent_init();
#if defined(KORE_USE_CURL)
@ -865,7 +866,6 @@ kore_server_start(int argc, char *argv[])
}
kore_platform_proctitle("[parent]");
kore_msg_init();
kore_worker_init();
/* Set worker_max_connections for kore_connection_init(). */
@ -885,6 +885,10 @@ kore_server_start(int argc, char *argv[])
kore_timer_add(kore_accesslog_run, 100, NULL, 0);
#endif
#if defined(KORE_USE_PYTHON)
kore_msg_unregister(KORE_PYTHON_SEND_OBJ);
#endif
while (quit != 1) {
if (sig_recv != 0) {
last_sig = sig_recv;

View File

@ -28,8 +28,6 @@ struct msg_type {
TAILQ_ENTRY(msg_type) list;
};
TAILQ_HEAD(, msg_type) msg_types;
static struct msg_type *msg_type_lookup(u_int8_t);
static int msg_recv_packet(struct netbuf *);
static int msg_recv_data(struct netbuf *);
@ -41,9 +39,18 @@ static void msg_type_shutdown(struct kore_msg *, const void *);
static void msg_type_websocket(struct kore_msg *, const void *);
#endif
static TAILQ_HEAD(, msg_type) msg_types;
static int cacheidx = 0;
static struct connection *conncache[KORE_WORKER_MAX];
void
kore_msg_init(void)
{
int i;
for (i = 0; i < KORE_WORKER_MAX; i++)
conncache[i] = NULL;
TAILQ_INIT(&msg_types);
}
@ -76,6 +83,11 @@ kore_msg_parent_add(struct kore_worker *kw)
kw->msg[0]->disconnect = msg_disconnected_worker;
kw->msg[0]->handle = kore_connection_handle;
if (cacheidx >= KORE_WORKER_MAX)
fatal("%s: too many workers", __func__);
conncache[cacheidx++] = kw->msg[0];
TAILQ_INSERT_TAIL(&connections, kw->msg[0], list);
kore_platform_event_all(kw->msg[0]->fd, kw->msg[0]);
@ -114,6 +126,18 @@ kore_msg_worker_init(void)
sizeof(struct kore_msg), 0, msg_recv_packet);
}
void
kore_msg_unregister(u_int8_t id)
{
struct msg_type *type;
if ((type = msg_type_lookup(id)) == NULL)
return;
TAILQ_REMOVE(&msg_types, type, list);
kore_free(type);
}
int
kore_msg_register(u_int8_t id, void (*cb)(struct kore_msg *, const void *))
{
@ -164,7 +188,9 @@ static int
msg_recv_data(struct netbuf *nb)
{
struct connection *c;
u_int8_t dst;
struct msg_type *type;
int deliver, i;
u_int16_t destination;
struct kore_msg *msg = (struct kore_msg *)nb->buf;
@ -182,12 +208,24 @@ msg_recv_data(struct netbuf *nb)
if (worker == NULL && type == NULL) {
destination = msg->dst;
TAILQ_FOREACH(c, &connections, list) {
for (i = 0; conncache[i] != NULL; i++) {
c = conncache[i];
if (c->proto != CONN_PROTO_MSG || c->hdlr_extra == NULL)
continue;
if (destination != KORE_MSG_WORKER_ALL &&
*(u_int8_t *)c->hdlr_extra != destination)
deliver = 1;
dst = *(u_int8_t *)c->hdlr_extra;
if (destination == KORE_MSG_WORKER_ALL) {
if (keymgr_active && dst == 0)
deliver = 0;
} else {
if (dst != destination)
deliver = 0;
}
if (deliver == 0)
continue;
/* This allows the worker to receive the correct id. */

View File

@ -56,6 +56,7 @@ static PyObject *python_import(const char *);
static PyObject *pyconnection_alloc(struct connection *);
static PyObject *python_callable(PyObject *, const char *);
static void python_split_arguments(char *, char **, size_t);
static void python_kore_recvobj(struct kore_msg *, const void *);
static const char *python_string_from_dict(PyObject *, const char *);
static int python_bool_from_dict(PyObject *, const char *, int *);
@ -255,9 +256,13 @@ static struct coro_list coro_suspended;
extern const char *__progname;
static PyObject *pickle = NULL;
static PyObject *pickle_dumps = NULL;
static PyObject *pickle_loads = NULL;
static PyObject *python_tracer = NULL;
/* XXX */
static struct python_coro *coro_running = NULL;
static PyObject *python_tracer = NULL;
#if !defined(KORE_SINGLE_BINARY)
const char *kore_pymodule = NULL;
@ -296,6 +301,8 @@ kore_python_init(void)
PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &allocator);
PyMem_SetupDebugHooks();
kore_msg_register(KORE_PYTHON_SEND_OBJ, python_kore_recvobj);
if (PyImport_AppendInittab("kore", &python_module_init) == -1)
fatal("kore_python_init: failed to add new module");
@ -307,6 +314,15 @@ kore_python_init(void)
Py_InitializeEx(0);
if ((pickle = PyImport_ImportModule("pickle")) == NULL)
fatal("failed to import pickle module");
if ((pickle_dumps = PyObject_GetAttrString(pickle, "dumps")) == NULL)
fatal("pickle module has no dumps method");
if ((pickle_loads = PyObject_GetAttrString(pickle, "loads")) == NULL)
fatal("pickle module has no loads method");
#if defined(__linux__)
kore_seccomp_filter("python", filter_python,
KORE_FILTER_LEN(filter_python));
@ -2055,6 +2071,83 @@ python_kore_setname(PyObject *self, PyObject *args)
Py_RETURN_NONE;
}
static PyObject *
python_kore_sendobj(PyObject *self, PyObject *args, PyObject *kwargs)
{
long val;
u_int16_t dst;
char *ptr;
Py_ssize_t length;
PyObject *object, *bytes;
if (!PyArg_ParseTuple(args, "O", &object))
return (NULL);
bytes = PyObject_CallFunctionObjArgs(pickle_dumps, object, NULL);
if (bytes == NULL)
return (NULL);
if (PyBytes_AsStringAndSize(bytes, &ptr, &length) == -1) {
Py_DECREF(bytes);
return (NULL);
}
dst = KORE_MSG_WORKER_ALL;
if (kwargs != NULL) {
if (python_long_from_dict(kwargs, "worker", &val)) {
if (val <= 0 || val > worker_count ||
val >= KORE_WORKER_MAX) {
PyErr_Format(PyExc_RuntimeError,
"worker %ld invalid", val);
Py_DECREF(bytes);
return (NULL);
}
dst = val;
}
}
kore_msg_send(dst, KORE_PYTHON_SEND_OBJ, ptr, length);
Py_DECREF(bytes);
Py_RETURN_NONE;
}
static void
python_kore_recvobj(struct kore_msg *msg, const void *data)
{
struct kore_runtime *rt;
PyObject *onmsg, *ret, *bytes, *obj;
if ((onmsg = kore_module_getsym("koreapp.onmsg", &rt)) == NULL)
return;
if (rt->type != KORE_RUNTIME_PYTHON)
return;
if ((bytes = PyBytes_FromStringAndSize(data, msg->length)) == NULL) {
Py_DECREF(onmsg);
kore_python_log_error("kore.recvobj");
return;
}
obj = PyObject_CallFunctionObjArgs(pickle_loads, bytes, NULL);
Py_DECREF(bytes);
if (obj == NULL) {
Py_DECREF(onmsg);
kore_python_log_error("kore.recvobj");
return;
}
ret = PyObject_CallFunctionObjArgs(onmsg, obj, NULL);
Py_DECREF(obj);
Py_DECREF(onmsg);
Py_XDECREF(ret);
}
static PyObject *
python_kore_suspend(PyObject *self, PyObject *args)
{