1 files changed, 54 insertions, 2 deletions

diff --git a/proxy_await.c b/proxy_await.c
index 1486068..5e379db 100644
--- a/proxy_await.c
+++ b/proxy_await.c
@@ -50,6 +50,7 @@ int mcplib_await(lua_State *L) {
             case AWAIT_OK:
             case AWAIT_FIRST:
             case AWAIT_FASTGOOD:
+            case AWAIT_BACKGROUND:
                 break;
             default:
                 proxy_lua_error(L, "invalid type argument tp mcp.await");
@@ -176,6 +177,40 @@ static void mcp_queue_await_io(conn *c, lua_State *Lc, mcp_request_t *rq, int aw
     return;
 }
 
+static void mcp_queue_await_dummy_io(conn *c, lua_State *Lc, int await_ref) {
+    io_queue_t *q = conn_io_queue_get(c, IO_QUEUE_PROXY);
+
+    io_pending_proxy_t *p = do_cache_alloc(c->thread->io_cache);
+    if (p == NULL) {
+        WSTAT_INCR(c, proxy_conn_oom, 1);
+        proxy_lua_error(Lc, "out of memory allocating from IO cache");
+        return;
+    }
+
+    // this is a re-cast structure, so assert that we never outsize it.
+    assert(sizeof(io_pending_t) >= sizeof(io_pending_proxy_t));
+    memset(p, 0, sizeof(io_pending_proxy_t));
+    // set up back references.
+    p->io_queue_type = IO_QUEUE_PROXY;
+    p->thread = c->thread;
+    p->c = c;
+    p->resp = NULL;
+
+    // await specific
+    p->is_await = true;
+    p->await_ref = await_ref;
+    p->await_background = true;
+
+    // Dummy IO has no backend, and no request attached.
+
+    // All we need to do is link into the batch chain.
+    p->next = q->stack_ctx;
+    q->stack_ctx = p;
+    P_DEBUG("%s: queued\n", __func__);
+
+    return;
+}
+
 // TODO (v2): need to get this code running under pcall().
 // It looks like a bulk of this code can move into mcplib_await(),
 // and then here post-yield we can add the conn and coro_ref to the right
@@ -224,6 +259,12 @@ int mcplib_await_run(conn *c, mc_resp *resp, lua_State *L, int coro_ref) {
     }
     P_DEBUG("%s: argtable len: %d\n", __func__, n);
 
+    if (aw->type == AWAIT_BACKGROUND) {
+        mcp_queue_await_dummy_io(c, L, await_ref);
+        aw->pending++;
+        aw->wait_for = 0;
+    }
+
     lua_pop(L, 1); // remove table key.
     aw->resp = resp; // cuddle the current mc_resp to fill later
 
@@ -262,7 +303,13 @@ int mcplib_await_return(io_pending_proxy_t *p) {
     // TODO (v2): for GOOD or OK cases, it might be better to return the
     // last object as valid if there are otherwise zero valids?
     // Think we just have to count valids...
-    if (!aw->completed) {
+    if (aw->type == AWAIT_BACKGROUND) {
+        // in the background case, we never want to collect responses.
+        if (p->await_background) {
+            // found the dummy IO, complete and return conn to worker.
+            completing = true;
+        }
+    } else if (!aw->completed) {
         valid = true; // always collect results unless we are completed.
         if (aw->wait_for > 0) {
             bool is_good = false;
@@ -298,6 +345,9 @@ int mcplib_await_return(io_pending_proxy_t *p) {
                         }
                     }
                     break;
+                case AWAIT_BACKGROUND:
+                    // In background mode we don't wait for any response.
+                    break;
             }
 
             if (is_good) {
@@ -335,7 +385,9 @@ int mcplib_await_return(io_pending_proxy_t *p) {
     }
 
     // lose our internal mcpres reference regardless.
-    luaL_unref(L, LUA_REGISTRYINDEX, p->mcpres_ref);
+    if (p->mcpres_ref) {
+        luaL_unref(L, LUA_REGISTRYINDEX, p->mcpres_ref);
+    }
     // our await_ref is shared, so we don't need to release it.
 
     if (completing) {