feat: show warning message when exiting neovim with active installations (#725)

5 files changed, 206 insertions, 1 deletions

diff --git a/lua/mason-core/package/init.lua b/lua/mason-core/package/init.lua
index 81d8ca31..e46219b7 100644
--- a/lua/mason-core/package/init.lua
+++ b/lua/mason-core/package/init.lua
@@ -80,7 +80,17 @@ function Package:new_handle()
     log.fmt_trace("Creating new handle for %s", self)
     local handle = InstallationHandle.new(self)
     self.handle = handle
+
+    -- First emit a private autocmd via the native event bus. This is to enable some internal perf improvements (helps avoid loading some Lua modules).
+    if vim.fn.has "nvim-0.8.0" == 1 then
+        vim.api.nvim_exec_autocmds("User", { pattern = "__MasonPackageHandle", data = self.name })
+    else
+        vim.api.nvim_exec_autocmds("User", { pattern = "__MasonPackageHandle" })
+    end
+
     self:emit("handle", handle)
+    registry:emit("package:handle", self, handle)
+
     return handle
 end
 
diff --git a/lua/mason/init.lua b/lua/mason/init.lua
index 37926acc..f6c46ef4 100644
--- a/lua/mason/init.lua
+++ b/lua/mason/init.lua
@@ -4,6 +4,24 @@ local platform = require "mason-core.platform"
 
 local M = {}
 
+local function setup_autocmds()
+    -- lazily set up terminator
+    vim.api.nvim_create_autocmd("User", {
+        pattern = "__MasonPackageHandle", -- private autocmd specific for this very use case
+        callback = function()
+            require("mason.terminator").setup()
+        end,
+        once = true,
+    })
+
+    vim.api.nvim_create_autocmd("VimLeavePre", {
+        callback = function()
+            require("mason.terminator").terminate()
+        end,
+        once = true,
+    })
+end
+
 ---@param config MasonSettings?
 function M.setup(config)
     if config then
@@ -17,6 +35,7 @@ function M.setup(config)
     end
 
     require "mason.api.command"
+    setup_autocmds()
 end
 
 return M
diff --git a/lua/mason/terminator.lua b/lua/mason/terminator.lua
new file mode 100644
index 00000000..7a9290a0
--- /dev/null
+++ b/lua/mason/terminator.lua
@@ -0,0 +1,101 @@
+local registry = require "mason-registry"
+local a = require "mason-core.async"
+
+-- Hasta la vista, baby.
+--                      ______
+--                    <((((((\\\
+--                    /      . }\
+--                    ;--..--._|}
+-- (\                 '--/\--'  )
+--  \\                | '-'  :'|
+--   \\               . -==- .-|
+--    \\               \.__.'   \--._
+--    [\\          __.--|       //  _/'--.
+--    \ \\       .'-._ ('-----'/ __/      \
+--     \ \\     /   __>|      | '--.       |
+--      \ \\   |   \   |     /    /       /
+--       \ '\ /     \  |     |  _/       /
+--        \  \       \ |     | /        /
+--  snd    \  \      \        /
+
+local M = {}
+
+---@async
+---@param handles InstallHandle[]
+local function terminate_handles(handles)
+    a.wait_all(vim.tbl_map(
+        ---@param handle InstallHandle
+        function(handle)
+            if not handle:is_closed() then
+                handle:terminate()
+            end
+            return function()
+                a.wait(function(resolve)
+                    if handle:is_closed() then
+                        resolve()
+                    else
+                        handle:once("closed", resolve)
+                    end
+                end)
+            end
+        end,
+        handles
+    ))
+end
+
+local active_handles = {}
+
+function M.setup()
+    registry:on("package:handle", function(_, handle)
+        if handle:is_closed() then
+            return
+        end
+        active_handles[handle] = true
+        handle:once("closed", function()
+            active_handles[handle] = nil
+        end)
+    end)
+end
+
+function M.terminate()
+    local handles = vim.tbl_keys(active_handles)
+    if #handles > 0 then
+        local package_names = vim.tbl_map(function(h)
+            return h.package.name
+        end, handles)
+        table.sort(package_names)
+
+        -- 1. Print warning message.
+        vim.api.nvim_echo({
+            {
+                "[mason.nvim] Neovim is exiting while packages are still installing. Terminating all installations…",
+                "WarningMsg",
+            },
+        }, true, {})
+        vim.cmd "redraw"
+
+        -- 2. Synchronously terminate all installation handles.
+        a.run_blocking(function()
+            a.wait_first {
+                function()
+                    a.sleep(5000)
+                end,
+                function()
+                    terminate_handles(handles)
+                end,
+            }
+        end)
+
+        -- 3. Schedule error message to be displayed so that Neovim prints it to the tty.
+        --    XXX: does this need to be conditional on which UIs are attached?
+        vim.schedule(function()
+            vim.api.nvim_err_writeln(
+                ("[mason.nvim] Neovim exited while the following packages were installing. Installation was aborted.\n- %s"):format(
+                    table.concat(package_names, #package_names > 5 and ", " or "\n- ")
+                )
+            )
+        end)
+    end
+end
+
+return M
diff --git a/tests/mason/terminator_spec.lua b/tests/mason/terminator_spec.lua
new file mode 100644
index 00000000..9f243086
--- /dev/null
+++ b/tests/mason/terminator_spec.lua
@@ -0,0 +1,75 @@
+local stub = require "luassert.stub"
+local spy = require "luassert.spy"
+local a = require "mason-core.async"
+local registry = require "mason-registry"
+local terminator = require "mason.terminator"
+local _ = require "mason-core.functional"
+local InstallHandle = require "mason-core.installer.handle"
+
+describe("terminator", function()
+    before_each(function()
+        terminator.setup()
+    end)
+
+    it(
+        "should terminate all active handles on nvim exit",
+        async_test(function()
+            local dummy = registry.get_package "dummy"
+            local dummy2 = registry.get_package "dummy2"
+            for _, pkg in ipairs { dummy, dummy2 } do
+                stub(pkg.spec, "install")
+                pkg.spec.install.invokes(function()
+                    a.sleep(10000)
+                end)
+            end
+
+            dummy:install()
+            dummy2:install()
+            spy.on(InstallHandle, "terminate")
+
+            terminator.terminate()
+            a.scheduler()
+
+            assert.spy(InstallHandle.terminate).was_called(2)
+        end)
+    )
+
+    it(
+        "should print warning messages",
+        async_test(function()
+            spy.on(vim.api, "nvim_echo")
+            spy.on(vim.api, "nvim_err_writeln")
+            local dummy = registry.get_package "dummy"
+            local dummy2 = registry.get_package "dummy2"
+            for _, pkg in ipairs { dummy, dummy2 } do
+                stub(pkg.spec, "install")
+                pkg.spec.install.invokes(function()
+                    a.sleep(10000)
+                end)
+            end
+
+            dummy:install()
+            dummy2:install()
+            spy.on(InstallHandle, "terminate")
+
+            terminator.terminate()
+
+            assert.spy(vim.api.nvim_echo).was_called(1)
+            assert.spy(vim.api.nvim_echo).was_called_with({
+                {
+                    "[mason.nvim] Neovim is exiting while packages are still installing. Terminating all installations…",
+                    "WarningMsg",
+                },
+            }, true, {})
+
+            a.scheduler()
+
+            assert.spy(vim.api.nvim_err_writeln).was_called(1)
+            assert.spy(vim.api.nvim_err_writeln).was_called_with(_.dedent [[
+                [mason.nvim] Neovim exited while the following packages were installing. Installation was aborted.
+                - dummy
+                - dummy2
+            ]])
+        end)
+    )
+end)
diff --git a/tests/minimal_init.vim b/tests/minimal_init.vim
index ab61d7ed..f099416c 100644
--- a/tests/minimal_init.vim
+++ b/tests/minimal_init.vim
@@ -18,7 +18,7 @@ lua require("test_helpers")
 lua <<EOF
 local index = require "mason-registry.index"
 index["dummy"] = "dummy_package"
-index["dummy2"] = "dummy_package"
+index["dummy2"] = "dummy2_package"
 
 require("mason").setup {
     install_root_dir = vim.env.INSTALL_ROOT_DIR,