From d2257402e310806617450d57b3d133108ebb20b9 Mon Sep 17 00:00:00 2001
From: zeertzjq <zeertzjq@outlook.com>
Date: Wed, 24 Sep 2025 07:25:29 +0800
Subject: [PATCH] test: don't block to wait for previous session's exit
 (#35885)

This reduces the time taken by clear() by over 20% with ASAN.
---
 test/client/rpc_stream.lua   |  4 +--
 test/client/session.lua      |  6 ++--
 test/client/uv_stream.lua    | 24 ++++++++++----
 test/functional/testnvim.lua | 63 +++++++++++++++++++++++++-----------
 4 files changed, 66 insertions(+), 31 deletions(-)

diff --git a/test/client/rpc_stream.lua b/test/client/rpc_stream.lua
index 9f2672bcf9..282cac2267 100644
--- a/test/client/rpc_stream.lua
+++ b/test/client/rpc_stream.lua
@@ -105,8 +105,8 @@ function RpcStream:read_stop()
   self._stream:read_stop()
 end
 
-function RpcStream:close(signal)
-  self._stream:close(signal)
+function RpcStream:close(signal, noblock)
+  self._stream:close(signal, noblock)
 end
 
 return RpcStream
diff --git a/test/client/session.lua b/test/client/session.lua
index a5839e012a..c4d8ceb197 100644
--- a/test/client/session.lua
+++ b/test/client/session.lua
@@ -13,7 +13,7 @@ local RpcStream = require('test.client.rpc_stream')
 --- @field private _prepare uv.uv_prepare_t
 --- @field private _timer uv.uv_timer_t
 --- @field private _is_running boolean true during `Session:run()` scope.
---- @field exec_lua_setup boolean
+--- @field data table Arbitrary user data.
 local Session = {}
 Session.__index = Session
 if package.loaded['jit'] then
@@ -172,14 +172,14 @@ function Session:stop()
   uv.stop()
 end
 
-function Session:close(signal)
+function Session:close(signal, noblock)
   if not self._timer:is_closing() then
     self._timer:close()
   end
   if not self._prepare:is_closing() then
     self._prepare:close()
   end
-  self._rpc_stream:close(signal)
+  self._rpc_stream:close(signal, noblock)
   self.closed = true
 end
 
diff --git a/test/client/uv_stream.lua b/test/client/uv_stream.lua
index 6e1a6995be..73e3255f6f 100644
--- a/test/client/uv_stream.lua
+++ b/test/client/uv_stream.lua
@@ -9,7 +9,7 @@ local uv = vim.uv
 --- @field write fun(self, data: string|string[])
 --- @field read_start fun(self, cb: fun(chunk: string))
 --- @field read_stop fun(self)
---- @field close fun(self, signal?: string)
+--- @field close fun(self, signal?: string, noblock?: boolean)
 
 --- Stream over given pipes.
 ---
@@ -126,6 +126,8 @@ end
 --- @field private _child_stdin uv.uv_pipe_t
 --- @field private _child_stdout uv.uv_pipe_t
 --- @field private _child_stderr uv.uv_pipe_t
+--- @field package _closed integer
+--- @field package _on_exit fun(closed: integer?)
 --- Collects stdout (if `collect_text=true`). Treats data as text (CRLF converted to LF).
 --- @field stdout string
 --- Collects stderr as raw data.
@@ -147,8 +149,10 @@ ProcStream.__index = ProcStream
 --- @param argv string[]
 --- @param env string[]?
 --- @param io_extra uv.uv_pipe_t?
+--- @param on_exit fun(closed: integer?)?  Called after the child process exits.
+--- `closed` is the timestamp (uv.now()) when close() was called, or nil if it wasn't.
 --- @return test.ProcStream
-function ProcStream.spawn(argv, env, io_extra)
+function ProcStream.spawn(argv, env, io_extra, on_exit)
   local self = setmetatable({
     collect_text = false,
     output = function(self)
@@ -165,6 +169,7 @@ function ProcStream.spawn(argv, env, io_extra)
     _child_stdout = assert(uv.new_pipe(false)),
     _child_stderr = assert(uv.new_pipe(false)),
     _exiting = false,
+    _on_exit = on_exit,
   }, ProcStream)
   local prog = argv[1]
   local args = {} --- @type string[]
@@ -181,6 +186,9 @@ function ProcStream.spawn(argv, env, io_extra)
     self.signal = signal
     -- "Abort" exit may not set status; force to nonzero in that case.
     self.status = (0 ~= (status or 0) or 0 == (signal or 0)) and status or (128 + (signal or 0))
+    if self._on_exit then
+      self._on_exit(self._closed)
+    end
   end)
 
   if not self._proc then
@@ -238,11 +246,11 @@ function ProcStream:read_stop()
   self._child_stderr:read_stop()
 end
 
-function ProcStream:close(signal)
+function ProcStream:close(signal, noblock)
   if self._closed then
     return
   end
-  self._closed = true
+  self._closed = uv.now()
   self:read_stop()
   self._child_stdin:close()
   self._child_stdout:close()
@@ -250,10 +258,12 @@ function ProcStream:close(signal)
   if type(signal) == 'string' then
     self._proc:kill('sig' .. signal)
   end
-  while self.status == nil do
-    uv.run 'once'
+  if not noblock then
+    while self.status == nil do
+      uv.run 'once'
+    end
+    return self.status, self.signal
   end
-  return self.status, self.signal
 end
 
 return {
diff --git a/test/functional/testnvim.lua b/test/functional/testnvim.lua
index 9a41395462..51ba6f1354 100644
--- a/test/functional/testnvim.lua
+++ b/test/functional/testnvim.lua
@@ -1,5 +1,6 @@
 local uv = vim.uv
 local t = require('test.testutil')
+local busted = require('busted')
 
 local Session = require('test.client.session')
 local uv_stream = require('test.client.uv_stream')
@@ -438,24 +439,12 @@ local function remove_args(args, args_rm)
   return new_args
 end
 
-function M.check_close()
+function M.check_close(noblock)
   if not session then
     return
   end
-  local start_time = uv.now()
-  session:close()
-  uv.update_time() -- Update cached value of luv.now() (libuv: uv_now()).
-  local end_time = uv.now()
-  local delta = end_time - start_time
-  if delta > 500 then
-    print(
-      'nvim took '
-        .. delta
-        .. ' milliseconds to exit after last test\n'
-        .. 'This indicates a likely problem with the test even if it passed!\n'
-    )
-    io.stdout:flush()
-  end
+
+  session:close(nil, noblock)
   session = nil
 end
 
@@ -492,6 +481,8 @@ function M.clear(...)
   return M.get_session()
 end
 
+local n_processes = 0
+
 --- Starts a new Nvim process with the given args and returns a msgpack-RPC session.
 ---
 --- Does not replace the current global session, unlike `clear()`.
@@ -501,16 +492,44 @@ end
 --- @return test.Session
 --- @overload fun(keep: boolean, opts: test.session.Opts): test.Session
 function M.new_session(keep, ...)
-  if not keep then
-    M.check_close()
+  local test_id = _G._nvim_test_id
+  if not keep and session ~= nil then
+    -- Don't block for the previous session's exit if it's from a different test.
+    session:close(nil, session.data and session.data.test_id ~= test_id)
+    session = nil
   end
 
   local argv, env, io_extra = M._new_argv(...)
 
-  local proc = ProcStream.spawn(argv, env, io_extra)
-  return Session.new(proc)
+  local proc = ProcStream.spawn(argv, env, io_extra, function(closed)
+    n_processes = n_processes - 1
+    local delta = 0
+    if closed then
+      uv.update_time() -- Update cached value of uv.now() (libuv: uv_now()).
+      delta = uv.now() - closed
+    end
+    if delta > 500 then
+      print(
+        ('Nvim session %s took %d milliseconds to exit\n'):format(test_id, delta)
+          .. 'This indicates a likely problem with the test even if it passed!\n'
+      )
+      io.stdout:flush()
+    end
+  end)
+  n_processes = n_processes + 1
+
+  local new_session = Session.new(proc)
+  new_session.data = { test_id = test_id }
+  return new_session
 end
 
+busted.subscribe({ 'suite', 'end' }, function()
+  M.check_close(true)
+  while n_processes > 0 do
+    uv.run('once')
+  end
+end)
+
 --- Starts a (non-RPC, `--headless --listen "Tx"`) Nvim process, waits for exit, and returns result.
 ---
 --- @param ... string Nvim CLI args, or `test.session.Opts` table.
@@ -1015,6 +1034,11 @@ return function()
 
   if after_each then
     after_each(function()
+      if not vim.endswith(_G._nvim_test_id, 'x') then
+        -- Use a different test ID for skipped tests as well as Nvim instances spawned
+        -- between this after_each() and the next before_each() (e.g. in setup()).
+        _G._nvim_test_id = _G._nvim_test_id .. 'x'
+      end
       check_logs()
       check_cores('build/bin/nvim')
       if session then
@@ -1027,5 +1051,6 @@ return function()
       end
     end)
   end
+
   return M
 end