Feature: Kill process on task cancel

lua/nio/process.lua

@@ -91,18 +91,30 @@ function nio.process.run(opts)
 
   local stdio = { stdin.pipe, stdout.pipe, stderr.pipe }
 
-  local handle, pid_or_error = vim.loop.spawn(cmd, {
-    args = args,
-    stdio = stdio,
-    env = opts.env,
-    cwd = opts.cwd,
-    uid = opts.uid,
-    gid = opts.gid,
-    verbatim = opts.verbatim,
-    detached = opts.detached,
-    hide = opts.hide,
-  }, function(code, _)
-    exit_code_future.set(code)
+  local handle, pid_or_error
+  task = require("nio").run(function()
+    handle, pid_or_error = vim.loop.spawn(cmd, {
+      args = args,
+      stdio = stdio,
+      env = opts.env,
+      cwd = opts.cwd,
+      uid = opts.uid,
+      gid = opts.gid,
+      verbatim = opts.verbatim,
+      detached = opts.detached,
+      hide = opts.hide,
+    }, function(code, _)
+      exit_code_future.set(code)
+    end)
+    exit_code_future.wait()
+  end)
+
+  local signal = function(signal)
+    vim.loop.process_kill(handle, signal)
+  end
+
+  task.on_cancel(function()
+    signal(15)
   end)
 
   if not handle then
@@ -125,9 +137,7 @@ function nio.process.run(opts)
   local process
   process = {
     pid = pid_or_error,
-    signal = function(signal)
-      vim.loop.process_kill(handle, signal)
-    end,
+    signal = signal,
     stdin = {
       write = stdin.write,
       fd = stdin_fd,
@@ -155,6 +165,7 @@ function nio.process.run(opts)
       return { stdin.close(), stdout.close(), stderr.close() }
     end,
   }
+
   return process
 end
 

lua/nio/tasks.lua

@@ -38,6 +38,7 @@ end
 ---@field cancel fun(): nil Cancels the task
 ---@field trace fun(): string Get the stack trace of the task
 ---@field wait async function Wait for the task to finish, returning any result
+---@field on_cancel fun(callback: fun()) Register a callback for when the task is cancelled
 
 ---@class nio.tasks.TaskError
 ---@field message string
@@ -60,6 +61,7 @@ end
 function nio.tasks.run(func, cb)
   local co = coroutine.create(func)
   local cancelled = false
+  local on_cancel
   local step
   local task = { parent = nio.tasks.current_task() }
   if task.parent then
@@ -68,15 +70,31 @@ function nio.tasks.run(func, cb)
   end
   local future = require("nio").control.future()
 
+  function task.on_cancel(callback)
+    on_cancel = callback
+  end
+
   function task.cancel()
-    if cancelled or coroutine.status(co) == "dead" then
+    if cancelled then
       return
     end
+
+    local coroutine_dead = coroutine.status(co) == "dead"
+    if on_cancel then
+      if not coroutine_dead then
+        on_cancel()
+      end
+      on_cancel = nil
+    end
     cancelled = true
+
     for _, child in pairs(child_tasks[task] or {}) do
       child.cancel()
     end
-    step()
+
+    if not coroutine_dead then
+      step()
+    end
   end
 
   function task.trace()

tests/process_spec.lua

@@ -204,4 +204,21 @@ describe("process", function()
     local exit_code = process.result()
     assert.equal(1, exit_code)
   end)
+
+  a.it("cancelling task, cancels loose process", function()
+    local process
+    task = nio.run(function()
+      local pipe = assert(vim.loop.new_pipe())
+      process = assert(nio.process.run({
+        cmd = "cat",
+        stdin = pipe,
+      }))
+      return true
+    end)
+    task.cancel()
+    local result = task.wait()
+    assert.equal(true, result)
+    local exit_code = process.result()
+    assert.equal(0, exit_code)
+  end)
 end)