save and restore envs for scheduler

3 files changed, 87 insertions, 15 deletions

diff --git a/core/src/xmake/os/setenv.c b/core/src/xmake/os/setenv.c
index 8240c872b..82a5fd48c 100644
--- a/core/src/xmake/os/setenv.c
+++ b/core/src/xmake/os/setenv.c
@@ -47,7 +47,5 @@ tb_int_t xm_os_setenv(lua_State* lua)
 
     // set it
     lua_pushboolean(lua, value? tb_environment_set(name, value) : tb_false);
-
-    // ok
     return 1;
 }
diff --git a/xmake/core/base/os.lua b/xmake/core/base/os.lua
index 44b373a66..e294fbcdb 100644
--- a/xmake/core/base/os.lua
+++ b/xmake/core/base/os.lua
@@ -181,6 +181,11 @@ function os._ramdir()
     return ramdir_root or nil
 end
 
+-- set on change environments callback for scheduler
+function os._sched_chenvs_set(envs)
+    os._SCHED_CHENVS = envs
+end
+
 -- set on change directory callback for scheduler
 function os._sched_chdir_set(chdir)
     os._SCHED_CHDIR = chdir
@@ -480,7 +485,7 @@ function os.cd(dir)
 
     -- do chdir callback for scheduler
     if os._SCHED_CHDIR then
-        os._SCHED_CHDIR(oldir, os.curdir())
+        os._SCHED_CHDIR(os.curdir())
     end
 
     -- ok
@@ -1001,42 +1006,63 @@ end
 -- set all current environment variables
 function os.setenvs(envs)
     if envs then
+        local changed = false
         -- remove new added values
         local curenvs = os.getenvs()
         for name, _ in pairs(curenvs) do
             if not envs[name] then
-                os.setenv(name, nil)
+                os._setenv(name, "")
+                changed = true
             end
         end
         -- change values
         for name, values in pairs(envs) do
-            os.setenv(name, values)
+            if curenvs[name] ~= values then
+                os._setenv(name, values)
+                changed = true
+            end
+        end
+        -- update envs for scheduler
+        if changed and os._SCHED_CHENVS then
+            os._SCHED_CHENVS(envs)
         end
     end
 end
 
 -- set values to environment variable
 function os.setenv(name, ...)
+    local ok
     local values = {...}
     if #values <= 1 then
         -- keep compatible with original implementation
-        return os._setenv(name, values[1] or "")
+        ok = os._setenv(name, values[1] or "")
     else
-        return os._setenv(name, path.joinenv(values))
+        ok = os._setenv(name, path.joinenv(values))
+    end
+    -- update envs for scheduler
+    if ok and os._SCHED_CHENVS then
+        os._SCHED_CHENVS()
     end
+    return ok
 end
 
 -- add values to environment variable
 function os.addenv(name, ...)
     local values = {...}
     if #values > 0 then
+        local ok
         local oldenv = os.getenv(name)
         local appendenv = path.joinenv(values)
         if oldenv == "" or oldenv == nil then
-            return os._setenv(name, appendenv)
+            ok = os._setenv(name, appendenv)
         else
-            return os._setenv(name, appendenv .. path.envsep() .. oldenv)
+            ok = os._setenv(name, appendenv .. path.envsep() .. oldenv)
+        end
+        -- update envs for scheduler
+        if ok and os._SCHED_CHENVS then
+            os._SCHED_CHENVS()
         end
+        return ok
     else
         return true
     end
@@ -1045,7 +1071,12 @@ end
 -- set values to environment variable with the given seperator
 function os.setenvp(name, values, sep)
     sep = sep or path.envsep()
-    return os._setenv(name, table.concat(table.wrap(values), sep))
+    local ok = os._setenv(name, table.concat(table.wrap(values), sep))
+    -- update envs for scheduler
+    if ok and os._SCHED_CHENVS then
+        os._SCHED_CHENVS()
+    end
+    return ok
 end
 
 -- add values to environment variable with the given seperator
@@ -1053,13 +1084,19 @@ function os.addenvp(name, values, sep)
     sep = sep or path.envsep()
     values = table.wrap(values)
     if #values > 0 then
+        local ok
         local oldenv = os.getenv(name)
         local appendenv = table.concat(values, sep)
         if oldenv == "" or oldenv == nil then
-            return os._setenv(name, appendenv)
+            ok = os._setenv(name, appendenv)
         else
-            return os._setenv(name, appendenv .. sep .. oldenv)
+            ok = os._setenv(name, appendenv .. sep .. oldenv)
+        end
+        -- update envs for scheduler
+        if ok and os._SCHED_CHENVS then
+            os._SCHED_CHENVS()
         end
+        return ok
     else
         return true
     end
diff --git a/xmake/core/base/scheduler.lua b/xmake/core/base/scheduler.lua
index 20f19adf9..bfdde4b83 100644
--- a/xmake/core/base/scheduler.lua
+++ b/xmake/core/base/scheduler.lua
@@ -241,7 +241,7 @@ function scheduler:_co_curdir_update(curdir)
     -- save the current directory hash
     curdir = curdir or os.curdir()
     local curdir_hash = hash.uuid4(path.absolute(curdir)):sub(1, 8)
-    self._CO_CURDIR = curdir_hash
+    self._CO_CURDIR_HASH = curdir_hash
 
     -- save the current directory for each coroutine
     local running = self:co_running()
@@ -255,6 +255,30 @@ function scheduler:_co_curdir_update(curdir)
     end
 end
 
+-- update the current environments hash of current coroutine
+function scheduler:_co_curenvs_update(envs)
+
+    -- save the current directory hash
+    local envs_hash = ""
+    envs = envs or os.getenvs()
+    for _, key in ipairs(table.orderkeys(envs)) do
+        envs_hash = envs_hash .. key:upper() .. envs[key]
+    end
+    envs_hash = hash.uuid4(envs_hash):sub(1, 8)
+    self._CO_CURENVS_HASH = envs_hash
+
+    -- save the current directory for each coroutine
+    local running = self:co_running()
+    if running then
+        local co_curenvs = self._CO_CURENVS
+        if not co_curenvs then
+            co_curenvs = {}
+            self._CO_CURENVS = co_curenvs
+        end
+        co_curenvs[running] = {envs_hash, envs}
+    end
+end
+
 -- resume it's waiting coroutine if all coroutines are dead in group
 function scheduler:_co_groups_resume()
 
@@ -314,6 +338,7 @@ function scheduler:co_start_named(coname, cotask, ...)
     local co
     co = _coroutine.new(coname, coroutine.create(function(...)
         self:_co_curdir_update()
+        self:_co_curenvs_update()
         cotask(...)
         self:co_tasks()[co:thread()] = nil
         if self:co_count() > 0 then
@@ -355,12 +380,19 @@ function scheduler:co_suspend(...)
 
     -- if the current directory has been changed? restore it
     local running = assert(self:co_running())
-    local curdir = self._CO_CURDIR
+    local curdir = self._CO_CURDIR_HASH
     local olddir = self._CO_CURDIRS and self._CO_CURDIRS[running] or nil
     if olddir and curdir ~= olddir[1] then -- hash changed?
         os.cd(olddir[2])
     end
 
+    -- if the current environments has been changed? restore it
+    local curenvs = self._CO_CURENVS_HASH
+    local oldenvs = self._CO_CURENVS and self._CO_CURENVS[running] or nil
+    if oldenvs and curenvs ~= oldenvs[1] then -- hash changed?
+        os.setenvs(oldenvs[2])
+    end
+
     -- return results
     return table.unpack(results)
 end
@@ -750,10 +782,15 @@ function scheduler:runloop()
     end
 
     -- set on change directory callback for scheduler
-    os._sched_chdir_set(function (oldir, curdir)
+    os._sched_chdir_set(function (curdir)
         self:_co_curdir_update(curdir)
     end)
 
+    -- set on change environments callback for scheduler
+    os._sched_chenvs_set(function (envs)
+        self:_co_curenvs_update(envs)
+    end)
+
     -- start all ready coroutine tasks
     local co_ready_tasks = self._CO_READY_TASKS
     if co_ready_tasks then