[WIP] cancellation support for async functions

lua/plenary/async/async.lua

@@ -3,11 +3,18 @@ local vararg = require('plenary.vararg')
 local errors = require('plenary.errors')
 local traceback_error = errors.traceback_error
 local f = require('plenary.functional')
+local Context = require('plenary.async.context')
 
 local M = {}
 
 ---because we can't store varargs
-local function callback_or_next(step, thread, callback, ...)
+local function callback_or_next(step, thread, callback, ctx, ...)
+  -- this means that we have canceled,
+  -- returning means that we will yield forever and never resume this coroutine
+  if ctx.canceled then return end
+
+  ctx:reset_cancel()
+
   local stat = f.first(...)
 
   if not stat then
@@ -20,10 +27,14 @@ local function callback_or_next(step, thread, callback, ...)
   else
     local returned_function = f.second(...)
     local nargs = f.third(...)
+    local canceller = select(4, ...)
+
     assert(type(returned_function) == "function", "type error :: expected func")
-    local rstat, msg = pcall(returned_function, vararg.rotate(nargs, step, select(4, ...)))
+    local rstat, res = pcall(returned_function, vararg.rotate(nargs, step, select(5, ...)))
     if not rstat then
-      error(('Failed to call leaf async function: %s'):format(msg))
+      error(('Failed to call leaf async function: %s'):format(res))
+    elseif canceller ~= nil then
+      ctx:set_cancel(canceller, res)
     end
   end
 end
@@ -34,14 +45,18 @@ end
 local execute = function(async_function, callback, ...)
   assert(type(async_function) == "function", "type error :: expected func")
 
+  local ctx = Context.new()
+
   local thread = co.create(async_function)
 
   local step
   step = function(...)
-    callback_or_next(step, thread, callback, co.resume(thread, ...))
+    callback_or_next(step, thread, callback, ctx, co.resume(thread, ...))
   end
 
   step(...)
+
+  return ctx
 end
 
 local add_leaf_function
@@ -69,7 +84,7 @@ end
 ---@param func function: A callback style function to be converted. The last argument must be the callback.
 ---@param argc number: The number of arguments of func. Must be included.
 ---@return function: Returns an async function
-M.wrap = function(func, argc)
+M.wrap = function(func, argc, canceller)
   if type(func) ~= "function" then
     traceback_error("type error :: expected func, got " .. type(func))
   end
@@ -84,7 +99,7 @@ M.wrap = function(func, argc)
     if nargs == argc then
       return func(...)
     else
-      return co.yield(func, argc, ...)
+      return co.yield(func, argc, canceller, ...)
     end
   end
 

lua/plenary/async/context.lua

@@ -0,0 +1,33 @@
+local Context = {}
+Context.__index = Context
+
+function Context.new()
+  return setmetatable({
+    canceled = false,
+  }, Context)
+end
+
+function Context:set_cancel(canceller, cancel_handle)
+  assert(type(canceller) == "function")
+  assert(type(cancel_handle) == "userdata")
+
+  self.canceller = canceller
+  self.cancel_handle = cancel_handle
+end
+
+function Context:reset_cancel()
+  self.canceller = nil
+  self.cancel_handle = nil
+end
+
+function Context:cancel()
+  if self.canceled == true then return end
+
+  self.canceled = true
+
+  if self.canceller ~= nil and self.cancel_handle ~= nil then
+    self.canceller(self.cancel_handle)
+  end
+end
+
+return Context

lua/plenary/async/uv_async.lua

@@ -3,80 +3,76 @@ local uv = vim.loop
 
 local M = {}
 
-local function add(name, argc)
-  local success, ret = pcall(a.wrap, uv[name], argc)
+M.close = a.wrap(uv.close, 4) -- close a handle
 
-  if not success then
-    error("Failed to add function with name " .. name)
-  end
-
-  M[name] = ret
+-- filesystem operations
+local fs_canceller = function(handle)
+  assert(uv.cancel(handle) == 0)
+  uv.close(handle)
+  -- handle
 end
 
-add('close', 4) -- close a handle
-
--- filesystem operations
-add('fs_open', 4)
-add('fs_read', 4)
-add('fs_close', 2)
-add('fs_unlink', 2)
-add('fs_write', 4)
-add('fs_mkdir', 3)
-add('fs_mkdtemp', 2)
+M.fs_open = a.wrap(uv.fs_open, 4, fs_canceller)
+M.fs_read = a.wrap(uv.fs_read, 4, fs_canceller)
+M.fs_close = a.wrap(uv.fs_close, 2, fs_canceller)
+M.fs_unlink = a.wrap(uv.fs_unlink, 2, fs_canceller)
+M.fs_write = a.wrap(uv.fs_write, 4, fs_canceller)
+M.fs_mkdir = a.wrap(uv.fs_mkdir, 3, fs_canceller)
+M.fs_mkdtemp = a.wrap(uv.fs_mkdtemp, 2, fs_canceller)
 -- 'fs_mkstemp',
-add('fs_rmdir', 2)
-add('fs_scandir', 2)
-add('fs_stat', 2)
-add('fs_fstat', 2)
-add('fs_lstat', 2)
-add('fs_rename', 3)
-add('fs_fsync', 2)
-add('fs_fdatasync', 2)
-add('fs_ftruncate', 3)
-add('fs_sendfile', 5)
-add('fs_access', 3)
-add('fs_chmod', 3)
-add('fs_fchmod', 3)
-add('fs_utime', 4)
-add('fs_futime', 4)
+M.fs_rmdir = a.wrap(uv.fs_rmdir, 2, fs_canceller)
+M.fs_scandir = a.wrap(uv.fs_scandir, 2, fs_canceller)
+M.fs_stat = a.wrap(uv.fs_stat, 2, fs_canceller)
+M.fs_fstat = a.wrap(uv.fs_fstat, 2, fs_canceller)
+M.fs_lstat = a.wrap(uv.fs_lstat, 2, fs_canceller)
+M.fs_rename = a.wrap(uv.fs_rename, 3, fs_canceller)
+M.fs_fsync = a.wrap(uv.fs_fsync, 2, fs_canceller)
+M.fs_fdatasync = a.wrap(uv.fs_fdatasync, 2, fs_canceller)
+M.fs_ftruncate = a.wrap(uv.fs_ftruncate, 3, fs_canceller)
+M.fs_sendfile = a.wrap(uv.fs_sendfile, 5, fs_canceller)
+M.fs_access = a.wrap(uv.fs_access, 3, fs_canceller)
+M.fs_chmod = a.wrap(uv.fs_chmod, 3, fs_canceller)
+M.fs_fchmod = a.wrap(uv.fs_fchmod, 3, fs_canceller)
+M.fs_utime = a.wrap(uv.fs_utime, 4, fs_canceller)
+M.fs_futime = a.wrap(uv.fs_futime, 4, fs_canceller)
 -- 'fs_lutime',
-add('fs_link', 3)
-add('fs_symlink', 4)
-add('fs_readlink', 2)
-add('fs_realpath', 2)
-add('fs_chown', 4)
-add('fs_fchown', 4)
+M.fs_link = a.wrap(uv.fs_link, 3, fs_canceller)
+M.fs_symlink = a.wrap(uv.fs_symlink, 4, fs_canceller)
+M.fs_readlink = a.wrap(uv.fs_readlink, 2, fs_canceller)
+M.fs_realpath = a.wrap(uv.fs_realpath, 2, fs_canceller)
+M.fs_chown = a.wrap(uv.fs_chown, 4, fs_canceller)
+M.fs_fchown = a.wrap(uv.fs_fchown, 4, fs_canceller)
 -- 'fs_lchown',
-add('fs_copyfile', 4)
+M.fs_copyfile = a.wrap(uv.fs_copyfile, 4, fs_canceller)
 -- add('fs_opendir', 3) -- TODO: fix this one
-add('fs_readdir', 2)
-add('fs_closedir', 2)
+M.fs_readdir = a.wrap(uv.fs_readdir, 2, fs_canceller)
+M.fs_closedir = a.wrap(uv.fs_closedir, 2, fs_canceller)
 -- 'fs_statfs',
 
 -- stream
-add('shutdown', 2)
-add('listen', 3)
+M.shutdown = a.wrap(uv.shutdown, 2)
+M.listen = a.wrap(uv.listen, 3)
 -- add('read_start', 2) -- do not do this one, the callback is made multiple times
-add('write', 3)
-add('write2', 4)
-add('shutdown', 2)
+M.write = a.wrap(uv.write, 3)
+M.write2 = a.wrap(uv.write2, 4)
+M.shutdown = a.wrap(uv.shutdown, 2)
 
 -- tcp
-add('tcp_connect', 4)
+M.tcp_connect = a.wrap(uv.tcp_connect, 4)
 -- 'tcp_close_reset',
 
 -- pipe
-add('pipe_connect', 3)
+M.pipe_connect = a.wrap(uv.pipe_connect, 3)
 
 -- udp
-add('udp_send', 5)
-add('udp_recv_start', 2)
+M.udp_send = a.wrap(uv.udp_send, 5)
+M.udp_recv_start = a.wrap(uv.udp_recv_start, 2)
 
 -- fs event (wip make into async await event)
 -- fs poll event (wip make into async await event)
 
 -- dns
-add('getaddrinfo', 4)
-add('getnameinfo', 2)
+M.getaddrinfo = a.wrap(uv.getaddrinfo, 4)
+M.getnameinfo = a.wrap(uv.getnameinfo, 2)
 
 return M