fix(async): slightly better support for nested coroutine contexts (#600)

* fix(async): only dispatch first failure in a.wait_all

* add InstallContext:run_concurrently

This is needed due to how multiple coroutine contexts are used in a hierchical
structure, and the async coroutine dispatcher loses this hierchical context
inside callback functions invoked via FFI (I… assume?).

6 files changed, 94 insertions, 21 deletions

diff --git a/lua/nvim-lsp-installer/core/async/init.lua b/lua/nvim-lsp-installer/core/async/init.lua
index 8537b909..27a4c7cb 100644
--- a/lua/nvim-lsp-installer/core/async/init.lua
+++ b/lua/nvim-lsp-installer/core/async/init.lua
@@ -147,6 +147,9 @@ local function oneshot_channel()
     local saved_callback
 
     return {
+        is_closed = function()
+            return has_sent
+        end,
         send = function(...)
             assert(not has_sent, "Oneshot channel can only send once.")
             has_sent = true
@@ -174,32 +177,31 @@ exports.wait_all = function(suspend_fns)
 
     do
         local results = {}
-        local threads = {}
+        local thread_cancellations = {}
         local count = #suspend_fns
         local completed = 0
 
-        local function callback(i)
-            return function(success, result)
+        for i, suspend_fn in ipairs(suspend_fns) do
+            thread_cancellations[i] = exports.run(suspend_fn, function(success, result)
+                completed = completed + 1
                 if not success then
-                    for _, cancel_thread in ipairs(threads) do
-                        cancel_thread()
+                    if not channel.is_closed() then
+                        for _, cancel_thread in ipairs(thread_cancellations) do
+                            cancel_thread()
+                        end
+                        channel.send(false, result)
+                        results = nil
+                        thread_cancellations = {}
                     end
-                    channel.send(false, result)
-                    results = nil
-                    threads = nil
                 else
                     results[i] = result
-                    completed = completed + 1
                     if completed >= count then
                         channel.send(true, results)
                         results = nil
-                        threads = nil
+                        thread_cancellations = {}
                     end
                 end
-            end
-        end
-        for i, suspend_fn in ipairs(suspend_fns) do
-            threads[i] = exports.run(suspend_fn, callback(i))
+            end)
         end
     end
 
diff --git a/lua/nvim-lsp-installer/core/installer/context.lua b/lua/nvim-lsp-installer/core/installer/context.lua
index 5705fa41..f97f334f 100644
--- a/lua/nvim-lsp-installer/core/installer/context.lua
+++ b/lua/nvim-lsp-installer/core/installer/context.lua
@@ -4,6 +4,8 @@ local fs = require "nvim-lsp-installer.core.fs"
 local path = require "nvim-lsp-installer.path"
 local platform = require "nvim-lsp-installer.platform"
 local receipt = require "nvim-lsp-installer.core.receipt"
+local installer = require "nvim-lsp-installer.core.installer"
+local a = require "nvim-lsp-installer.core.async"
 
 ---@class ContextualSpawn
 ---@field cwd CwdManager
@@ -139,4 +141,17 @@ function InstallContext:promote_cwd()
     self.cwd:set(self.destination_dir)
 end
 
+---Runs the provided async functions concurrently and returns their result, once all are resolved.
+---This is really just a wrapper around a.wait_all() that makes sure to patch the coroutine context before creating the
+---new async execution contexts.
+---@async
+---@param suspend_fns async fun(ctx: InstallContext)[]
+function InstallContext:run_concurrently(suspend_fns)
+    return a.wait_all(vim.tbl_map(function(suspend_fn)
+        return function()
+            return installer.run_installer(self, suspend_fn)
+        end
+    end, suspend_fns))
+end
+
 return InstallContext
diff --git a/lua/nvim-lsp-installer/core/installer/init.lua b/lua/nvim-lsp-installer/core/installer/init.lua
index dc9143a2..888415f2 100644
--- a/lua/nvim-lsp-installer/core/installer/init.lua
+++ b/lua/nvim-lsp-installer/core/installer/init.lua
@@ -34,6 +34,7 @@ end
 function M.run_installer(context, installer)
     local thread = coroutine.create(installer)
     local step
+    local ret_val
     step = function(...)
         local ok, result = coroutine.resume(thread, ...)
         if not ok then
@@ -43,9 +44,12 @@ function M.run_installer(context, installer)
         elseif coroutine.status(thread) == "suspended" then
             -- yield to parent coroutine
             step(coroutine.yield(result))
+        else
+            ret_val = result
         end
     end
     step(context)
+    return ret_val
 end
 
 ---@async
diff --git a/lua/nvim-lsp-installer/core/managers/powershell/init.lua b/lua/nvim-lsp-installer/core/managers/powershell/init.lua
index a246e24c..94046b05 100644
--- a/lua/nvim-lsp-installer/core/managers/powershell/init.lua
+++ b/lua/nvim-lsp-installer/core/managers/powershell/init.lua
@@ -11,6 +11,7 @@ local PWSHOPT = {
 ---@param script string
 ---@param opts JobSpawnOpts
 ---@param custom_spawn JobSpawn
+---@return Result
 function M.script(script, opts, custom_spawn)
     opts = opts or {}
     ---@type JobSpawn
@@ -31,6 +32,7 @@ end
 ---@param command string
 ---@param opts JobSpawnOpts
 ---@param custom_spawn JobSpawn
+---@return Result
 function M.command(command, opts, custom_spawn)
     opts = opts or {}
     ---@type JobSpawn
diff --git a/tests/core/async/async_spec.lua b/tests/core/async/async_spec.lua
index 9795adce..c0e8992e 100644
--- a/tests/core/async/async_spec.lua
+++ b/tests/core/async/async_spec.lua
@@ -136,22 +136,28 @@ describe("async", function()
         "should run all suspending functions concurrently",
         async_test(function()
             local start = timestamp()
-            local function sleep(ms)
+            local function sleep(ms, ret_val)
                 return function()
                     a.sleep(ms)
+                    return ret_val
                 end
             end
-            a.wait_all {
-                sleep(100),
-                sleep(100),
-                sleep(100),
-                sleep(100),
-                sleep(100),
+            local one, two, three, four, five = a.wait_all {
+                sleep(100, 1),
+                sleep(100, "two"),
+                sleep(100, "three"),
+                sleep(100, 4),
+                sleep(100, 5),
             }
             local grace = 50
             local delta = timestamp() - start
             assert.is_true(delta <= (100 + grace))
             assert.is_true(delta >= (100 - grace))
+            assert.equals(1, one)
+            assert.equals("two", two)
+            assert.equals("three", three)
+            assert.equals(4, four)
+            assert.equals(5, five)
         end)
     )
 end)
diff --git a/tests/core/installer_spec.lua b/tests/core/installer_spec.lua
index 75911416..cf8e8795 100644
--- a/tests/core/installer_spec.lua
+++ b/tests/core/installer_spec.lua
@@ -1,11 +1,17 @@
 local spy = require "luassert.spy"
 local match = require "luassert.match"
 local fs = require "nvim-lsp-installer.core.fs"
+local a = require "nvim-lsp-installer.core.async"
 local installer = require "nvim-lsp-installer.core.installer"
 local InstallContext = require "nvim-lsp-installer.core.installer.context"
 local process = require "nvim-lsp-installer.process"
 local Optional = require "nvim-lsp-installer.core.optional"
 
+local function timestamp()
+    local seconds, microseconds = vim.loop.gettimeofday()
+    return (seconds * 1000) + math.floor(microseconds / 1000)
+end
+
 describe("installer", function()
     it(
         "should call installer",
@@ -89,4 +95,42 @@ describe("installer", function()
             )
         end)
     )
+
+    it(
+        "should run async functions concurrently",
+        async_test(function()
+            spy.on(fs, "write_file")
+            local destination_dir = ("%s/installer_spec_concurrent"):format(os.getenv "INSTALL_ROOT_DIR")
+            local ctx = InstallContext.new {
+                name = "installer_spec_receipt",
+                destination_dir = destination_dir,
+                boundary_path = os.getenv "INSTALL_ROOT_DIR",
+                stdio_sink = process.empty_sink(),
+                requested_version = Optional.empty(),
+            }
+            local capture = spy.new()
+            local start = timestamp()
+            installer.run_installer(ctx, function(c)
+                capture(c:run_concurrently {
+                    function()
+                        a.sleep(100)
+                        return installer.context()
+                    end,
+                    function()
+                        a.sleep(100)
+                        return "two"
+                    end,
+                    function()
+                        a.sleep(100)
+                        return "three"
+                    end,
+                })
+                c.receipt:with_primary_source(c.receipt.npm "my-pkg")
+            end)
+            local stop = timestamp()
+            local grace_ms = 25
+            assert.is_true((stop - start) >= (100 - grace_ms))
+            assert.spy(capture).was_called_with(match.is_ref(ctx), "two", "three")
+        end)
+    )
 end)