Python: support for multiple targets.

6 files changed, 250 insertions, 70 deletions

diff --git a/src/python/nxt_python.c b/src/python/nxt_python.c
index d8204937..588a147a 100644
--- a/src/python/nxt_python.c
+++ b/src/python/nxt_python.c
@@ -24,6 +24,8 @@ typedef struct {
 
 static nxt_int_t nxt_python_start(nxt_task_t *task,
     nxt_process_data_t *data);
+static nxt_int_t nxt_python_set_target(nxt_task_t *task,
+    nxt_python_target_t *target, nxt_conf_value_t *conf);
 static nxt_int_t nxt_python_set_path(nxt_task_t *task, nxt_conf_value_t *value);
 static int nxt_python_init_threads(nxt_python_app_conf_t *c);
 static int nxt_python_ready_handler(nxt_unit_ctx_t *ctx);
@@ -49,7 +51,7 @@ NXT_EXPORT nxt_app_module_t  nxt_app_module = {
 };
 
 static PyObject           *nxt_py_stderr_flush;
-PyObject                  *nxt_py_application;
+nxt_python_targets_t      *nxt_py_targets;
 
 #if PY_MAJOR_VERSION == 3
 static wchar_t            *nxt_py_home;
@@ -66,18 +68,19 @@ static nxt_int_t
 nxt_python_start(nxt_task_t *task, nxt_process_data_t *data)
 {
     int                    rc;
-    char                   *nxt_py_module;
-    size_t                 len;
+    size_t                 len, size;
+    uint32_t               next;
     PyObject               *obj, *module;
-    nxt_str_t              proto;
-    const char             *callable;
+    nxt_str_t              proto, probe_proto, name;
+    nxt_int_t              ret, n, i;
     nxt_unit_ctx_t         *unit_ctx;
     nxt_unit_init_t        python_init;
+    nxt_conf_value_t       *cv;
+    nxt_python_targets_t   *targets;
     nxt_common_app_conf_t  *app_conf;
     nxt_python_app_conf_t  *c;
 #if PY_MAJOR_VERSION == 3
     char                   *path;
-    size_t                 size;
     nxt_int_t              pep405;
 
     static const char pyvenv[] = "/pyvenv.cfg";
@@ -190,38 +193,42 @@ nxt_python_start(nxt_task_t *task, nxt_process_data_t *data)
 
     Py_CLEAR(obj);
 
-    nxt_py_module = nxt_alloca(c->module.length + 1);
-    nxt_memcpy(nxt_py_module, c->module.start, c->module.length);
-    nxt_py_module[c->module.length] = '\0';
+    n = (c->targets != NULL ? nxt_conf_object_members_count(c->targets) : 1);
 
-    module = PyImport_ImportModule(nxt_py_module);
-    if (nxt_slow_path(module == NULL)) {
-        nxt_alert(task, "Python failed to import module \"%s\"", nxt_py_module);
-        nxt_python_print_exception();
+    size = sizeof(nxt_python_targets_t) + n * sizeof(nxt_python_target_t);
+
+    targets = nxt_unit_malloc(NULL, size);
+    if (nxt_slow_path(targets == NULL)) {
+        nxt_alert(task, "Could not allocate targets");
         goto fail;
     }
 
-    callable = (c->callable != NULL) ? c->callable : "application";
+    memset(targets, 0, size);
 
-    obj = PyDict_GetItemString(PyModule_GetDict(module), callable);
-    if (nxt_slow_path(obj == NULL)) {
-        nxt_alert(task, "Python failed to get \"%s\" "
-                  "from module \"%s\"", callable, nxt_py_module);
-        goto fail;
-    }
+    targets->count = n;
+    nxt_py_targets = targets;
 
-    if (nxt_slow_path(PyCallable_Check(obj) == 0)) {
-        nxt_alert(task, "\"%s\" in module \"%s\" "
-                  "is not a callable object", callable, nxt_py_module);
-        goto fail;
-    }
+    if (c->targets != NULL) {
+        next = 0;
 
-    nxt_py_application = obj;
-    obj = NULL;
+        for (i = 0; /* void */; i++) {
+            cv = nxt_conf_next_object_member(c->targets, &name, &next);
+            if (cv == NULL) {
+                break;
+            }
 
-    Py_INCREF(nxt_py_application);
+            ret = nxt_python_set_target(task, &targets->target[i], cv);
+            if (nxt_slow_path(ret != NXT_OK)) {
+                goto fail;
+            }
+        }
 
-    Py_CLEAR(module);
+    } else {
+        ret = nxt_python_set_target(task, &targets->target[0], app_conf->self);
+        if (nxt_slow_path(ret != NXT_OK)) {
+            goto fail;
+        }
+    }
 
     nxt_unit_default_init(task, &python_init);
 
@@ -232,7 +239,18 @@ nxt_python_start(nxt_task_t *task, nxt_process_data_t *data)
     proto = c->protocol;
 
     if (proto.length == 0) {
-        proto = nxt_python_asgi_check(nxt_py_application) ? asgi : wsgi;
+        proto = nxt_python_asgi_check(targets->target[0].application)
+                ? asgi : wsgi;
+
+        for (i = 1; i < targets->count; i++) {
+            probe_proto = nxt_python_asgi_check(targets->target[i].application)
+                          ? asgi : wsgi;
+            if (probe_proto.start != proto.start) {
+                nxt_alert(task, "A mix of ASGI & WSGI targets is forbidden, "
+                                "specify protocol in config if incorrect");
+                goto fail;
+            }
+        }
     }
 
     if (nxt_strstr_eq(&proto, &asgi)) {
@@ -299,6 +317,81 @@ fail:
 
 
 static nxt_int_t
+nxt_python_set_target(nxt_task_t *task, nxt_python_target_t *target,
+    nxt_conf_value_t *conf)
+{
+    char              *callable, *module_name;
+    PyObject          *module, *obj;
+    nxt_str_t         str;
+    nxt_conf_value_t  *value;
+
+    static nxt_str_t  module_str = nxt_string("module");
+    static nxt_str_t  callable_str = nxt_string("callable");
+
+    module = obj = NULL;
+
+    value = nxt_conf_get_object_member(conf, &module_str, NULL);
+    if (nxt_slow_path(value == NULL)) {
+        goto fail;
+    }
+
+    nxt_conf_get_string(value, &str);
+
+    module_name = nxt_alloca(str.length + 1);
+    nxt_memcpy(module_name, str.start, str.length);
+    module_name[str.length] = '\0';
+
+    module = PyImport_ImportModule(module_name);
+    if (nxt_slow_path(module == NULL)) {
+        nxt_alert(task, "Python failed to import module \"%s\"", module_name);
+        nxt_python_print_exception();
+        goto fail;
+    }
+
+    value = nxt_conf_get_object_member(conf, &callable_str, NULL);
+    if (value == NULL) {
+        callable = nxt_alloca(12);
+        nxt_memcpy(callable, "application", 12);
+
+    } else {
+        nxt_conf_get_string(value, &str);
+
+        callable = nxt_alloca(str.length + 1);
+        nxt_memcpy(callable, str.start, str.length);
+        callable[str.length] = '\0';
+    }
+
+    obj = PyDict_GetItemString(PyModule_GetDict(module), callable);
+    if (nxt_slow_path(obj == NULL)) {
+        nxt_alert(task, "Python failed to get \"%s\" from module \"%s\"",
+                  callable, module);
+        goto fail;
+    }
+
+    if (nxt_slow_path(PyCallable_Check(obj) == 0)) {
+        nxt_alert(task, "\"%s\" in module \"%s\" is not a callable object",
+                  callable, module);
+        goto fail;
+    }
+
+    target->application = obj;
+    obj = NULL;
+
+    Py_INCREF(target->application);
+    Py_CLEAR(module);
+
+    return NXT_OK;
+
+fail:
+
+    Py_XDECREF(obj);
+    Py_XDECREF(module);
+
+    return NXT_ERROR;
+}
+
+
+static nxt_int_t
 nxt_python_set_path(nxt_task_t *task, nxt_conf_value_t *value)
 {
     int               ret;
@@ -596,12 +689,21 @@ nxt_python_done_strings(nxt_python_string_t *pstr)
 static void
 nxt_python_atexit(void)
 {
+    nxt_int_t  i;
+
     if (nxt_py_proto.done != NULL) {
         nxt_py_proto.done();
     }
 
     Py_XDECREF(nxt_py_stderr_flush);
-    Py_XDECREF(nxt_py_application);
+
+    if (nxt_py_targets != NULL) {
+        for (i = 0; i < nxt_py_targets->count; i++) {
+            Py_XDECREF(nxt_py_targets->target[i].application);
+        }
+
+        nxt_unit_free(NULL, nxt_py_targets);
+    }
 
     Py_Finalize();
 
diff --git a/src/python/nxt_python.h b/src/python/nxt_python.h
index b581dd46..a5c1d9a6 100644
--- a/src/python/nxt_python.h
+++ b/src/python/nxt_python.h
@@ -37,13 +37,28 @@
 #define NXT_HAVE_ASGI  1
 #endif
 
-extern PyObject  *nxt_py_application;
+
+typedef struct {
+    PyObject    *application;
+    nxt_bool_t  asgi_legacy;
+} nxt_python_target_t;
+
+
+typedef struct {
+    nxt_int_t            count;
+    nxt_python_target_t  target[0];
+} nxt_python_targets_t;
+
+
+extern nxt_python_targets_t  *nxt_py_targets;
+
 
 typedef struct {
     nxt_str_t  string;
     PyObject   **object_p;
 } nxt_python_string_t;
 
+
 typedef struct {
     int   (*ctx_data_alloc)(void **pdata);
     void  (*ctx_data_free)(void *data);
diff --git a/src/python/nxt_python_asgi.c b/src/python/nxt_python_asgi.c
index a6f94507..1d220678 100644
--- a/src/python/nxt_python_asgi.c
+++ b/src/python/nxt_python_asgi.c
@@ -43,8 +43,6 @@ static void nxt_py_asgi_shm_ack_handler(nxt_unit_ctx_t *ctx);
 static PyObject *nxt_py_asgi_port_read(PyObject *self, PyObject *args);
 static void nxt_python_asgi_done(void);
 
-
-int                       nxt_py_asgi_legacy;
 static PyObject           *nxt_py_port_read;
 static nxt_unit_port_t    *nxt_py_shared_port;
 
@@ -137,6 +135,7 @@ int
 nxt_python_asgi_init(nxt_unit_init_t *init, nxt_python_proto_t *proto)
 {
     PyObject      *func;
+    nxt_int_t     i;
     PyCodeObject  *code;
 
     nxt_unit_debug(NULL, "asgi_init");
@@ -161,21 +160,23 @@ nxt_python_asgi_init(nxt_unit_init_t *init, nxt_python_proto_t *proto)
         return NXT_UNIT_ERROR;
     }
 
-    func = nxt_python_asgi_get_func(nxt_py_application);
-    if (nxt_slow_path(func == NULL)) {
-        nxt_unit_alert(NULL, "Python cannot find function for callable");
-        return NXT_UNIT_ERROR;
-    }
+    for (i = 0; i < nxt_py_targets->count; i++) {
+        func = nxt_python_asgi_get_func(nxt_py_targets->target[i].application);
+        if (nxt_slow_path(func == NULL)) {
+            nxt_unit_alert(NULL, "Python cannot find function for callable");
+            return NXT_UNIT_ERROR;
+        }
 
-    code = (PyCodeObject *) PyFunction_GET_CODE(func);
+        code = (PyCodeObject *) PyFunction_GET_CODE(func);
 
-    if ((code->co_flags & CO_COROUTINE) == 0) {
-        nxt_unit_debug(NULL, "asgi: callable is not a coroutine function "
-                             "switching to legacy mode");
-        nxt_py_asgi_legacy = 1;
-    }
+        if ((code->co_flags & CO_COROUTINE) == 0) {
+            nxt_unit_debug(NULL, "asgi: callable is not a coroutine function "
+                                 "switching to legacy mode");
+            nxt_py_targets->target[i].asgi_legacy = 1;
+        }
 
-    Py_DECREF(func);
+        Py_DECREF(func);
+    }
 
     init->callbacks.request_handler = nxt_py_asgi_request_handler;
     init->callbacks.data_handler = nxt_py_asgi_http_data_handler;
@@ -408,6 +409,7 @@ nxt_py_asgi_request_handler(nxt_unit_request_info_t *req)
 {
     PyObject                *scope, *res, *task, *receive, *send, *done, *asgi;
     PyObject                *stage2;
+    nxt_python_target_t     *target;
     nxt_py_asgi_ctx_data_t  *ctx_data;
 
     if (req->request->websocket_handshake) {
@@ -456,17 +458,18 @@ nxt_py_asgi_request_handler(nxt_unit_request_info_t *req)
     }
 
     req->data = asgi;
+    target = &nxt_py_targets->target[req->request->app_target];
 
-    if (!nxt_py_asgi_legacy) {
+    if (!target->asgi_legacy) {
         nxt_unit_req_debug(req, "Python call ASGI 3.0 application");
 
-        res = PyObject_CallFunctionObjArgs(nxt_py_application,
+        res = PyObject_CallFunctionObjArgs(target->application,
                                            scope, receive, send, NULL);
 
     } else {
         nxt_unit_req_debug(req, "Python call legacy application");
 
-        res = PyObject_CallFunctionObjArgs(nxt_py_application, scope, NULL);
+        res = PyObject_CallFunctionObjArgs(target->application, scope, NULL);
 
         if (nxt_slow_path(res == NULL)) {
             nxt_unit_req_error(req, "Python failed to call legacy app stage1");
diff --git a/src/python/nxt_python_asgi.h b/src/python/nxt_python_asgi.h
index 37f2a099..20702065 100644
--- a/src/python/nxt_python_asgi.h
+++ b/src/python/nxt_python_asgi.h
@@ -33,7 +33,7 @@ typedef struct {
     PyObject         *loop_remove_reader;
     PyObject         *quit_future;
     PyObject         *quit_future_set_result;
-    PyObject         *lifespan;
+    PyObject         **target_lifespans;
     nxt_unit_port_t  *port;
 } nxt_py_asgi_ctx_data_t;
 
@@ -69,6 +69,4 @@ int nxt_py_asgi_lifespan_startup(nxt_py_asgi_ctx_data_t *ctx_data);
 int nxt_py_asgi_lifespan_shutdown(nxt_unit_ctx_t *ctx);
 
 
-extern int  nxt_py_asgi_legacy;
-
 #endif  /* _NXT_PYTHON_ASGI_H_INCLUDED_ */
diff --git a/src/python/nxt_python_asgi_lifespan.c b/src/python/nxt_python_asgi_lifespan.c
index 506eaf4d..1fc0e6b7 100644
--- a/src/python/nxt_python_asgi_lifespan.c
+++ b/src/python/nxt_python_asgi_lifespan.c
@@ -27,7 +27,10 @@ typedef struct  {
     PyObject                *receive_future;
 } nxt_py_asgi_lifespan_t;
 
-
+static PyObject *nxt_py_asgi_lifespan_target_startup(
+    nxt_py_asgi_ctx_data_t *ctx_data, nxt_python_target_t *target);
+static int nxt_py_asgi_lifespan_target_shutdown(
+    nxt_py_asgi_lifespan_t *lifespan);
 static PyObject *nxt_py_asgi_lifespan_receive(PyObject *self, PyObject *none);
 static PyObject *nxt_py_asgi_lifespan_send(PyObject *self, PyObject *dict);
 static PyObject *nxt_py_asgi_lifespan_send_startup(
@@ -69,24 +72,60 @@ static PyTypeObject nxt_py_asgi_lifespan_type = {
 int
 nxt_py_asgi_lifespan_startup(nxt_py_asgi_ctx_data_t *ctx_data)
 {
-    int                     rc;
+    size_t               size;
+    PyObject             *lifespan;
+    PyObject             **target_lifespans;
+    nxt_int_t            i;
+    nxt_python_target_t  *target;
+
+    size = nxt_py_targets->count * sizeof(PyObject*);
+
+    target_lifespans = nxt_unit_malloc(NULL, size);
+    if (nxt_slow_path(target_lifespans == NULL)) {
+        nxt_unit_alert(NULL, "Failed to allocate lifespan data");
+        return NXT_UNIT_ERROR;
+    }
+
+    memset(target_lifespans, 0, size);
+
+    for (i = 0; i < nxt_py_targets->count; i++) {
+        target = &nxt_py_targets->target[i];
+
+        lifespan = nxt_py_asgi_lifespan_target_startup(ctx_data, target);
+        if (nxt_slow_path(lifespan == NULL)) {
+            return NXT_UNIT_ERROR;
+        }
+
+        target_lifespans[i] = lifespan;
+    }
+
+    ctx_data->target_lifespans = target_lifespans;
+
+    return NXT_UNIT_OK;
+}
+
+
+static PyObject *
+nxt_py_asgi_lifespan_target_startup(nxt_py_asgi_ctx_data_t *ctx_data,
+    nxt_python_target_t *target)
+{
     PyObject                *scope, *res, *py_task, *receive, *send, *done;
     PyObject                *stage2;
-    nxt_py_asgi_lifespan_t  *lifespan;
+    nxt_py_asgi_lifespan_t  *lifespan, *ret;
 
     if (nxt_slow_path(PyType_Ready(&nxt_py_asgi_lifespan_type) != 0)) {
         nxt_unit_alert(NULL,
                  "Python failed to initialize the 'asgi_lifespan' type object");
-        return NXT_UNIT_ERROR;
+        return NULL;
     }
 
     lifespan = PyObject_New(nxt_py_asgi_lifespan_t, &nxt_py_asgi_lifespan_type);
     if (nxt_slow_path(lifespan == NULL)) {
         nxt_unit_alert(NULL, "Python failed to create lifespan object");
-        return NXT_UNIT_ERROR;
+        return NULL;
     }
 
-    rc = NXT_UNIT_ERROR;
+    ret = NULL;
 
     receive = PyObject_GetAttrString((PyObject *) lifespan, "receive");
     if (nxt_slow_path(receive == NULL)) {
@@ -130,23 +169,25 @@ nxt_py_asgi_lifespan_startup(nxt_py_asgi_ctx_data_t *ctx_data)
         goto release_future;
     }
 
-    if (!nxt_py_asgi_legacy) {
+    if (!target->asgi_legacy) {
         nxt_unit_req_debug(NULL, "Python call ASGI 3.0 application");
 
-        res = PyObject_CallFunctionObjArgs(nxt_py_application,
+        res = PyObject_CallFunctionObjArgs(target->application,
                                            scope, receive, send, NULL);
 
     } else {
         nxt_unit_req_debug(NULL, "Python call legacy application");
 
-        res = PyObject_CallFunctionObjArgs(nxt_py_application, scope, NULL);
+        res = PyObject_CallFunctionObjArgs(target->application, scope, NULL);
         if (nxt_slow_path(res == NULL)) {
             nxt_unit_log(NULL, NXT_UNIT_LOG_INFO,
                          "ASGI Lifespan processing exception");
             nxt_python_print_exception();
 
             lifespan->disabled = 1;
-            rc = NXT_UNIT_OK;
+
+            Py_INCREF(lifespan);
+            ret = lifespan;
 
             goto release_scope;
         }
@@ -211,10 +252,9 @@ nxt_py_asgi_lifespan_startup(nxt_py_asgi_ctx_data_t *ctx_data)
     Py_DECREF(res);
 
     if (lifespan->startup_sent == 1 || lifespan->disabled) {
-        ctx_data->lifespan = (PyObject *) lifespan;
-        Py_INCREF(ctx_data->lifespan);
+        Py_INCREF(lifespan);
 
-        rc = NXT_UNIT_OK;
+        ret = lifespan;
     }
 
 release_task:
@@ -232,20 +272,41 @@ release_receive:
 release_lifespan:
     Py_DECREF(lifespan);
 
-    return rc;
+    return (PyObject *) ret;
 }
 
 
 int
 nxt_py_asgi_lifespan_shutdown(nxt_unit_ctx_t *ctx)
 {
-    PyObject                *msg, *future, *res;
+    nxt_int_t               i, ret;
     nxt_py_asgi_lifespan_t  *lifespan;
     nxt_py_asgi_ctx_data_t  *ctx_data;
 
     ctx_data = ctx->data;
 
-    lifespan = (nxt_py_asgi_lifespan_t *) ctx_data->lifespan;
+    for (i = 0; i < nxt_py_targets->count; i++) {
+        lifespan = (nxt_py_asgi_lifespan_t *)ctx_data->target_lifespans[i];
+
+        ret = nxt_py_asgi_lifespan_target_shutdown(lifespan);
+        if (nxt_slow_path(ret != NXT_UNIT_OK)) {
+            return NXT_UNIT_ERROR;
+        }
+    }
+
+    nxt_unit_free(NULL, ctx_data->target_lifespans);
+
+    return NXT_UNIT_OK;
+}
+
+
+static int
+nxt_py_asgi_lifespan_target_shutdown(nxt_py_asgi_lifespan_t *lifespan)
+{
+    PyObject                *msg, *future, *res;
+    nxt_py_asgi_ctx_data_t  *ctx_data;
+
+    ctx_data = lifespan->ctx_data;
 
     if (nxt_slow_path(lifespan == NULL || lifespan->disabled)) {
         return NXT_UNIT_OK;
diff --git a/src/python/nxt_python_wsgi.c b/src/python/nxt_python_wsgi.c
index 77c45af5..b80d10fa 100644
--- a/src/python/nxt_python_wsgi.c
+++ b/src/python/nxt_python_wsgi.c
@@ -302,7 +302,7 @@ nxt_python_request_handler(nxt_unit_request_info_t *req)
 {
     int               rc;
     PyObject          *environ, *args, *response, *iterator, *item;
-    PyObject          *close, *result;
+    PyObject          *close, *result, *application;
     nxt_bool_t        prepare_environ;
     nxt_python_ctx_t  *pctx;
 
@@ -348,7 +348,8 @@ nxt_python_request_handler(nxt_unit_request_info_t *req)
     Py_INCREF(pctx->start_resp);
     PyTuple_SET_ITEM(args, 1, pctx->start_resp);
 
-    response = PyObject_CallObject(nxt_py_application, args);
+    application = nxt_py_targets->target[req->request->app_target].application;
+    response = PyObject_CallObject(application, args);
 
     Py_DECREF(args);