fix(rpc)!: preseve files when stdio channel is closed (#22137)

BREAKING CHANGE: Unsaved changes are now preserved rather than discarded
when stdio channel is closed.

10 files changed, 99 insertions, 42 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index c8d40094f3..2afb22bb43 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -47,6 +47,9 @@ The following changes may require adaptations in user config or plugins.
 
 • libiconv is now a required build dependency.
 
+• Unsaved changes are now preserved rather than discarded when |channel-stdio|
+  is closed.
+
 ==============================================================================
 NEW FEATURES                                                    *news-features*
 
diff --git a/src/nvim/event/process.c b/src/nvim/event/process.c
index 1a524a56ca..1219566e9b 100644
--- a/src/nvim/event/process.c
+++ b/src/nvim/event/process.c
@@ -422,7 +422,12 @@ static void exit_event(void **argv)
   }
 
   if (!exiting) {
-    os_exit(status);
+    if (ui_client_channel_id) {
+      os_exit(status);
+    } else {
+      assert(status == 0);  // Called from rpc_close(), which passes 0 as status.
+      preserve_exit(NULL);
+    }
   }
 }
 
diff --git a/src/nvim/main.c b/src/nvim/main.c
index 4ee02b4e1b..f36730a8b3 100644
--- a/src/nvim/main.c
+++ b/src/nvim/main.c
@@ -674,6 +674,7 @@ void os_exit(int r)
 void getout(int exitval)
   FUNC_ATTR_NORETURN
 {
+  assert(!ui_client_channel_id);
   exiting = true;
 
   // On error during Ex mode, exit with a non-zero code.
@@ -794,6 +795,7 @@ void getout(int exitval)
 }
 
 /// Preserve files, print contents of `errmsg`, and exit 1.
+/// @param errmsg  If NULL, this function will not print anything.
 ///
 /// May be called from deadly_signal().
 void preserve_exit(const char *errmsg)
@@ -819,19 +821,21 @@ void preserve_exit(const char *errmsg)
     // For TUI: exit alternate screen so that the error messages can be seen.
     ui_client_stop();
   }
-  os_errmsg(errmsg);
-  os_errmsg("\n");
+  if (errmsg != NULL) {
+    os_errmsg(errmsg);
+    os_errmsg("\n");
+  }
   if (ui_client_channel_id) {
     os_exit(1);
   }
-  ui_flush();
 
   ml_close_notmod();                // close all not-modified buffers
 
   FOR_ALL_BUFFERS(buf) {
     if (buf->b_ml.ml_mfp != NULL && buf->b_ml.ml_mfp->mf_fname != NULL) {
-      os_errmsg("Vim: preserving files...\r\n");
-      ui_flush();
+      if (errmsg != NULL) {
+        os_errmsg("Vim: preserving files...\r\n");
+      }
       ml_sync_all(false, false, true);  // preserve all swap files
       break;
     }
@@ -839,7 +843,9 @@ void preserve_exit(const char *errmsg)
 
   ml_close_all(false);              // close all memfiles, without deleting
 
-  os_errmsg("Vim: Finished.\r\n");
+  if (errmsg != NULL) {
+    os_errmsg("Vim: Finished.\r\n");
+  }
 
   getout(1);
 }
diff --git a/src/nvim/msgpack_rpc/channel.c b/src/nvim/msgpack_rpc/channel.c
index 34c8f89e3f..f8be0d4c8d 100644
--- a/src/nvim/msgpack_rpc/channel.c
+++ b/src/nvim/msgpack_rpc/channel.c
@@ -550,6 +550,10 @@ void rpc_close(Channel *channel)
 
   if (channel->streamtype == kChannelStreamStdio
       || (channel->id == ui_client_channel_id && channel->streamtype != kChannelStreamProc)) {
+    if (channel->streamtype == kChannelStreamStdio) {
+      // Avoid hanging when there are no other UIs and a prompt is triggered on exit.
+      remote_ui_disconnect(channel->id);
+    }
     exit_from_channel(0);
   }
 }
diff --git a/test/functional/ex_cmds/oldfiles_spec.lua b/test/functional/ex_cmds/oldfiles_spec.lua
index 5f87c3cdd9..19611429e0 100644
--- a/test/functional/ex_cmds/oldfiles_spec.lua
+++ b/test/functional/ex_cmds/oldfiles_spec.lua
@@ -2,6 +2,8 @@ local Screen = require('test.functional.ui.screen')
 local helpers = require('test.functional.helpers')(after_each)
 
 local clear = helpers.clear
+local command = helpers.command
+local expect_exit = helpers.expect_exit
 local buf, eq, feed_command = helpers.curbufmeths, helpers.eq, helpers.feed_command
 local feed, poke_eventloop = helpers.feed, helpers.poke_eventloop
 local ok = helpers.ok
@@ -19,6 +21,7 @@ describe(':oldfiles', function()
   before_each(_clear)
 
   after_each(function()
+    expect_exit(command, 'qall!')
     os.remove(shada_file)
   end)
 
@@ -42,6 +45,7 @@ describe(':oldfiles', function()
                                                                                                           |
       Press ENTER or type command to continue^                                                             |
     ]])
+    feed('<CR>')
   end)
 
   it('can be filtered with :filter', function()
@@ -107,6 +111,7 @@ describe(':browse oldfiles', function()
   end)
 
   after_each(function()
+    expect_exit(command, 'qall!')
     os.remove(shada_file)
   end)
 
diff --git a/test/functional/ex_cmds/profile_spec.lua b/test/functional/ex_cmds/profile_spec.lua
index 2b92f8d0de..bf045a4d1d 100644
--- a/test/functional/ex_cmds/profile_spec.lua
+++ b/test/functional/ex_cmds/profile_spec.lua
@@ -30,6 +30,7 @@ describe(':profile', function()
   before_each(helpers.clear)
 
   after_each(function()
+    helpers.expect_exit(command, 'qall!')
     if lfs.attributes(tempfile, 'uid') ~= nil then
       os.remove(tempfile)
     end
diff --git a/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua b/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
index 8eed00c973..ad59025d47 100644
--- a/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
+++ b/test/functional/ex_cmds/swapfile_preserve_recover_spec.lua
@@ -13,6 +13,7 @@ local nvim_prog = helpers.nvim_prog
 local ok = helpers.ok
 local rmdir = helpers.rmdir
 local new_argv = helpers.new_argv
+local new_pipename = helpers.new_pipename
 local pesc = helpers.pesc
 local os_kill = helpers.os_kill
 local set_session = helpers.set_session
@@ -37,10 +38,21 @@ describe(':recover', function()
 
 end)
 
-describe(':preserve', function()
+describe("preserve and (R)ecover with custom 'directory'", function()
   local swapdir = lfs.currentdir()..'/Xtest_recover_dir'
+  local testfile = 'Xtest_recover_file1'
+  -- Put swapdir at the start of the 'directory' list. #1836
+  -- Note: `set swapfile` *must* go after `set directory`: otherwise it may
+  -- attempt to create a swapfile in different directory.
+  local init = [[
+    set directory^=]]..swapdir:gsub([[\]], [[\\]])..[[//
+    set swapfile fileformat=unix undolevels=-1
+  ]]
+
+  local nvim1
   before_each(function()
-    clear()
+    nvim1 = spawn(new_argv())
+    set_session(nvim1)
     rmdir(swapdir)
     lfs.mkdir(swapdir)
   end)
@@ -49,25 +61,15 @@ describe(':preserve', function()
     rmdir(swapdir)
   end)
 
-  it("saves to custom 'directory' and (R)ecovers #1836", function()
-    local testfile = 'Xtest_recover_file1'
-    -- Put swapdir at the start of the 'directory' list. #1836
-    -- Note: `set swapfile` *must* go after `set directory`: otherwise it may
-    -- attempt to create a swapfile in different directory.
-    local init = [[
-      set directory^=]]..swapdir:gsub([[\]], [[\\]])..[[//
-      set swapfile fileformat=unix undolevels=-1
-    ]]
-
+  local function setup_swapname()
     exec(init)
     command('edit! '..testfile)
     feed('isometext<esc>')
-    command('preserve')
     exec('redir => g:swapname | silent swapname | redir END')
+    return eval('g:swapname')
+  end
 
-    local swappath1 = eval('g:swapname')
-
-    os_kill(eval('getpid()'))
+  local function test_recover(swappath1)
     -- Start another Nvim instance.
     local nvim2 = spawn({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--embed'},
                                 true)
@@ -90,6 +92,35 @@ describe(':preserve', function()
     -- Verify that :swapname was not truncated (:help 'shortmess').
     ok(nil == string.find(swappath1, '%.%.%.'))
     ok(nil == string.find(swappath2, '%.%.%.'))
+  end
+
+  it('with :preserve and SIGKILL', function()
+    local swappath1 = setup_swapname()
+    command('preserve')
+    os_kill(eval('getpid()'))
+    test_recover(swappath1)
+  end)
+
+  it('closing stdio channel without :preserve #22096', function()
+    local swappath1 = setup_swapname()
+    nvim1:close()
+    test_recover(swappath1)
+  end)
+
+  it('killing TUI process without :preserve #22096', function()
+    helpers.skip(helpers.is_os('win'))
+    local screen = Screen.new()
+    screen:attach()
+    local child_server = new_pipename()
+    funcs.termopen({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--listen', child_server})
+    screen:expect({any = pesc('[No Name]')})  -- Wait for the child process to start.
+    local child_session = helpers.connect(child_server)
+    set_session(child_session)
+    local swappath1 = setup_swapname()
+    set_session(nvim1)
+    command('call chanclose(&channel)')  -- Kill the child process.
+    screen:expect({any = pesc('[Process exited 1]')})  -- Wait for the child process to stop.
+    test_recover(swappath1)
   end)
 
 end)
diff --git a/test/functional/shada/helpers.lua b/test/functional/shada/helpers.lua
index fb3ec4a87c..cd99d38345 100644
--- a/test/functional/shada/helpers.lua
+++ b/test/functional/shada/helpers.lua
@@ -33,6 +33,7 @@ local function reset(o)
 end
 
 local clear = function()
+  helpers.expect_exit(helpers.command, 'qall!')
   os.remove(tmpname)
 end
 
diff --git a/test/functional/shada/history_spec.lua b/test/functional/shada/history_spec.lua
index d1daf6c7cc..aa4106ad63 100644
--- a/test/functional/shada/history_spec.lua
+++ b/test/functional/shada/history_spec.lua
@@ -32,7 +32,6 @@ describe('ShaDa support code', function()
     nvim_command('rshada')
     eq('" Test 2', funcs.histget(':', -1))
     eq('" Test', funcs.histget(':', -2))
-    expect_exit(nvim_command, 'qall')
   end)
 
   it('respects &history when dumping',
diff --git a/test/functional/shada/shada_spec.lua b/test/functional/shada/shada_spec.lua
index 7b1ce5d0ca..24bd363e95 100644
--- a/test/functional/shada/shada_spec.lua
+++ b/test/functional/shada/shada_spec.lua
@@ -136,24 +136,6 @@ describe('ShaDa support code', function()
     eq(#msgpack, found)
   end)
 
-  it('does not write NONE file', function()
-    local session = spawn({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--embed',
-                           '--headless', '--cmd', 'qall'}, true)
-    session:close()
-    eq(nil, lfs.attributes('NONE'))
-    eq(nil, lfs.attributes('NONE.tmp.a'))
-  end)
-
-  it('does not read NONE file', function()
-    write_file('NONE', '\005\001\015\131\161na\162rX\194\162rc\145\196\001-')
-    local session = spawn({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--embed',
-                           '--headless'}, true)
-    set_session(session)
-    eq('', funcs.getreg('a'))
-    session:close()
-    os.remove('NONE')
-  end)
-
   local marklike = {[7]=true, [8]=true, [10]=true, [11]=true}
   local find_file = function(fname)
     local found = {}
@@ -263,3 +245,23 @@ describe('ShaDa support code', function()
     meths.set_option('shada', '')
   end)
 end)
+
+describe('ShaDa support code', function()
+  it('does not write NONE file', function()
+    local session = spawn({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--embed',
+                           '--headless', '--cmd', 'qall'}, true)
+    session:close()
+    eq(nil, lfs.attributes('NONE'))
+    eq(nil, lfs.attributes('NONE.tmp.a'))
+  end)
+
+  it('does not read NONE file', function()
+    write_file('NONE', '\005\001\015\131\161na\162rX\194\162rc\145\196\001-')
+    local session = spawn({nvim_prog, '-u', 'NONE', '-i', 'NONE', '--embed',
+                           '--headless'}, true)
+    set_session(session)
+    eq('', funcs.getreg('a'))
+    session:close()
+    os.remove('NONE')
+  end)
+end)