refactor(vim.system): factor out on_exit handling

6 files changed, 87 insertions, 73 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index 61aee559ef..71f001f4fc 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -1792,7 +1792,7 @@ vim.system({cmd}, {opts}, {on_exit})                            *vim.system()*
                    object, see return of SystemObj:wait().
 
     Return: ~
-        SystemObj Object with the fields:
+        vim.SystemObj Object with the fields:
         • pid (integer) Process ID
         • wait (fun(timeout: integer|nil): SystemCompleted) Wait for the
           process to complete. Upon timeout the process is sent the KILL
@@ -2545,7 +2545,7 @@ vim.ui.open({path})                                            *vim.ui.open()*
       • {path}  (string) Path or URL to open
 
     Return (multiple): ~
-        SystemCompleted|nil Command result, or nil if not found.
+        vim.SystemCompleted|nil Command result, or nil if not found.
         (string|nil) Error message on failure
 
     See also: ~
diff --git a/runtime/lua/vim/_editor.lua b/runtime/lua/vim/_editor.lua
index ba0345cdd5..68992a16bb 100644
--- a/runtime/lua/vim/_editor.lua
+++ b/runtime/lua/vim/_editor.lua
@@ -117,7 +117,7 @@ vim.log = {
 --- @param on_exit (function|nil) Called when subprocess exits. When provided, the command runs
 ---   asynchronously. Receives SystemCompleted object, see return of SystemObj:wait().
 ---
---- @return SystemObj Object with the fields:
+--- @return vim.SystemObj Object with the fields:
 ---   - pid (integer) Process ID
 ---   - wait (fun(timeout: integer|nil): SystemCompleted) Wait for the process to complete. Upon
 ---     timeout the process is sent the KILL signal (9) and the exit code is set to 124.
diff --git a/runtime/lua/vim/_system.lua b/runtime/lua/vim/_system.lua
index 7bd7ff91f9..5ea0071215 100644
--- a/runtime/lua/vim/_system.lua
+++ b/runtime/lua/vim/_system.lua
@@ -11,13 +11,13 @@ local uv = vim.uv
 --- @field timeout? integer Timeout in ms
 --- @field detach? boolean
 
---- @class SystemCompleted
+--- @class vim.SystemCompleted
 --- @field code integer
 --- @field signal integer
 --- @field stdout? string
 --- @field stderr? string
 
---- @class SystemState
+--- @class vim.SystemState
 --- @field handle? uv.uv_process_t
 --- @field timer?  uv.uv_timer_t
 --- @field pid? integer
@@ -26,7 +26,9 @@ local uv = vim.uv
 --- @field stdin? uv.uv_stream_t
 --- @field stdout? uv.uv_stream_t
 --- @field stderr? uv.uv_stream_t
---- @field result? SystemCompleted
+--- @field stdout_data? string[]
+--- @field stderr_data? string[]
+--- @field result? vim.SystemCompleted
 
 --- @enum vim.SystemSig
 local SIG = {
@@ -44,7 +46,7 @@ local function close_handle(handle)
   end
 end
 
----@param state SystemState
+---@param state vim.SystemState
 local function close_handles(state)
   close_handle(state.handle)
   close_handle(state.stdin)
@@ -53,17 +55,17 @@ local function close_handles(state)
   close_handle(state.timer)
 end
 
---- @class SystemObj
+--- @class vim.SystemObj
 --- @field pid integer
---- @field private _state SystemState
---- @field wait fun(self: SystemObj, timeout?: integer): SystemCompleted
---- @field kill fun(self: SystemObj, signal: integer|string)
---- @field write fun(self: SystemObj, data?: string|string[])
---- @field is_closing fun(self: SystemObj): boolean?
+--- @field private _state vim.SystemState
+--- @field wait fun(self: vim.SystemObj, timeout?: integer): vim.SystemCompleted
+--- @field kill fun(self: vim.SystemObj, signal: integer|string)
+--- @field write fun(self: vim.SystemObj, data?: string|string[])
+--- @field is_closing fun(self: vim.SystemObj): boolean?
 local SystemObj = {}
 
---- @param state SystemState
---- @return SystemObj
+--- @param state vim.SystemState
+--- @return vim.SystemObj
 local function new_systemobj(state)
   return setmetatable({
     pid = state.pid,
@@ -86,7 +88,7 @@ end
 local MAX_TIMEOUT = 2 ^ 31
 
 --- @param timeout? integer
---- @return SystemCompleted
+--- @return vim.SystemCompleted
 function SystemObj:wait(timeout)
   local state = self._state
 
@@ -254,12 +256,53 @@ local function timer_oneshot(timeout, cb)
   return timer
 end
 
+--- @param state vim.SystemState
+--- @param code integer
+--- @param signal integer
+--- @param on_exit fun(result: vim.SystemCompleted)?
+local function _on_exit(state, code, signal, on_exit)
+  close_handles(state)
+
+  local check = assert(uv.new_check())
+  check:start(function()
+    for _, pipe in pairs({ state.stdin, state.stdout, state.stderr }) do
+      if not pipe:is_closing() then
+        return
+      end
+    end
+    check:stop()
+    check:close()
+
+    if state.done == nil then
+      state.done = true
+    end
+
+    if code == 0 and state.done == 'timeout' then
+      code = 124
+    end
+
+    local stdout_data = state.stdout_data
+    local stderr_data = state.stderr_data
+
+    state.result = {
+      code = code,
+      signal = signal,
+      stdout = stdout_data and table.concat(stdout_data) or nil,
+      stderr = stderr_data and table.concat(stderr_data) or nil,
+    }
+
+    if on_exit then
+      on_exit(state.result)
+    end
+  end)
+end
+
 --- Run a system command
 ---
 --- @param cmd string[]
 --- @param opts? SystemOpts
---- @param on_exit? fun(out: SystemCompleted)
---- @return SystemObj
+--- @param on_exit? fun(out: vim.SystemCompleted)
+--- @return vim.SystemObj
 function M.run(cmd, opts, on_exit)
   vim.validate({
     cmd = { cmd, 'table' },
@@ -273,7 +316,7 @@ function M.run(cmd, opts, on_exit)
   local stderr, stderr_handler = setup_output(opts.stderr)
   local stdin, towrite = setup_input(opts.stdin)
 
-  --- @type SystemState
+  --- @type vim.SystemState
   local state = {
     done = false,
     cmd = cmd,
@@ -283,11 +326,6 @@ function M.run(cmd, opts, on_exit)
     stderr = stderr,
   }
 
-  -- Define data buckets as tables and concatenate the elements at the end as
-  -- one operation.
-  --- @type string[], string[]
-  local stdout_data, stderr_data
-
   state.handle, state.pid = spawn(cmd[1], {
     args = vim.list_slice(cmd, 2),
     stdio = { stdin, stdout, stderr },
@@ -296,50 +334,19 @@ function M.run(cmd, opts, on_exit)
     detached = opts.detach,
     hide = true,
   }, function(code, signal)
-    close_handles(state)
-
-    local check = assert(uv.new_check())
-
-    check:start(function()
-      for _, pipe in pairs({ state.stdin, state.stdout, state.stderr }) do
-        if not pipe:is_closing() then
-          return
-        end
-      end
-      check:stop()
-      check:close()
-
-      if state.done == nil then
-        state.done = true
-      end
-
-      if code == 0 and state.done == 'timeout' then
-        code = 124
-      end
-
-      state.result = {
-        code = code,
-        signal = signal,
-        stdout = stdout_data and table.concat(stdout_data) or nil,
-        stderr = stderr_data and table.concat(stderr_data) or nil,
-      }
-
-      if on_exit then
-        on_exit(state.result)
-      end
-    end)
+    _on_exit(state, code, signal, on_exit)
   end, function()
     close_handles(state)
   end)
 
   if stdout then
-    stdout_data = {}
-    stdout:read_start(stdout_handler or default_handler(stdout, opts.text, stdout_data))
+    state.stdout_data = {}
+    stdout:read_start(stdout_handler or default_handler(stdout, opts.text, state.stdout_data))
   end
 
   if stderr then
-    stderr_data = {}
-    stderr:read_start(stderr_handler or default_handler(stderr, opts.text, stderr_data))
+    state.stderr_data = {}
+    stderr:read_start(stderr_handler or default_handler(stderr, opts.text, state.stderr_data))
   end
 
   local obj = new_systemobj(state)
diff --git a/runtime/lua/vim/lsp/rpc.lua b/runtime/lua/vim/lsp/rpc.lua
index a77af937b3..6ab5708721 100644
--- a/runtime/lua/vim/lsp/rpc.lua
+++ b/runtime/lua/vim/lsp/rpc.lua
@@ -648,7 +648,7 @@ function M.start(cmd, cmd_args, dispatchers, extra_spawn_params)
 
   dispatchers = merge_dispatchers(dispatchers)
 
-  local sysobj ---@type SystemObj
+  local sysobj ---@type vim.SystemObj
 
   local client = new_client(dispatchers, {
     write = function(msg)
@@ -708,7 +708,7 @@ function M.start(cmd, cmd_args, dispatchers, extra_spawn_params)
     return
   end
 
-  sysobj = sysobj_or_err --[[@as SystemObj]]
+  sysobj = sysobj_or_err --[[@as vim.SystemObj]]
 
   return public_client(client)
 end
diff --git a/runtime/lua/vim/ui.lua b/runtime/lua/vim/ui.lua
index 87b52787a0..cd90886489 100644
--- a/runtime/lua/vim/ui.lua
+++ b/runtime/lua/vim/ui.lua
@@ -118,7 +118,7 @@ end
 ---
 ---@param path string Path or URL to open
 ---
----@return SystemCompleted|nil # Command result, or nil if not found.
+---@return vim.SystemCompleted|nil # Command result, or nil if not found.
 ---@return string|nil # Error message on failure
 ---
 ---@see |vim.system()|
diff --git a/test/functional/lua/system_spec.lua b/test/functional/lua/system_spec.lua
index 9321468f84..a988d3f0d7 100644
--- a/test/functional/lua/system_spec.lua
+++ b/test/functional/lua/system_spec.lua
@@ -5,13 +5,20 @@ local eq = helpers.eq
 
 local function system_sync(cmd, opts)
   return exec_lua([[
+    local cmd, opts = ...
     local obj = vim.system(...)
-    local pid = obj.pid
+
+    if opts.timeout then
+      -- Minor delay before calling wait() so the timeout uv timer can have a headstart over the
+      -- internal call to vim.wait() in wait().
+      vim.wait(10)
+    end
+
     local res = obj:wait()
 
     -- Check the process is no longer running
-    vim.fn.systemlist({'ps', 'p', tostring(pid)})
-    assert(vim.v.shell_error == 1, 'process still exists')
+    local proc = vim.api.nvim_get_proc(obj.pid)
+    assert(not proc, 'process still exists')
 
     return res
   ]], cmd, opts)
@@ -26,15 +33,15 @@ local function system_async(cmd, opts)
       _G.ret = obj
     end)
 
-    local done = vim.wait(10000, function()
+    local ok = vim.wait(10000, function()
       return _G.done
     end)
 
-    assert(done, 'process did not exit')
+    assert(ok, 'process did not exit')
 
     -- Check the process is no longer running
-    vim.fn.systemlist({'ps', 'p', tostring(obj.pid)})
-    assert(vim.v.shell_error == 1, 'process still exists')
+    local proc = vim.api.nvim_get_proc(obj.pid)
+    assert(not proc, 'process still exists')
 
     return _G.ret
   ]], cmd, opts)
@@ -61,7 +68,7 @@ describe('vim.system', function()
           signal = 15,
           stdout = '',
           stderr = ''
-        }, system({ 'sleep', '10' }, { timeout = 1 }))
+        }, system({ 'sleep', '10' }, { timeout = 1000 }))
       end)
     end)
   end
@@ -83,8 +90,8 @@ describe('vim.system', function()
       assert(done, 'process did not exit')
 
       -- Check the process is no longer running
-      vim.fn.systemlist({'ps', 'p', tostring(cmd.pid)})
-      assert(vim.v.shell_error == 1, 'dwqdqd '..vim.v.shell_error)
+      local proc = vim.api.nvim_get_proc(cmd.pid)
+      assert(not proc, 'process still exists')
 
       assert(signal == 2)
     ]])