From 1fad61712bd3937dda925775a7736b8efbcbf1a7 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Sun, 30 Jun 2024 13:35:11 +0200
Subject: [PATCH] fix(async): make asyncs abortable

---
 lua/lazy/async.lua           | 21 +++++++++++++++++----
 lua/lazy/core/util.lua       |  4 ++++
 lua/lazy/manage/process.lua  |  2 +-
 lua/lazy/view/float.lua      |  3 ++-
 lua/lazy/view/init.lua       |  1 +
 tests/manage/runner_spec.lua |  3 ++-
 tests/manage/task_spec.lua   |  3 ++-
 7 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/lua/lazy/async.lua b/lua/lazy/async.lua
index 5cbdb1bd..7a865b36 100644
--- a/lua/lazy/async.lua
+++ b/lua/lazy/async.lua
@@ -79,7 +79,7 @@ end
 function Async:suspend(yield)
   self._suspended = true
   if coroutine.running() == self._co and yield ~= false then
-    coroutine.yield()
+    M.yield()
   end
 end
 
@@ -132,12 +132,25 @@ function Async:step()
   return self:running()
 end
 
+function M.abort()
+  for _, async in ipairs(M._active) do
+    coroutine.resume(async._co, "abort")
+  end
+end
+
+function M.yield()
+  if coroutine.yield() == "abort" then
+    error("aborted", 2)
+  end
+end
+
 function M.step()
   local start = vim.uv.hrtime()
   for _ = 1, #M._active do
-    if vim.uv.hrtime() - start > M.BUDGET * 1e6 then
+    if Util.exiting() or vim.uv.hrtime() - start > M.BUDGET * 1e6 then
       break
     end
+
     local state = table.remove(M._active, 1)
     if state:step() then
       if state._suspended then
@@ -153,7 +166,7 @@ function M.step()
   end
 
   -- M.debug()
-  if #M._active == 0 then
+  if #M._active == 0 or Util.exiting() then
     return M._executor:stop()
   end
 end
@@ -183,7 +196,7 @@ function M.add(async)
 end
 
 function M._run()
-  if not M._executor:is_active() then
+  if not Util.exiting() and not M._executor:is_active() then
     M._executor:start(vim.schedule_wrap(M.step))
   end
 end
diff --git a/lua/lazy/core/util.lua b/lua/lazy/core/util.lua
index f42cf6a7..34ca1d6e 100644
--- a/lua/lazy/core/util.lua
+++ b/lua/lazy/core/util.lua
@@ -29,6 +29,10 @@ function M.track(data, time)
   end
 end
 
+function M.exiting()
+  return vim.v.exiting ~= vim.NIL
+end
+
 ---@generic T
 ---@param list T[]
 ---@param fn fun(v: T):boolean?
diff --git a/lua/lazy/manage/process.lua b/lua/lazy/manage/process.lua
index c4a4baf5..0a67a14c 100644
--- a/lua/lazy/manage/process.lua
+++ b/lua/lazy/manage/process.lua
@@ -93,7 +93,7 @@ function Process:_run()
     end)
     self:suspend()
     while not (self.handle:is_closing() and stdout:is_closing() and stderr:is_closing()) do
-      coroutine.yield()
+      Async.yield()
     end
   else
     self.data = "Failed to spawn process " .. self.cmd .. " " .. vim.inspect(self.opts)
diff --git a/lua/lazy/view/float.lua b/lua/lazy/view/float.lua
index d57b384f..3a590694 100644
--- a/lua/lazy/view/float.lua
+++ b/lua/lazy/view/float.lua
@@ -252,7 +252,7 @@ end
 ---@param fn fun(self?)
 ---@param desc? string
 ---@param mode? string[]
-function M:on_key(key, fn, desc,mode)
+function M:on_key(key, fn, desc, mode)
   vim.keymap.set(mode or "n", key, function()
     fn(self)
   end, {
@@ -295,6 +295,7 @@ function M:close(opts)
       vim.diagnostic.reset(Config.ns, buf)
       vim.api.nvim_buf_delete(buf, { force = true })
     end
+    vim.cmd.redraw()
   end)
 end
 
diff --git a/lua/lazy/view/init.lua b/lua/lazy/view/init.lua
index 468ac8a8..313d3adc 100644
--- a/lua/lazy/view/init.lua
+++ b/lua/lazy/view/init.lua
@@ -83,6 +83,7 @@ function M.create()
 
   vim.keymap.set("n", ViewConfig.keys.abort, function()
     require("lazy.manage.process").abort()
+    require("lazy.async").abort()
     return ViewConfig.keys.abort
   end, { silent = true, buffer = self.buf, expr = true })
 
diff --git a/tests/manage/runner_spec.lua b/tests/manage/runner_spec.lua
index 925aaf14..a2507c9f 100644
--- a/tests/manage/runner_spec.lua
+++ b/tests/manage/runner_spec.lua
@@ -1,3 +1,4 @@
+local Async = require("lazy.async")
 local Runner = require("lazy.manage.runner")
 
 describe("runner", function()
@@ -33,7 +34,7 @@ describe("runner", function()
       ---@async
       ---@param task LazyTask
       run = function(task)
-        coroutine.yield()
+        Async.yield()
         table.insert(runs, { plugin = task.plugin.name, task = task.name })
       end,
     }
diff --git a/tests/manage/task_spec.lua b/tests/manage/task_spec.lua
index fdd9c35c..f41eb80f 100644
--- a/tests/manage/task_spec.lua
+++ b/tests/manage/task_spec.lua
@@ -1,4 +1,5 @@
 --# selene:allow(incorrect_standard_library_use)
+local Async = require("lazy.async")
 local Task = require("lazy.manage.task")
 
 describe("task", function()
@@ -42,7 +43,7 @@ describe("task", function()
     local running = true
     ---@async
     local task = Task.new(plugin, "test", function()
-      coroutine.yield()
+      Async.yield()
       running = false
     end, opts)
     assert(task:running())