Support specifying an exit status on `quit` commands

The current client exit status can be specified as an optional
parameter, is nothing is given the exit status will be 0.

Fixes #1230

8 files changed, 82 insertions, 49 deletions

diff --git a/src/client.cc b/src/client.cc
index 22e574bb..dea670d6 100644
--- a/src/client.cc
+++ b/src/client.cc
@@ -26,8 +26,10 @@ Client::Client(std::unique_ptr<UserInterface>&& ui,
                std::unique_ptr<Window>&& window,
                SelectionList selections,
                EnvVarMap env_vars,
-               String name)
+               String name,
+               OnExitCallback on_exit)
     : m_ui{std::move(ui)}, m_window{std::move(window)},
+      m_on_exit{std::move(on_exit)},
       m_input_handler{std::move(selections), Context::Flags::None,
                       std::move(name)},
       m_env_vars(std::move(env_vars))
diff --git a/src/client.hh b/src/client.hh
index daf8aa02..d870871b 100644
--- a/src/client.hh
+++ b/src/client.hh
@@ -21,15 +21,17 @@ enum class EventMode;
 enum class InfoStyle;
 enum class MenuStyle;
 
-
 class Client : public SafeCountable, public OptionManagerWatcher
 {
 public:
+    using OnExitCallback = std::function<void (int status)>;
+
     Client(std::unique_ptr<UserInterface>&& ui,
            std::unique_ptr<Window>&& window,
            SelectionList selections,
            EnvVarMap env_vars,
-           String name);
+           String name,
+           OnExitCallback on_exit);
     ~Client();
 
     Client(Client&&) = delete;
@@ -65,6 +67,8 @@ public:
     Buffer* last_buffer() const { return m_last_buffer.get(); }
     void set_last_buffer(Buffer* last_buffer) { m_last_buffer = last_buffer; }
 
+    void exit(int status) { m_on_exit(status); }
+
 private:
     void on_option_changed(const Option& option) override;
 
@@ -79,6 +83,8 @@ private:
     std::unique_ptr<UserInterface> m_ui;
     std::unique_ptr<Window> m_window;
 
+    OnExitCallback m_on_exit;
+
     EnvVarMap m_env_vars;
 
     InputHandler m_input_handler;
diff --git a/src/client_manager.cc b/src/client_manager.cc
index 4eeb4f60..6525b854 100644
--- a/src/client_manager.cc
+++ b/src/client_manager.cc
@@ -40,13 +40,14 @@ String ClientManager::generate_name() const
 
 Client* ClientManager::create_client(std::unique_ptr<UserInterface>&& ui,
                                      EnvVarMap env_vars, StringView init_cmds,
-                                     Optional<BufferCoord> init_coord)
+                                     Optional<BufferCoord> init_coord,
+                                     Client::OnExitCallback on_exit)
 {
     Buffer& buffer = BufferManager::instance().get_first_buffer();
     WindowAndSelections ws = get_free_window(buffer);
     Client* client = new Client{std::move(ui), std::move(ws.window),
                                 std::move(ws.selections), std::move(env_vars),
-                                generate_name()};
+                                generate_name(), std::move(on_exit)};
     m_clients.emplace_back(client);
 
     if (init_coord)
@@ -88,7 +89,7 @@ void ClientManager::process_pending_inputs() const
     }
 }
 
-void ClientManager::remove_client(Client& client, bool graceful)
+void ClientManager::remove_client(Client& client, bool graceful, int status)
 {
     auto it = find(m_clients, &client);
     if (it == m_clients.end())
@@ -96,6 +97,7 @@ void ClientManager::remove_client(Client& client, bool graceful)
         kak_assert(contains(m_client_trash, &client));
         return;
     }
+    client.exit(status);
     m_client_trash.push_back(std::move(*it));
     m_clients.erase(it);
 
diff --git a/src/client_manager.hh b/src/client_manager.hh
index 35a1376f..f0a3a59d 100644
--- a/src/client_manager.hh
+++ b/src/client_manager.hh
@@ -21,7 +21,8 @@ public:
 
     Client* create_client(std::unique_ptr<UserInterface>&& ui,
                           EnvVarMap env_vars, StringView init_cmds,
-                          Optional<BufferCoord> init_coord);
+                          Optional<BufferCoord> init_coord,
+                          Client::OnExitCallback on_exit);
 
     bool   empty() const { return m_clients.empty(); }
     size_t count() const { return m_clients.size(); }
@@ -39,7 +40,7 @@ public:
     Client*  get_client_ifp(StringView name);
     Client&  get_client(StringView name);
     bool validate_client_name(StringView name) const;
-    void remove_client(Client& client, bool graceful);
+    void remove_client(Client& client, bool graceful, int status);
 
     using ClientList = Vector<std::unique_ptr<Client>, MemoryDomain::Client>;
     using iterator = ClientList::const_iterator;
diff --git a/src/commands.cc b/src/commands.cc
index fcce5e0f..848d4406 100644
--- a/src/commands.cc
+++ b/src/commands.cc
@@ -423,24 +423,26 @@ const CommandDesc force_kill_cmd = {
 };
 
 template<bool force>
-void quit(Context& context)
+void quit(const ParametersParser& parser, Context& context, const ShellContext&)
 {
     if (not force and ClientManager::instance().count() == 1)
         ensure_all_buffers_are_saved();
 
-    ClientManager::instance().remove_client(context.client(), true);
+    const int status = parser.positional_count() > 0 ? str_to_int(parser[0]) : 0;
+    ClientManager::instance().remove_client(context.client(), true, status);
 }
 
 const CommandDesc quit_cmd = {
     "quit",
     "q",
     "quit current client, and the kakoune session if the client is the last "
-    "(if not running in daemon mode)",
-    no_params,
+    "(if not running in daemon mode). An optional integer parameter can set "
+    "the client exit status",
+    { {}, ParameterDesc::Flags::SwitchesAsPositional, 0, 1 },
     CommandFlags::None,
     CommandHelper{},
     CommandCompleter{},
-    [](const ParametersParser&, Context& context, const ShellContext&){ quit<false>(context); }
+    quit<false>
 };
 
 const CommandDesc force_quit_cmd = {
@@ -448,27 +450,29 @@ const CommandDesc force_quit_cmd = {
     "q!",
     "quit current client, and the kakoune session if the client is the last "
     "(if not running in daemon mode). force quit even if the client is the "
-    "last and some buffers are not saved.",
-    no_params,
+    "last and some buffers are not saved. An optional integer parameter can "
+    "set the client exit status",
+    { {}, ParameterDesc::Flags::SwitchesAsPositional, 0, 1 },
     CommandFlags::None,
     CommandHelper{},
     CommandCompleter{},
-    [](const ParametersParser&, Context& context, const ShellContext&){ quit<true>(context); }
+    quit<true>
 };
 
 template<bool force>
 void write_quit(const ParametersParser& parser, Context& context,
                 const ShellContext& shell_context)
 {
-    write_buffer(parser, context, shell_context);
-    quit<force>(context);
+    write_buffer({{}, {}}, context, shell_context);
+    quit<force>(parser, context, shell_context);
 }
 
 const CommandDesc write_quit_cmd = {
     "write-quit",
     "wq",
-    "write current buffer and quit current client",
-    no_params,
+    "write current buffer and quit current client. An optional integer"
+    "parameter can set the client exit status",
+    { {}, ParameterDesc::Flags::SwitchesAsPositional, 0, 1 },
     CommandFlags::None,
     CommandHelper{},
     CommandCompleter{},
@@ -479,8 +483,8 @@ const CommandDesc force_write_quit_cmd = {
     "write-quit!",
     "wq!",
     "write current buffer and quit current client, even if other buffers are "
-    "not saved",
-    no_params,
+    "not saved. An optional integer parameter can set the client exit status",
+    { {}, ParameterDesc::Flags::SwitchesAsPositional, 0, 1 },
     CommandFlags::None,
     CommandHelper{},
     CommandCompleter{},
@@ -490,15 +494,16 @@ const CommandDesc force_write_quit_cmd = {
 const CommandDesc write_all_quit_cmd = {
     "write-all-quit",
     "waq",
-    "write all buffers associated to a file and quit current client",
-    no_params,
+    "write all buffers associated to a file and quit current client. An "
+    "optional integer parameter can set the client exit status",
+    { {}, ParameterDesc::Flags::SwitchesAsPositional, 0, 1 },
     CommandFlags::None,
     CommandHelper{},
     CommandCompleter{},
-    [](const ParametersParser& parser, Context& context, const ShellContext&)
+    [](const ParametersParser& parser, Context& context, const ShellContext& shell_context)
     {
         write_all_buffers(context);
-        quit<false>(context);
+        quit<false>(parser, context, shell_context);
     }
 };
 
diff --git a/src/main.cc b/src/main.cc
index 33263227..0d4f0c78 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -328,6 +328,7 @@ void register_options()
 }
 
 static Client* local_client = nullptr;
+static int local_client_exit = 0;
 static UserInterface* local_ui = nullptr;
 static bool convert_to_client_pending = false;
 
@@ -403,7 +404,7 @@ std::unique_ptr<UserInterface> create_local_ui(UIType ui_type)
             kak_assert(not local_ui);
             local_ui = this;
             m_old_sighup = set_signal_handler(SIGHUP, [](int) {
-                ClientManager::instance().remove_client(*local_client, false);
+                ClientManager::instance().remove_client(*local_client, false, -1);
                 static_cast<LocalUI*>(local_ui)->on_sighup();
             });
 
@@ -441,7 +442,7 @@ std::unique_ptr<UserInterface> create_local_ui(UIType ui_type)
                 if (fork_server_to_background())
                 {
                     this->NCursesUI::~NCursesUI();
-                    exit(0);
+                    exit(local_client_exit);
                 }
             }
         }
@@ -475,16 +476,15 @@ int run_client(StringView session, StringView client_init,
     {
         EventManager event_manager;
         RemoteClient client{session, make_ui(ui_type), get_env_vars(), client_init, std::move(init_coord)};
-        while (true)
+        while (not client.exit_status())
             event_manager.handle_next_events(EventMode::Normal);
+        return *client.exit_status();
     }
     catch (disconnected& e)
     {
-        if (!e.m_graceful)
-            write_stderr(format("{}\ndisconnecting\n", e.what()));
-        return e.m_graceful ? 0 : -1;
+        write_stderr(format("{}\ndisconnecting\n", e.what()));
+        return -1;
     }
-    return 0;
 }
 
 struct convert_to_client_mode
@@ -614,7 +614,8 @@ int run_server(StringView session, StringView server_init,
         if (not (flags & ServerFlags::Daemon))
         {
             local_client = client_manager.create_client(
-                 create_local_ui(ui_type), get_env_vars(), client_init, std::move(init_coord));
+                 create_local_ui(ui_type), get_env_vars(), client_init, std::move(init_coord),
+                 [](int status) { local_client_exit = status; });
 
             if (startup_error)
                 local_client->print_status({
@@ -640,7 +641,7 @@ int run_server(StringView session, StringView server_init,
                 String buffer_name = local_client->context().buffer().name();
                 String selections = selection_list_to_string(local_client->context().selections());
 
-                ClientManager::instance().remove_client(*local_client, true);
+                ClientManager::instance().remove_client(*local_client, true, 0);
                 client_manager.clear_client_trash();
                 convert_to_client_pending = false;
 
@@ -660,7 +661,7 @@ int run_server(StringView session, StringView server_init,
         global_scope.hooks().run_hook("KakEnd", "", empty_context);
     }
 
-    return 0;
+    return local_client_exit;
 }
 
 int run_filter(StringView keystr, StringView commands, ConstArrayView<StringView> files, bool quiet, StringView suffix_backup)
@@ -743,7 +744,7 @@ int run_pipe(StringView session)
     catch (disconnected& e)
     {
         write_stderr(format("{}\ndisconnecting\n", e.what()));
-        return e.m_graceful ? 0 : -1;
+        return -1;
     }
     return 0;
 }
diff --git a/src/remote.cc b/src/remote.cc
index 7b3ffccb..d88195f8 100644
--- a/src/remote.cc
+++ b/src/remote.cc
@@ -38,7 +38,8 @@ enum class MessageType : uint8_t
     SetCursor,
     Refresh,
     SetOptions,
-    Key
+    Exit,
+    Key,
 };
 
 class MsgWriter
@@ -252,8 +253,7 @@ private:
         kak_assert(m_write_pos + size <= m_stream.size());
         int res = ::read(sock, m_stream.data() + m_write_pos, size);
         if (res <= 0)
-            throw disconnected{res ? format("socket read failed: {}", strerror(errno))
-                                   : "peer disconnected", res == 0};
+            throw disconnected{format("socket read failed: {}", strerror(errno))};
         m_write_pos += res;
     }
 
@@ -352,6 +352,8 @@ public:
     void set_client(Client* client) { m_client = client; }
     Client* client() const { return m_client.get(); }
 
+    void exit(int status);
+
 private:
     FDWatcher     m_socket_watcher;
     MsgReader     m_reader;
@@ -368,8 +370,7 @@ static bool send_data(int fd, RemoteBuffer& buffer)
     {
       int res = ::write(fd, buffer.data(), buffer.size());
       if (res <= 0)
-          throw disconnected{res ? format("socket write failed: {}", strerror(errno))
-                                 : "peer disconnected", res == 0};
+          throw disconnected{format("socket write failed: {}", strerror(errno))};
       buffer.erase(buffer.begin(), buffer.begin() + res);
     }
     return buffer.empty();
@@ -393,7 +394,7 @@ RemoteUI::RemoteUI(int socket, DisplayCoord dimensions)
 
                    if (m_reader.type() != MessageType::Key)
                    {
-                      ClientManager::instance().remove_client(*m_client, false);
+                      ClientManager::instance().remove_client(*m_client, false, -1);
                       return;
                    }
 
@@ -407,7 +408,7 @@ RemoteUI::RemoteUI(int socket, DisplayCoord dimensions)
           catch (const disconnected& err)
           {
               write_to_debug_buffer(format("Error while transfering remote messages: {}", err.what()));
-              ClientManager::instance().remove_client(*m_client, false);
+              ClientManager::instance().remove_client(*m_client, false, -1);
           }
       }),
       m_dimensions(dimensions)
@@ -417,6 +418,9 @@ RemoteUI::RemoteUI(int socket, DisplayCoord dimensions)
 
 RemoteUI::~RemoteUI()
 {
+    // Try to send the remaining data if possible, as it might contain the desired exit status
+    send_data(m_socket_watcher.fd(), m_send_buffer);
+
     write_to_debug_buffer(format("remote client disconnected: {}", m_socket_watcher.fd()));
     m_socket_watcher.close_fd();
 }
@@ -510,6 +514,13 @@ void RemoteUI::set_ui_options(const Options& options)
     m_socket_watcher.events() |= FdEvents::Write;
 }
 
+void RemoteUI::exit(int status)
+{
+    MsgWriter msg{m_send_buffer, MessageType::Exit};
+    msg.write(status);
+    m_socket_watcher.events() |= FdEvents::Write;
+}
+
 static sockaddr_un session_addr(StringView session)
 {
     sockaddr_un addr;
@@ -637,6 +648,11 @@ RemoteClient::RemoteClient(StringView session, std::unique_ptr<UserInterface>&&
             case MessageType::SetOptions:
                 m_ui->set_ui_options(reader.read_hash_map<String, String, MemoryDomain::Options>());
                 break;
+            case MessageType::Exit:
+                m_exit_status = reader.read<int>();
+                m_socket_watcher->close_fd();
+                m_socket_watcher.reset();
+                return; // This lambda is now dead
             default:
                 kak_assert(false);
             }
@@ -697,7 +713,8 @@ private:
                 auto* ui = new RemoteUI{sock, dimensions};
                 if (auto* client = ClientManager::instance().create_client(
                                        std::unique_ptr<UserInterface>(ui),
-                                       std::move(env_vars), init_cmds, init_coord))
+                                       std::move(env_vars), init_cmds, init_coord,
+                                       [ui](int status) { ui->exit(status); }))
                     ui->set_client(client);
 
                 Server::instance().remove_accepter(this);
diff --git a/src/remote.hh b/src/remote.hh
index b0cad37c..84382dc4 100644
--- a/src/remote.hh
+++ b/src/remote.hh
@@ -13,10 +13,7 @@ namespace Kakoune
 
 struct disconnected : runtime_error
 {
-    disconnected(String what, bool graceful = false)
-      : runtime_error{std::move(what)}, m_graceful{graceful} {}
-
-    const bool m_graceful;
+    using runtime_error::runtime_error;
 };
 
 class FDWatcher;
@@ -36,10 +33,12 @@ public:
                  const EnvVarMap& env_vars, StringView init_command,
                  Optional<BufferCoord> init_coord);
 
+    const Optional<int>& exit_status() const { return m_exit_status; }
 private:
     std::unique_ptr<UserInterface> m_ui;
     std::unique_ptr<FDWatcher>     m_socket_watcher;
     RemoteBuffer                   m_send_buffer;
+    Optional<int>                  m_exit_status;
 };
 
 void send_command(StringView session, StringView command);