From 2eb2cae50391e10170c5489c888d527db55bb8fe Mon Sep 17 00:00:00 2001
From: Appla <bhg@live.it>
Date: Fri, 17 Mar 2023 11:38:33 +0800
Subject: [PATCH] feature: supported wait multi uthreads many times safety for
 ngx.thread.wait

---
 src/ngx_http_lua_uthread.c | 12 +++++----
 t/098-uthread-wait.t       | 51 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/src/ngx_http_lua_uthread.c b/src/ngx_http_lua_uthread.c
index 2208b2a819..427b5c2f51 100644
--- a/src/ngx_http_lua_uthread.c
+++ b/src/ngx_http_lua_uthread.c
@@ -104,7 +104,7 @@ ngx_http_lua_uthread_spawn(lua_State *L)
 static int
 ngx_http_lua_uthread_wait(lua_State *L)
 {
-    int                          i, nargs, nrets;
+    int                          i, nargs, nrets, dead_count;
     lua_State                   *sub_co;
     ngx_http_request_t          *r;
     ngx_http_lua_ctx_t          *ctx;
@@ -129,6 +129,8 @@ ngx_http_lua_uthread_wait(lua_State *L)
         return luaL_error(L, "at least one coroutine should be specified");
     }
 
+    dead_count = 0;
+
     for (i = 1; i <= nargs; i++) {
         sub_co = lua_tothread(L, i);
 
@@ -172,11 +174,11 @@ ngx_http_lua_uthread_wait(lua_State *L)
             return nrets;
 
         case NGX_HTTP_LUA_CO_DEAD:
-            dd("uthread already waited: %p (parent %p)", sub_coctx,
-               coctx);
+            dd("uthread already waited: %p (parent %p), dead/total: %d/%d",
+               sub_coctx, coctx, dead_count + 1, nargs);
 
-            if (i < nargs) {
-                /* just ignore it if it is not the last one */
+            if (++dead_count < nargs) {
+                /* just ignore it if not all threads are dead */
                 continue;
             }
 
diff --git a/t/098-uthread-wait.t b/t/098-uthread-wait.t
index e2818e008e..1ad5e62633 100644
--- a/t/098-uthread-wait.t
+++ b/t/098-uthread-wait.t
@@ -1340,3 +1340,54 @@ GET /lua
 at least one coroutine should be specified
 --- no_error_log
 [crit]
+
+
+
+=== TEST 24: wait any uthreads
+--- config
+    location /lua {
+        content_by_lua '
+            local function f(seconds, tag)
+                -- ngx.say("hello in thread")
+                ngx.sleep(seconds)
+                return tag, seconds
+            end
+            local threads_args = {
+                { 1, 't1 wait 1 seconds'},
+                { 0.2, 't2 wait 0.2 seconds'},
+                { 3, 't3 wait 3 seconds'},
+                { 2, 't4 wait 2 seconds'},
+                { 0.5, 't5 wait 0.5 seconds'},
+            }
+            local threads = {}
+            for i, args in ipairs(threads_args) do
+                local t, err = ngx.thread.spawn(f, args[1], args[2])
+                if not t then
+                    ngx.say("failed to spawn thread: ", err)
+                    break
+                end
+                threads[i] = t
+            end
+
+            for i = 1, #threads + 1 do
+                local ok, res = ngx.thread.wait(unpack(threads))
+                if not ok then
+                    ngx.say("failed to run thread: ", res)
+                    break
+                end
+
+                ngx.say(i, ": ", res)
+            end
+        ';
+    }
+--- request
+GET /lua
+--- response_body
+1: t2 wait 0.2 seconds
+2: t5 wait 0.5 seconds
+3: t1 wait 1 seconds
+4: t4 wait 2 seconds
+5: t3 wait 3 seconds
+failed to run thread: already waited or killed
+--- no_error_log
+[error]
\ No newline at end of file