lazy.nvim/lua/lazy/async.lua

223 lines
4.3 KiB
Lua
Raw Permalink Normal View History

local Util = require("lazy.core.util")
2024-06-26 21:11:31 +08:00
local M = {}
---@type Async[]
M._active = {}
---@type Async[]
M._suspended = {}
M._executor = assert(vim.loop.new_check())
M.BUDGET = 10
2024-06-28 22:08:26 +08:00
---@type table<thread, Async>
M._threads = setmetatable({}, { __mode = "k" })
---@alias AsyncEvent "done" | "error" | "yield" | "ok"
2024-06-26 21:11:31 +08:00
---@class Async
2024-06-28 22:08:26 +08:00
---@field _co thread
---@field _fn fun()
---@field _suspended? boolean
---@field _on table<AsyncEvent, fun(res:any, async:Async)[]>
2024-06-26 21:11:31 +08:00
local Async = {}
---@param fn async fun()
---@return Async
2024-06-28 22:08:26 +08:00
function Async.new(fn)
2024-06-26 21:11:31 +08:00
local self = setmetatable({}, { __index = Async })
2024-06-28 22:08:26 +08:00
return self:init(fn)
end
---@param fn async fun()
---@return Async
function Async:init(fn)
self._fn = fn
self._on = {}
self._co = coroutine.create(function()
local ok, err = pcall(self._fn)
if not ok then
self:_emit("error", err)
end
self:_emit("done")
end)
M._threads[self._co] = self
return M.add(self)
end
---@param event AsyncEvent
---@param cb async fun(res:any, async:Async)
function Async:on(event, cb)
self._on[event] = self._on[event] or {}
table.insert(self._on[event], cb)
2024-06-26 21:11:31 +08:00
return self
end
2024-06-28 22:08:26 +08:00
---@private
---@param event AsyncEvent
---@param res any
function Async:_emit(event, res)
for _, cb in ipairs(self._on[event] or {}) do
cb(res, self)
end
end
2024-06-26 21:11:31 +08:00
function Async:running()
2024-06-28 22:08:26 +08:00
return coroutine.status(self._co) ~= "dead"
2024-06-26 21:11:31 +08:00
end
2024-06-28 22:08:26 +08:00
---@async
function Async:sleep(ms)
vim.defer_fn(function()
self:resume()
end, ms)
self:suspend()
end
2024-06-28 22:08:26 +08:00
---@async
---@param yield? boolean
function Async:suspend(yield)
2024-06-28 22:08:26 +08:00
self._suspended = true
if coroutine.running() == self._co and yield ~= false then
2024-06-30 19:35:11 +08:00
M.yield()
2024-06-28 22:08:26 +08:00
end
end
function Async:resume()
2024-06-28 22:08:26 +08:00
self._suspended = false
M._run()
2024-06-28 22:08:26 +08:00
end
---@async
---@param yield? boolean
function Async:wake(yield)
2024-06-28 22:08:26 +08:00
local async = M.running()
assert(async, "Not in an async context")
self:on("done", function()
async:resume()
end)
async:suspend(yield)
end
---@async
function Async:wait()
2024-06-28 22:08:26 +08:00
if coroutine.running() == self._co then
error("Cannot wait on self")
end
local async = M.running()
if async then
self:wake()
else
while self:running() do
2024-06-28 22:08:26 +08:00
vim.wait(10)
end
end
return self
end
2024-06-26 21:11:31 +08:00
function Async:step()
2024-06-28 22:08:26 +08:00
if self._suspended then
return true
end
2024-06-28 22:08:26 +08:00
local status = coroutine.status(self._co)
2024-06-26 21:11:31 +08:00
if status == "suspended" then
2024-06-28 22:08:26 +08:00
local ok, res = coroutine.resume(self._co)
2024-06-26 21:11:31 +08:00
if not ok then
2024-06-28 22:08:26 +08:00
error(res)
2024-06-26 21:11:31 +08:00
elseif res then
2024-06-28 22:08:26 +08:00
self:_emit("yield", res)
2024-06-26 21:11:31 +08:00
end
end
2024-06-28 22:08:26 +08:00
return self:running()
2024-06-26 21:11:31 +08:00
end
2024-06-30 19:35:11 +08:00
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
2024-06-26 21:11:31 +08:00
function M.step()
local start = vim.uv.hrtime()
for _ = 1, #M._active do
2024-06-30 19:35:11 +08:00
if Util.exiting() or vim.uv.hrtime() - start > M.BUDGET * 1e6 then
2024-06-26 21:11:31 +08:00
break
end
2024-06-30 19:35:11 +08:00
local state = table.remove(M._active, 1)
if state:step() then
if state._suspended then
table.insert(M._suspended, state)
else
table.insert(M._active, state)
end
end
2024-06-26 21:11:31 +08:00
end
for _ = 1, #M._suspended do
local state = table.remove(M._suspended, 1)
table.insert(state._suspended and M._suspended or M._active, state)
end
2024-06-30 15:13:04 +08:00
-- M.debug()
2024-06-30 19:35:11 +08:00
if #M._active == 0 or Util.exiting() then
2024-06-26 21:11:31 +08:00
return M._executor:stop()
end
end
2024-06-30 15:13:04 +08:00
function M.debug()
local lines = {
"- active: " .. #M._active,
"- suspended: " .. #M._suspended,
}
for _, async in ipairs(M._active) do
local info = debug.getinfo(async._fn)
local file = vim.fn.fnamemodify(info.short_src:sub(1), ":~:.")
table.insert(lines, ("%s:%d"):format(file, info.linedefined))
if #lines > 10 then
break
end
end
local msg = table.concat(lines, "\n")
M._notif = vim.notify(msg, nil, { replace = M._notif })
end
2024-06-26 21:11:31 +08:00
---@param async Async
function M.add(async)
table.insert(M._active, async)
M._run()
return async
end
function M._run()
2024-06-30 19:35:11 +08:00
if not Util.exiting() and not M._executor:is_active() then
M._executor:start(vim.schedule_wrap(M.step))
2024-06-26 21:11:31 +08:00
end
end
2024-06-28 22:08:26 +08:00
function M.running()
local co = coroutine.running()
if co then
2024-06-29 16:36:35 +08:00
return M._threads[co]
2024-06-26 21:11:31 +08:00
end
end
2024-06-27 20:43:35 +08:00
---@async
---@param ms number
function M.sleep(ms)
2024-06-28 22:08:26 +08:00
local async = M.running()
assert(async, "Not in an async context")
async:sleep(ms)
2024-06-27 20:43:35 +08:00
end
2024-06-28 22:08:26 +08:00
M.Async = Async
M.new = Async.new
2024-06-26 21:11:31 +08:00
return M