feat: add session name template support

This patch adds support for session name templates so
that you can switch multiple shpool sessions all at once.

In some sense this is a super-set of the 'shpool switch' FR.
I did lay some groundwork for implementing support for that
in this change, though I'm starting to wonder if templates are
good enough on their own. The only extra thing that switch would
bring is the ability to switch sessions you don't pre-declare as
switchable with a dedicated variable up front.

One thing worth bikeshedding: At first I was thinking '${var}' for
the substitution syntax, then realized that would be weird about
nesting when it comes to shells, so I switched to '#{var}' syntax,
but then I realized that's the comment char in shells, and wound up
on '@{var}'. I'm open to other symbol/syntax ideas.

BREAKING: this breaks shpool-protocol since we have a new chunk kind.

Author: ethanpailes

libshpool/src/attach.rs

@@ -35,9 +35,10 @@ use anyhow::{anyhow, Context};
 use nix::unistd;
 use shpool_protocol::{
     AttachHeader, AttachReplyHeader, AttachStatus, ConnectHeader, DetachReply, DetachRequest,
-    KillReply, KillRequest, ListReply, LogLevel, ResizeReply, Session, SessionMessageDetachReply,
-    SessionMessageReply, SessionMessageRequest, SessionMessageRequestPayload, SessionStatus,
-    SetLogLevelReply, SetLogLevelRequest, VersionHeader,
+    KillReply, KillRequest, ListReply, LogLevel, MaybeSwitch, ModifyVarReply, ModifyVarRequest,
+    ResizeReply, Session, SessionMessageDetachReply, SessionMessageReply, SessionMessageRequest,
+    SessionMessageRequestPayload, SessionStatus, SetLogLevelReply, SetLogLevelRequest,
+    VersionHeader,
 };
 use tracing::{debug, error, info, instrument, span, warn, Level};
 
@@ -79,6 +80,7 @@ pub struct Server {
         tracing_subscriber::filter::LevelFilter,
         tracing_subscriber::registry::Registry,
     >,
+    vars: Mutex<HashMap<String, String>>,
 }
 
 impl Server {
@@ -112,6 +114,7 @@ impl Server {
             hooks,
             daily_messenger,
             log_level_handle,
+            vars: HashMap::new().into(),
         }))
     }
 
@@ -209,6 +212,8 @@ impl Server {
             ConnectHeader::List => self.handle_list(stream),
             ConnectHeader::SessionMessage(header) => self.handle_session_message(stream, header),
             ConnectHeader::SetLogLevel(r) => self.handle_set_log_level(stream, r),
+            ConnectHeader::GetVars => self.handle_get_vars(stream),
+            ConnectHeader::ModifyVar(r) => self.handle_modify_var(stream, r),
         }
     }
 
@@ -609,6 +614,57 @@ impl Server {
         Ok(())
     }
 
+    #[instrument(skip_all)]
+    fn handle_get_vars(&self, mut stream: UnixStream) -> anyhow::Result<()> {
+        let maybe_switch = {
+            let var_map = self.vars.lock().unwrap();
+            let vars: Vec<(String, String)> =
+                var_map.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
+            shpool_protocol::MaybeSwitch { switch_to: None, vars }
+        };
+
+        write_reply(&mut stream, maybe_switch).context("writing maybe_switch reply")?;
+        Ok(())
+    }
+
+    #[instrument(skip_all)]
+    fn handle_modify_var(
+        &self,
+        mut stream: UnixStream,
+        request: ModifyVarRequest,
+    ) -> anyhow::Result<()> {
+        let maybe_switch = {
+            let mut vars = self.vars.lock().unwrap();
+            if let Some(val) = request.val {
+                vars.insert(request.var, val);
+            } else {
+                vars.remove(&request.var);
+            }
+
+            MaybeSwitch {
+                switch_to: None,
+                vars: vars.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
+            }
+        };
+
+        let mut ctls = Vec::new();
+        {
+            let shells = self.shells.lock().unwrap();
+            for (_, session) in shells.iter() {
+                ctls.push(Arc::clone(&session.shell_to_client_ctl));
+            }
+        }
+        for ctl in ctls.into_iter() {
+            let ctl = ctl.lock().unwrap();
+            ctl.maybe_switch
+                .send_timeout(maybe_switch.clone(), SESSION_MSG_TIMEOUT)
+                .context("broadcasting maybe_switch")?;
+        }
+
+        write_reply(&mut stream, ModifyVarReply {}).context("writing modify var reply")?;
+        Ok(())
+    }
+
     #[instrument(skip_all)]
     fn handle_kill(&self, mut stream: UnixStream, request: KillRequest) -> anyhow::Result<()> {
         let mut not_found_sessions = vec![];
@@ -983,13 +1039,18 @@ impl Server {
         let (heartbeat_tx, heartbeat_rx) = crossbeam_channel::bounded(0);
         let (heartbeat_ack_tx, heartbeat_ack_rx) = crossbeam_channel::bounded(0);
 
-        let shell_to_client_ctl = Arc::new(Mutex::new(shell::ReaderCtl {
+        // We make this buffered to avoid blocking during a broadcast. There is
+        // no ack chan so we can afford to buffer a bit.
+        let (maybe_switch_tx, maybe_switch_rx) = crossbeam_channel::bounded(10);
+
+        let shell_to_client_ctl = Arc::new(Mutex::new(shell::ShellToClientCtl {
             client_connection: client_connection_tx,
             client_connection_ack: client_connection_ack_rx,
             tty_size_change: tty_size_change_tx,
             tty_size_change_ack: tty_size_change_ack_rx,
             heartbeat: heartbeat_tx,
             heartbeat_ack: heartbeat_ack_rx,
+            maybe_switch: maybe_switch_tx,
         }));
 
         let mut session_inner = shell::SessionInner {
@@ -1023,6 +1084,7 @@ impl Server {
                 tty_size_change_ack: tty_size_change_ack_tx,
                 heartbeat: heartbeat_rx,
                 heartbeat_ack: heartbeat_ack_tx,
+                maybe_switch: maybe_switch_rx,
                 child_exit_notifier: shell_to_client_child_exit_notifier,
             })?);
 

libshpool/src/daemon/server.rs

@@ -28,12 +28,13 @@ use std::{
 
 use anyhow::{anyhow, Context};
 use nix::{poll, poll::PollFlags, sys::signal, unistd::Pid};
-use shpool_protocol::{Chunk, ChunkKind, TtySize};
+use shpool_protocol::{Chunk, ChunkKind, MaybeSwitch, TtySize};
 use tracing::{debug, error, info, instrument, span, trace, warn, Level};
 
 use crate::{
     common, consts,
     daemon::{config, exit_notify::ExitNotifier, keybindings, pager::PagerCtl, prompt, show_motd},
+    protocol,
     protocol::ChunkExt as _,
     session_restore, test_hooks,
     tty::TtySizeExt as _,
@@ -74,7 +75,7 @@ pub struct Session {
     pub lifecycle_timestamps: Mutex<SessionLifecycleTimestamps>,
     pub child_pid: libc::pid_t,
     pub child_exit_notifier: Arc<ExitNotifier>,
-    pub shell_to_client_ctl: Arc<Mutex<ReaderCtl>>,
+    pub shell_to_client_ctl: Arc<Mutex<ShellToClientCtl>>,
     pub pager_ctl: Arc<Mutex<Option<PagerCtl>>>,
     /// Mutable state with the lock held by the servicing handle_attach thread
     /// while a tty is attached to the session. Probing the mutex can be used
@@ -109,7 +110,7 @@ impl Session {
 #[derive(Debug)]
 pub struct SessionInner {
     pub name: String, // to improve logging
-    pub shell_to_client_ctl: Arc<Mutex<ReaderCtl>>,
+    pub shell_to_client_ctl: Arc<Mutex<ShellToClientCtl>>,
     pub pty_master: shpool_pty::fork::Fork,
     pub client_stream: Option<UnixStream>,
     pub config: config::Manager,
@@ -207,6 +208,7 @@ pub struct ShellToClientArgs {
     pub tty_size_change: crossbeam_channel::Receiver<TtySize>,
     pub tty_size_change_ack: crossbeam_channel::Sender<()>,
     pub heartbeat: crossbeam_channel::Receiver<()>,
+    pub maybe_switch: crossbeam_channel::Receiver<MaybeSwitch>,
     // true if the client is still live, false if it has hung up on us
     pub heartbeat_ack: crossbeam_channel::Sender<bool>,
     pub child_exit_notifier: Arc<ExitNotifier>,
@@ -393,6 +395,42 @@ impl SessionInner {
                         args.heartbeat_ack.send(client_present)
                             .context("sending heartbeat ack")?;
                     }
+                    recv(args.maybe_switch) -> maybe_switch => {
+                        let maybe_switch = match maybe_switch {
+                            Ok(ms) => ms,
+                            Err(e) => {
+                                error!("error recving MaybeSwitch: {:?}", e);
+                                continue;
+                            },
+                        };
+
+                        let conn = if let ClientConnectionMsg::New(c) = &mut client_conn {
+                            c
+                        } else {
+                            info!("got MaybeSwitch, but no attached client, dropping");
+                            continue;
+                        };
+
+                        let mut encoded = Vec::new();
+                        if let Err(e) = protocol::encode_to(&maybe_switch, &mut encoded) {
+                            error!("error encoding MaybeSwitch: {:?}", e);
+                            continue;
+                        }
+
+                        let chunk = Chunk { kind: ChunkKind::MaybeSwitch, buf: &encoded[..] };
+                        match chunk.write_to(&mut conn.sink).and_then(|_| conn.sink.flush()) {
+                            Ok(_) => {
+                                trace!("wrote MaybeSwitch");
+                            }
+                            Err(e) if e.kind() == io::ErrorKind::BrokenPipe => {
+                                trace!("writing MaybeSwitch: client hangup: {:?}", e);
+                            }
+                            Err(e) => {
+                                error!("unexpected IO error while writing heartbeat: {}", e);
+                                return Err(e).context("writing MaybeSwitch")?;
+                            }
+                        }
+                    }
 
                     // make this select non-blocking so we spend most of our time parked
                     // in poll
@@ -1003,7 +1041,7 @@ impl SessionInner {
 /// Shared between the session struct (for calls originating with the cli)
 /// and the session inner struct (for calls resulting from keybindings).
 #[derive(Debug)]
-pub struct ReaderCtl {
+pub struct ShellToClientCtl {
     /// A control channel for the shell->client thread. Whenever a new client
     /// dials in, the output stream for that client must be attached to the
     /// shell->client thread by sending it down this channel. A disconnect
@@ -1028,6 +1066,12 @@ pub struct ReaderCtl {
     // True if the client is still listening, false if it has hung up
     // on us.
     pub heartbeat_ack: crossbeam_channel::Receiver<bool>,
+
+    /// A control channel telling the shell->client thread to
+    /// broadcast the given MaybeSwitch. There is no ack channel
+    /// because we just blast this out and the caller doesn't need
+    /// to know about completion.
+    pub maybe_switch: crossbeam_channel::Sender<MaybeSwitch>,
 }
 
 /// Given a buffer, a length after which the data is not valid, a list of

libshpool/src/daemon/shell.rs

@@ -43,9 +43,11 @@ mod list;
 mod protocol;
 mod session_restore;
 mod set_log_level;
+mod template;
 mod test_hooks;
 mod tty;
 mod user;
+mod var;
 
 /// The command line arguments that shpool expects.
 /// These can be directly parsed with clap or manually
@@ -210,6 +212,54 @@ needs debugging, but would be clobbered by a restart.")]
         #[clap(help = "new log level")]
         level: shpool_protocol::LogLevel,
     },
+
+    #[clap(about = "Manipulate template variables
+
+shpool session names can include #{variables} which are resolved via
+an environment stored globally in the shpool daemon. This command
+manipulates that environment.
+
+The main usecase for templated session names is the ability to switch
+multiple shpool sessions to new targets at the same time. For example,
+you might have a `shpool attach -f '#{workspace}-edit'` session and
+a `shpool attach -f '#{workspace}-term'` session. To switch both
+sessions from the fun-feature workspace to the key-bugfix workspace,
+you could just do `shpool var set workspace key-bugfix`.
+")]
+    #[non_exhaustive]
+    Var {
+        #[clap(subcommand)]
+        command: VarCommands,
+    },
+}
+
+/// The subcommds of the var command.
+#[derive(Subcommand, Debug)]
+#[non_exhaustive]
+pub enum VarCommands {
+    #[clap(about = "List the variables
+
+This command dumps out the whole variable list with
+both vars and values in a JSON object using vars as keys.")]
+    List {
+        #[clap(short, long, help = "Output as JSON")]
+        json: bool,
+    },
+    #[clap(about = "Get a variable
+
+This returns the raw value of the given variable.")]
+    #[non_exhaustive]
+    Get { var: String },
+    #[clap(about = "Set a variable
+
+This updates the value of the given variable.")]
+    #[non_exhaustive]
+    Set { var: String, val: String },
+    #[clap(about = "Unset a variable
+
+This removes the given variable from the environment.")]
+    #[non_exhaustive]
+    Unset { var: String },
 }
 
 impl Args {
@@ -382,6 +432,7 @@ pub fn run(args: Args, hooks: Option<Box<dyn hooks::Hooks + Send + Sync>>) -> an
         Commands::Kill { sessions } => kill::run(sessions, socket),
         Commands::List { json } => list::run(socket, json),
         Commands::SetLogLevel { level } => set_log_level::run(level, socket),
+        Commands::Var { command } => var::run(socket, command),
     };
 
     if let Err(err) = res {