refactor(registry): parallelize registry installation

5 files changed, 74 insertions, 31 deletions

diff --git a/lua/mason-registry/init.lua b/lua/mason-registry/init.lua
index 5e4d124f..ad8eef5a 100644
--- a/lua/mason-registry/init.lua
+++ b/lua/mason-registry/init.lua
@@ -132,26 +132,40 @@ function Registry.refresh(callback)
     local installer = require "mason-registry.installer"
 
     local state = installer.get_registry_state()
-    local is_state_ok = state
-        and (os.time() - state.timestamp) <= REGISTRY_STORE_TTL
-        and state.checksum == Registry.sources:checksum()
+    if state and Registry.sources:is_all_installed() then
+        local registry_age = os.time() - state.timestamp
 
-    if is_state_ok and Registry.sources:is_all_installed() then
-        log.debug "Registry refresh not necessary."
-        if callback then
-            callback(true, {})
+        if registry_age <= REGISTRY_STORE_TTL and state.checksum == Registry.sources:checksum() then
+            log.fmt_debug(
+                "Registry refresh is not necessary yet. Registry age=%d, checksum=%s",
+                registry_age,
+                state.checksum
+            )
+            if callback then
+                callback(true, {})
+            end
+            return
         end
-        return
     end
 
-    if not callback then
-        return a.run_blocking(function()
-            return a.wait(Registry.update)
+    local function async_update()
+        return a.wait(function(resolve, reject)
+            Registry.update(function(success, result)
+                if success then
+                    resolve(result)
+                else
+                    reject(result)
+                end
+            end)
         end)
+    end
+
+    if not callback then
+        -- We don't want to error in the synchronous version because of how this function is recommended to be used in
+        -- 3rd party code. If accessing the update result is required, users are recommended to pass a callback.
+        pcall(a.run_blocking, async_update)
     else
-        a.run(function()
-            return a.wait(Registry.update)
-        end, callback)
+        a.run(async_update, callback)
     end
 end
 
diff --git a/lua/mason-registry/installer.lua b/lua/mason-registry/installer.lua
index a1984921..5abbcd74 100644
--- a/lua/mason-registry/installer.lua
+++ b/lua/mason-registry/installer.lua
@@ -37,22 +37,38 @@ function M.install(sources)
     log.debug("Installing registries.", sources)
     assert(not M.channel, "Cannot install when channel is active.")
     M.channel = OneShotChannel:new()
-    local result = Result.try(function(try)
-        local updated_sources = {}
-        for source in sources:iterate { include_uninstalled = true } do
-            log.trace("Installing registry.", source)
-            try(source:install():map_err(function(err)
-                return ("%s failed to install: %s"):format(source, err)
-            end))
-            table.insert(updated_sources, source)
-        end
+
+    local results = {
+        a.wait_all(_.map(
+            ---@param source RegistrySource
+            function(source)
+                return function()
+                    log.trace("Installing registry.", source)
+                    return source:install():map(_.always(source)):map_err(function(err)
+                        return ("%s failed to install: %s"):format(source, err)
+                    end)
+                end
+            end,
+            sources:to_list { include_uninstalled = true }
+        )),
+    }
+
+    local any_failed = _.any(Result.is_failure, results)
+
+    if any_failed then
+        local unwrap_failures = _.compose(_.map(Result.err_or_nil), _.filter(Result.is_failure))
+        local result = Result.failure(unwrap_failures(results))
+        M.channel:send(result)
+        M.channel = nil
+        return result
+    else
+        local result = Result.success(_.map(Result.get_or_nil, results))
         a.scheduler()
         update_registry_state(sources, os.time())
-        return updated_sources
-    end)
-    M.channel:send(result)
-    M.channel = nil
-    return result
+        M.channel:send(result)
+        M.channel = nil
+        return result
+    end
 end
 
 return M
diff --git a/lua/mason-registry/sources/init.lua b/lua/mason-registry/sources/init.lua
index fd9cc573..ccc501fe 100644
--- a/lua/mason-registry/sources/init.lua
+++ b/lua/mason-registry/sources/init.lua
@@ -160,6 +160,16 @@ function LazySourceCollection:iterate(opts)
     end
 end
 
+---@param opts? { include_uninstalled?: boolean }
+function LazySourceCollection:to_list(opts)
+    opts = opts or {}
+    local list = {}
+    for source in self:iterate(opts) do
+        table.insert(list, source)
+    end
+    return list
+end
+
 function LazySourceCollection:__tostring()
     return ("LazySourceCollection(list={%s})"):format(table.concat(vim.tbl_map(tostring, self.list), ", "))
 end
diff --git a/lua/mason/ui/init.lua b/lua/mason/ui/init.lua
index 340434eb..bbe534ad 100644
--- a/lua/mason/ui/init.lua
+++ b/lua/mason/ui/init.lua
@@ -7,7 +7,6 @@ end
 
 function M.open()
     local api = require "mason.ui.instance"
-    require("mason-registry").refresh(function() end)
     api.window.open()
 end
 
diff --git a/lua/mason/ui/instance.lua b/lua/mason/ui/instance.lua
index 70c85858..4302685a 100644
--- a/lua/mason/ui/instance.lua
+++ b/lua/mason/ui/instance.lua
@@ -706,13 +706,17 @@ local function setup_packages(packages)
 end
 
 setup_packages(registry.get_all_packages())
+update_registry_info()
 
-registry:on("update", function()
+registry.refresh(function()
     setup_packages(registry.get_all_packages())
     update_registry_info()
 end)
 
-update_registry_info()
+registry:on("update", function()
+    setup_packages(registry.get_all_packages())
+    update_registry_info()
+end)
 
 window.init {
     effects = effects,