Fix tmux panic (#165)

Fix `tmux` sessions

This will create a new session named `topgrade`, `topgrade-1`,
`topgrade-2`, using the first nonexistent session name it finds. That
session will have a window in it named `topgrade` in which `topgrade` is
run. If `topgrade --tmux` is being run from within tmux, it won't attach
to the new tmux session. If the user is not currently in tmux, it will
attach to the newly-created session.

Co-authored-by: Thomas Schönauer <37108907+DottoDev@users.noreply.github.com>

src/execution_context.rs

@@ -6,6 +6,7 @@ use crate::{config::Config, executor::Executor};
 use color_eyre::eyre::Result;
 use directories::BaseDirs;
 use std::path::{Path, PathBuf};
+use std::sync::Mutex;
 
 pub struct ExecutionContext<'a> {
     run_type: RunType,
@@ -13,6 +14,10 @@ pub struct ExecutionContext<'a> {
     git: &'a Git,
     config: &'a Config,
     base_dirs: &'a BaseDirs,
+    /// Name of a tmux session to execute commands in, if any.
+    /// This is used in `./steps/remote/ssh.rs`, where we want to run `topgrade` in a new
+    /// tmux window for each remote.
+    tmux_session: Mutex<Option<String>>,
 }
 
 impl<'a> ExecutionContext<'a> {
@@ -22,13 +27,14 @@ impl<'a> ExecutionContext<'a> {
         git: &'a Git,
         config: &'a Config,
         base_dirs: &'a BaseDirs,
-    ) -> ExecutionContext<'a> {
-        ExecutionContext {
+    ) -> Self {
+        Self {
             run_type,
             sudo,
             git,
             config,
             base_dirs,
+            tmux_session: Mutex::new(None),
         }
     }
 
@@ -67,4 +73,12 @@ impl<'a> ExecutionContext<'a> {
     pub fn base_dirs(&self) -> &BaseDirs {
         self.base_dirs
     }
+
+    pub fn set_tmux_session(&self, session_name: String) {
+        self.tmux_session.lock().unwrap().replace(session_name);
+    }
+
+    pub fn get_tmux_session(&self) -> Option<String> {
+        self.tmux_session.lock().unwrap().clone()
+    }
 }

src/main.rs

@@ -82,7 +82,8 @@ fn run() -> Result<()> {
     if config.run_in_tmux() && env::var("TOPGRADE_INSIDE_TMUX").is_err() {
         #[cfg(unix)]
         {
-            tmux::run_in_tmux(config.tmux_arguments()?);
+            tmux::run_in_tmux(config.tmux_arguments()?)?;
+            return Ok(());
         }
     }
 

src/steps/remote/ssh.rs

@@ -26,7 +26,7 @@ pub fn ssh_step(ctx: &ExecutionContext, hostname: &str) -> Result<()> {
         #[cfg(unix)]
         {
             prepare_async_ssh_command(&mut args);
-            crate::tmux::run_command(ctx, &shell_words::join(args))?;
+            crate::tmux::run_command(ctx, hostname, &shell_words::join(args))?;
             Err(SkipStep(String::from("Remote Topgrade launched in Tmux")).into())
         }
 

src/steps/tmux.rs

@@ -1,7 +1,9 @@
 use std::env;
 use std::path::PathBuf;
-use std::process::{exit, Command};
+use std::process::Command;
 
+use color_eyre::eyre::eyre;
+use color_eyre::eyre::Context;
 use color_eyre::eyre::Result;
 use directories::BaseDirs;
 
@@ -57,58 +59,121 @@ impl Tmux {
             .success())
     }
 
-    fn new_session(&self, session_name: &str) -> Result<bool> {
-        Ok(self
+    /// Create a new tmux session with the given name, running the given command.
+    /// The command is passed to `sh` (see "shell-command arguments are sh(1) commands" in the
+    /// `tmux(1)` man page).
+    fn new_session(&self, session_name: &str, window_name: &str, command: &str) -> Result<()> {
+        let _ = self
             .build()
-            .args(["new-session", "-d", "-s", session_name, "-n", "dummy"])
-            .output_checked_with(|_| Ok(()))?
-            .status
-            .success())
+            // `-d`: initial size comes from the global `default-size` option (instead
+            //       of passing `-x` and `-y` arguments.
+            //       (What do those even do?)
+            // `-s`: session name
+            // `-n`: window name (always `topgrade`)
+            .args(["new-session", "-d", "-s", session_name, "-n", window_name, command])
+            .output_checked()?;
+        Ok(())
     }
 
-    fn run_in_session(&self, command: &str) -> Result<()> {
+    /// Like [`new_session`] but it appends a digit to the session name (if necessary) to
+    /// avoid duplicate session names.
+    ///
+    /// The session name is returned.
+    fn new_unique_session(&self, session_name: &str, window_name: &str, command: &str) -> Result<String> {
+        let mut session = session_name.to_owned();
+        for i in 1.. {
+            if !self
+                .has_session(&session)
+                .context("Error determining if a tmux session exists")?
+            {
+                self.new_session(&session, window_name, command)
+                    .context("Error running Topgrade in tmux")?;
+                return Ok(session);
+            }
+            session = format!("{session_name}-{i}");
+        }
+        unreachable!()
+    }
+
+    /// Create a new window in the given tmux session, running the given command.
+    fn new_window(&self, session_name: &str, window_name: &str, command: &str) -> Result<()> {
         self.build()
-            .args(["new-window", "-t", "topgrade", command])
+            // `-d`: initial size comes from the global `default-size` option (instead
+            //       of passing `-x` and `-y` arguments.
+            //       (What do those even do?)
+            // `-s`: session name
+            // `-n`: window name
+            .args([
+                "new-window",
+                "-a",
+                "-t",
+                &format!("{session_name}:{window_name}"),
+                "-n",
+                window_name,
+                command,
+            ])
+            .env_remove("TMUX")
             .status_checked()
     }
+
+    fn window_indices(&self, session_name: &str) -> Result<Vec<usize>> {
+        self.build()
+            .args(["list-windows", "-F", "#{window_index}", "-t", session_name])
+            .output_checked_utf8()?
+            .stdout
+            .lines()
+            .map(|l| l.parse())
+            .collect::<Result<Vec<usize>, _>>()
+            .context("Failed to compute tmux windows")
+    }
 }
 
-pub fn run_in_tmux(args: Vec<String>) -> ! {
+pub fn run_in_tmux(args: Vec<String>) -> Result<()> {
     let command = {
         let mut command = vec![
             String::from("env"),
             String::from("TOPGRADE_KEEP_END=1"),
             String::from("TOPGRADE_INSIDE_TMUX=1"),
         ];
+        // TODO: Should we use `topgrade` instead of the first argument here, which may be
+        // a local path?
         command.extend(env::args());
         shell_words::join(command)
     };
 
     let tmux = Tmux::new(args);
 
-    if !tmux.has_session("topgrade").expect("Error detecting a tmux session") {
-        tmux.new_session("topgrade").expect("Error creating a tmux session");
-    }
-
-    tmux.run_in_session(&command).expect("Error running Topgrade in tmux");
-    tmux.build()
-        .args(["kill-window", "-t", "topgrade:dummy"])
-        .output_checked()
-        .expect("Error killing the dummy tmux window");
+    // Find an unused session and run `topgrade` in it with the current command's arguments.
+    let session_name = "topgrade";
+    let window_name = "topgrade";
+    let session = tmux.new_unique_session(session_name, window_name, &command)?;
 
+    // Only attach to the newly-created session if we're not currently in a tmux session.
     if env::var("TMUX").is_err() {
-        let err = tmux.build().args(["attach", "-t", "topgrade"]).exec();
-        panic!("{:?}", err);
+        let err = tmux.build().args(["attach-session", "-t", &session]).exec();
+        Err(eyre!("{err}")).context("Failed to `execvp(3)` tmux")
     } else {
         println!("Topgrade launched in a new tmux session");
-        exit(0);
+        Ok(())
     }
 }
 
-pub fn run_command(ctx: &ExecutionContext, command: &str) -> Result<()> {
-    Tmux::new(ctx.config().tmux_arguments()?)
-        .build()
-        .args(["new-window", "-a", "-t", "topgrade:1", command])
-        .env_remove("TMUX")
-        .status_checked()
+pub fn run_command(ctx: &ExecutionContext, window_name: &str, command: &str) -> Result<()> {
+    let tmux = Tmux::new(ctx.config().tmux_arguments()?);
+
+    match ctx.get_tmux_session() {
+        Some(session_name) => {
+            let indices = tmux.window_indices(&session_name)?;
+            let last_window = indices
+                .iter()
+                .last()
+                .ok_or_else(|| eyre!("tmux session {session_name} has no windows"))?;
+            tmux.new_window(&session_name, &format!("{last_window}"), command)?;
+        }
+        None => {
+            let name = tmux.new_unique_session("topgrade", window_name, command)?;
+            ctx.set_tmux_session(name);
+        }
+    }
+    Ok(())
 }