improve co_group_begin

5 files changed, 105 insertions, 40 deletions

diff --git a/tests/modules/process/test.lua b/tests/modules/process/test.lua
index 30a42fd99..3835272b4 100644
--- a/tests/modules/process/test.lua
+++ b/tests/modules/process/test.lua
@@ -25,10 +25,11 @@ function test_sched_process(t)
         stderr:close()
         count = count + 1
     end
-    local cotasks = {}
-    for i = 1, 3 do
-        table.insert(cotasks, scheduler.co_start(_session))
-    end
-    scheduler.co_waitexit(cotasks)
+    scheduler.co_group_begin("test", function ()
+        for i = 1, 3 do
+            scheduler.co_start(_session)
+        end
+    end)
+    scheduler.co_group_wait("test")
     t:are_equal(count, 3)
 end
diff --git a/tests/modules/scheduler/test.lua b/tests/modules/scheduler/test.lua
index d6e090c6a..de62e47ed 100644
--- a/tests/modules/scheduler/test.lua
+++ b/tests/modules/scheduler/test.lua
@@ -7,11 +7,12 @@ function test_runjobs(t)
         t:are_equal(a, "xmake!")
         count = count + 1
     end
-    local cotasks = {}
-    for i = 1, 100 do
-        table.insert(cotasks, scheduler.co_start(task, "xmake!"))
-    end
-    scheduler.co_waitexit(cotasks)
+    scheduler.co_group_begin("test", function ()
+        for i = 1, 100 do
+            scheduler.co_start(task, "xmake!")
+        end
+    end)
+    scheduler.co_group_wait("test")
     t:are_equal(count, 100)
 end
 
@@ -37,10 +38,11 @@ function test_yield(t)
         scheduler.co_yield()
         count = count + 1
     end
-    local cotasks = {}
-    for i = 1, 10 do
-        table.insert(cotasks, scheduler.co_start(task))
-    end
-    scheduler.co_waitexit(cotasks)
+    scheduler.co_group_begin("test", function ()
+        for i = 1, 10 do
+            scheduler.co_start(task)
+        end
+    end)
+    scheduler.co_group_wait("test")
     t:are_equal(count, 10)
 end
diff --git a/tests/modules/socket/sched_tcp/echo_client.lua b/tests/modules/socket/sched_tcp/echo_client.lua
index acc681584..b77487090 100644
--- a/tests/modules/socket/sched_tcp/echo_client.lua
+++ b/tests/modules/socket/sched_tcp/echo_client.lua
@@ -42,20 +42,23 @@ function _session(addr, port)
     if sock then
         print("%s: connected!", sock)
         table.insert(socks, sock)
-        scheduler.co_start(_session_recv, sock)
-        scheduler.co_start(_session_send, sock)
+        scheduler.co_group_begin("test", function ()
+            scheduler.co_start(_session_recv, sock)
+            scheduler.co_start(_session_send, sock)
+        end)
     else
         print("connect %s:%d failed", addr, port)
     end
 end
 
 function main(count)
-    local cotasks = {}
     count = count and tonumber(count) or 1
-    for i = 1, count do
-        table.insert(cotasks, scheduler.co_start(_session, "127.0.0.1", 9001))
-    end
-    scheduler.co_waitexit(cotasks)
+    scheduler.co_group_begin("test", function ()
+        for i = 1, count do
+            scheduler.co_start(_session, "127.0.0.1", 9001)
+        end
+    end)
+    scheduler.co_group_wait("test")
     for _, sock in ipairs(socks) do
         sock:close()
     end
diff --git a/xmake/core/base/scheduler.lua b/xmake/core/base/scheduler.lua
index d53f9363a..12b3eda3f 100644
--- a/xmake/core/base/scheduler.lua
+++ b/xmake/core/base/scheduler.lua
@@ -24,6 +24,7 @@ local _coroutine = _coroutine or {}
 
 -- load modules
 local table     = require("base/table")
+local utils     = require("base/utils")
 local option    = require("base/option")
 local string    = require("base/string")
 local poller    = require("base/poller")
@@ -277,6 +278,14 @@ function scheduler:co_start_named(coname, cotask, ...)
         self._CO_READY_TASKS = self._CO_READY_TASKS or {}
         table.insert(self._CO_READY_TASKS, {co, table.pack(...)})
     end
+
+    -- add this coroutine to the pending groups
+    local co_groups_pending = self._CO_GROUPS_PENDING
+    if co_groups_pending then
+        for _, co_group_pending in pairs(co_groups_pending) do
+            table.insert(co_group_pending, co)
+        end
+    end
     return co
 end
 
@@ -292,7 +301,7 @@ end
 
 -- yield the current coroutine
 function scheduler:co_yield()
-    return scheduler.co_sleep(1)
+    return scheduler.co_sleep(self, 1)
 end
 
 -- sleep some times (ms)
@@ -330,13 +339,48 @@ function scheduler:co_sleep(ms)
     return true
 end
 
--- wait for exiting the given coroutine tasks
-function scheduler:co_waitexit(cotasks)
+-- get the given coroutine group
+function scheduler:co_group(name)
+    return self._CO_GROUPS and self._CO_GROUPS[name]
+end
+
+-- begin coroutine group
+function scheduler:co_group_begin(name, scopefunc)
+    
+    -- enter groups
+    self._CO_GROUPS = self._CO_GROUPS or {}
+    self._CO_GROUPS_PENDING = self._CO_GROUPS_PENDING or {}
+    if self._CO_GROUPS_PENDING[name] then
+        return false, string.format("co_group(%s): already exists!", name)
+    end
+    self._CO_GROUPS_PENDING[name] = self._CO_GROUPS_PENDING[name] or {}
+
+    -- call the scope function
+    local ok, errors = utils.trycall(scopefunc)
+    if not ok then
+        return false, errors
+    end
+
+    -- leave groups
+    self._CO_GROUPS[name] = self._CO_GROUPS[name] or {}
+    table.join2(self._CO_GROUPS[name], self._CO_GROUPS_PENDING[name])
+    self._CO_GROUPS_PENDING[name] = nil
+    return true
+end
+
+-- wait for finishing the given coroutine group 
+function scheduler:co_group_wait(name)
+
+    -- get coroutine group
+    local co_group = self:co_group(name)
+    if not co_group then
+        return false, string.format("co_group(%s): not found!", name)
+    end
 
     -- get the running coroutine
     local running = self:co_running()
     if not running then
-        return false, "we must call waitexit() in coroutine with scheduler!"
+        return false, "we must call co_group_wait() in coroutine with scheduler!"
     end
 
     -- is stopped?
@@ -345,20 +389,22 @@ function scheduler:co_waitexit(cotasks)
     end
 
     -- wait it
-    cotasks = table.copy(table.wrap(cotasks))
-    while #cotasks > 0 do
-        for i = #cotasks, 1, -1 do
-            local co = cotasks[i]
+    local count
+    repeat 
+        count = 0
+        for _, co in ipairs(co_group) do
             if co:is_dead() then
-                table.remove(cotasks, i)
-            else
-                local ok, errors = self:co_sleep(100)
-                if not ok then
-                    return false, errors
-                end
+                count = count + 1
             end
         end
-    end
+        if count ~= #co_group then
+            local ok, errors = scheduler.co_sleep(self, 100)
+            if not ok then
+                return false, errors
+            end
+        end
+    until count == #co_group 
+    self._CO_GROUPS[name] = nil
     return true
 end
 
diff --git a/xmake/core/sandbox/modules/import/core/base/scheduler.lua b/xmake/core/sandbox/modules/import/core/base/scheduler.lua
index 62bfcbb5c..c899fcd34 100644
--- a/xmake/core/sandbox/modules/import/core/base/scheduler.lua
+++ b/xmake/core/sandbox/modules/import/core/base/scheduler.lua
@@ -69,9 +69,22 @@ function sandbox_core_base_scheduler.co_sleep(ms)
     end
 end
 
--- wait for exiting the given coroutine tasks
-function sandbox_core_base_scheduler.co_waitexit(cotasks)
-    local ok, errors = scheduler:co_waitexit(cotasks)
+-- get coroutine group with the given name
+function sandbox_core_base_scheduler.co_group(name)
+    return scheduler:co_group(name)
+end
+
+-- begin coroutine group with the given name
+function sandbox_core_base_scheduler.co_group_begin(name, scopefunc)
+    local ok, errors = scheduler:co_group_begin(name, scopefunc)
+    if not ok then
+        raise(errors)
+    end
+end
+
+-- wait for finishing the given coroutine group
+function sandbox_core_base_scheduler.co_group_wait(name)
+    local ok, errors = scheduler:co_group_wait(name)
     if not ok then
         raise(errors)
     end