Split pty::Handler into 2 traits

src/cmd/session.rs

@@ -29,7 +29,7 @@ use crate::logger;
 use crate::notifier::{self, Notifier, NullNotifier};
 use crate::pty;
 use crate::server;
-use crate::session::{self, KeyBindings, Session};
+use crate::session::{self, KeyBindings, SessionStarter};
 use crate::stream::Stream;
 use crate::tty::{DevTty, FixedSizeTty, NullTty};
 use crate::util;
@@ -143,9 +143,9 @@ impl cli::Session {
         let exec_extra_env = build_exec_extra_env(relay_id.as_ref());
 
         {
-            let mut session = Session::new(outputs, record_input, keys, notifier);
+            let starter = SessionStarter::new(outputs, record_input, keys, notifier);
             let mut tty = self.get_tty()?;
-            pty::exec(&exec_command, &exec_extra_env, &mut tty, &mut session)?;
+            pty::exec(&exec_command, &exec_extra_env, &mut tty, starter)?;
         }
 
         runtime.block_on(async {

src/pty.rs

@@ -8,7 +8,7 @@ use std::os::fd::{BorrowedFd, OwnedFd};
 use std::os::unix::io::{AsRawFd, FromRawFd};
 use std::time::{Duration, Instant};
 
-use anyhow::{bail, Result};
+use anyhow::bail;
 use nix::errno::Errno;
 use nix::libc::EIO;
 use nix::sys::select::{select, FdSet};
@@ -24,26 +24,33 @@ use crate::tty::{Theme, Tty, TtySize};
 
 type ExtraEnv = HashMap<String, String>;
 
+pub trait HandlerStarter<H: Handler> {
+    fn start(self, tty_size: TtySize, theme: Option<Theme>) -> H;
+}
+
 pub trait Handler {
-    fn start(&mut self, tty_size: TtySize, theme: Option<Theme>);
     fn output(&mut self, time: Duration, data: &[u8]) -> bool;
     fn input(&mut self, time: Duration, data: &[u8]) -> bool;
     fn resize(&mut self, time: Duration, tty_size: TtySize) -> bool;
+    fn stop(self) -> Self;
 }
 
-pub fn exec<S: AsRef<str>, T: Tty + ?Sized, H: Handler>(
+pub fn exec<S: AsRef<str>, T: Tty, H: Handler, R: HandlerStarter<H>>(
     command: &[S],
     extra_env: &ExtraEnv,
     tty: &mut T,
-    handler: &mut H,
-) -> Result<i32> {
+    handler_starter: R,
+) -> anyhow::Result<(i32, H)> {
     let winsize = tty.get_size();
     let epoch = Instant::now();
-    handler.start(winsize.into(), tty.get_theme());
+    let mut handler = handler_starter.start(winsize.into(), tty.get_theme());
     let result = unsafe { pty::forkpty(Some(&winsize), None) }?;
 
     match result.fork_result {
-        ForkResult::Parent { child } => handle_parent(result.master, child, tty, handler, epoch),
+        ForkResult::Parent { child } => {
+            handle_parent(result.master, child, tty, &mut handler, epoch)
+                .map(|code| (code, handler.stop()))
+        }
 
         ForkResult::Child => {
             handle_child(command, extra_env)?;
@@ -52,13 +59,13 @@ pub fn exec<S: AsRef<str>, T: Tty + ?Sized, H: Handler>(
     }
 }
 
-fn handle_parent<T: Tty + ?Sized, H: Handler>(
+fn handle_parent<T: Tty, H: Handler>(
     master_fd: OwnedFd,
     child: unistd::Pid,
     tty: &mut T,
     handler: &mut H,
     epoch: Instant,
-) -> Result<i32> {
+) -> anyhow::Result<i32> {
     let wait_result = match copy(master_fd, child, tty, handler, epoch) {
         Ok(Some(status)) => Ok(status),
         Ok(None) => wait::waitpid(child, None),
@@ -79,13 +86,13 @@ fn handle_parent<T: Tty + ?Sized, H: Handler>(
 
 const BUF_SIZE: usize = 128 * 1024;
 
-fn copy<T: Tty + ?Sized, H: Handler>(
+fn copy<T: Tty, H: Handler>(
     master_fd: OwnedFd,
     child: unistd::Pid,
     tty: &mut T,
     handler: &mut H,
     epoch: Instant,
-) -> Result<Option<WaitStatus>> {
+) -> anyhow::Result<Option<WaitStatus>> {
     let mut master = File::from(master_fd);
     let master_raw_fd = master.as_raw_fd();
     let mut buf = [0u8; BUF_SIZE];
@@ -274,7 +281,7 @@ fn copy<T: Tty + ?Sized, H: Handler>(
     }
 }
 
-fn handle_child<S: AsRef<str>>(command: &[S], extra_env: &ExtraEnv) -> Result<()> {
+fn handle_child<S: AsRef<str>>(command: &[S], extra_env: &ExtraEnv) -> anyhow::Result<()> {
     use signal::{SigHandler, Signal};
 
     let command = command
@@ -336,7 +343,7 @@ struct SignalFd {
 }
 
 impl SignalFd {
-    fn open(signal: libc::c_int) -> Result<Self> {
+    fn open(signal: libc::c_int) -> anyhow::Result<Self> {
         let (rx, tx) = unistd::pipe()?;
         set_non_blocking(&rx)?;
         set_non_blocking(&tx)?;
@@ -377,22 +384,29 @@ impl Drop for SignalFd {
 
 #[cfg(test)]
 mod tests {
-    use super::Handler;
+    use super::{Handler, HandlerStarter};
     use crate::pty::ExtraEnv;
     use crate::tty::{FixedSizeTty, NullTty, Theme, TtySize};
     use std::time::Duration;
 
+    struct TestHandlerStarter;
+
     #[derive(Default)]
     struct TestHandler {
-        tty_size: Option<TtySize>,
+        tty_size: TtySize,
         output: Vec<Vec<u8>>,
     }
 
-    impl Handler for TestHandler {
-        fn start(&mut self, tty_size: TtySize, _theme: Option<Theme>) {
-            self.tty_size = Some(tty_size);
+    impl HandlerStarter<TestHandler> for TestHandlerStarter {
+        fn start(self, tty_size: TtySize, _theme: Option<Theme>) -> TestHandler {
+            TestHandler {
+                tty_size,
+                output: Vec::new(),
+            }
         }
+    }
 
+    impl Handler for TestHandler {
         fn output(&mut self, _time: Duration, data: &[u8]) -> bool {
             self.output.push(data.into());
 
@@ -406,6 +420,10 @@ mod tests {
         fn resize(&mut self, _time: Duration, _size: TtySize) -> bool {
             true
         }
+
+        fn stop(self) -> Self {
+            self
+        }
     }
 
     impl TestHandler {
@@ -419,7 +437,7 @@ mod tests {
 
     #[test]
     fn exec_basic() {
-        let mut handler = TestHandler::default();
+        let starter = TestHandlerStarter;
 
         let code = r#"
 import sys;
@@ -430,27 +448,27 @@ time.sleep(0.1);
 sys.stdout.write('bar');
 "#;
 
-        super::exec(
+        let (_code, handler) = super::exec(
             &["python3", "-c", code],
             &ExtraEnv::new(),
             &mut NullTty::open().unwrap(),
-            &mut handler,
+            starter,
         )
         .unwrap();
 
         assert_eq!(handler.output(), vec!["foo", "bar"]);
-        assert_eq!(handler.tty_size, Some(TtySize(80, 24)));
+        assert_eq!(handler.tty_size, TtySize(80, 24));
     }
 
     #[test]
     fn exec_no_output() {
-        let mut handler = TestHandler::default();
+        let starter = TestHandlerStarter;
 
-        super::exec(
+        let (_code, handler) = super::exec(
             &["true"],
             &ExtraEnv::new(),
             &mut NullTty::open().unwrap(),
-            &mut handler,
+            starter,
         )
         .unwrap();
 
@@ -459,13 +477,13 @@ sys.stdout.write('bar');
 
     #[test]
     fn exec_quick() {
-        let mut handler = TestHandler::default();
+        let starter = TestHandlerStarter;
 
-        super::exec(
+        let (_code, handler) = super::exec(
             &["printf", "hello world\n"],
             &ExtraEnv::new(),
             &mut NullTty::open().unwrap(),
-            &mut handler,
+            starter,
         )
         .unwrap();
 
@@ -474,16 +492,16 @@ sys.stdout.write('bar');
 
     #[test]
     fn exec_extra_env() {
-        let mut handler = TestHandler::default();
+        let starter = TestHandlerStarter;
 
         let mut env = ExtraEnv::new();
         env.insert("ASCIINEMA_TEST_FOO".to_owned(), "bar".to_owned());
 
-        super::exec(
+        let (_code, handler) = super::exec(
             &["sh", "-c", "echo -n $ASCIINEMA_TEST_FOO"],
             &env,
             &mut NullTty::open().unwrap(),
-            &mut handler,
+            starter,
         )
         .unwrap();
 
@@ -492,16 +510,16 @@ sys.stdout.write('bar');
 
     #[test]
     fn exec_winsize_override() {
-        let mut handler = TestHandler::default();
+        let starter = TestHandlerStarter;
 
-        super::exec(
+        let (_code, handler) = super::exec(
             &["true"],
             &ExtraEnv::new(),
             &mut FixedSizeTty::new(NullTty::open().unwrap(), Some(100), Some(50)),
-            &mut handler,
+            starter,
         )
         .unwrap();
 
-        assert_eq!(handler.tty_size, Some(TtySize(100, 50)));
+        assert_eq!(handler.tty_size, TtySize(100, 50));
     }
 }

src/session.rs

@@ -1,5 +1,5 @@
 use std::io;
-use std::sync::mpsc::{self, Receiver};
+use std::sync::mpsc;
 use std::thread;
 use std::time::{Duration, SystemTime};
 
@@ -9,20 +9,11 @@ use crate::pty;
 use crate::tty;
 use crate::util::{JoinHandle, Utf8Decoder};
 
-pub struct Session<N> {
+pub struct SessionStarter<N> {
     outputs: Vec<Box<dyn Output + Send>>,
-    input_decoder: Utf8Decoder,
-    output_decoder: Utf8Decoder,
-    tty_size: tty::TtySize,
     record_input: bool,
     keys: KeyBindings,
     notifier: N,
-    sender: mpsc::Sender<Event>,
-    receiver: Option<Receiver<Event>>,
-    handle: Option<JoinHandle>,
-    time_offset: u64,
-    pause_time: Option<u64>,
-    prefix_mode: bool,
 }
 
 pub trait Output {
@@ -44,32 +35,71 @@ pub enum Event {
     Marker(u64, String),
 }
 
-impl<N: Notifier> Session<N> {
+impl<N: Notifier> SessionStarter<N> {
     pub fn new(
         outputs: Vec<Box<dyn Output + Send>>,
         record_input: bool,
         keys: KeyBindings,
         notifier: N,
     ) -> Self {
-        let (sender, receiver) = mpsc::channel();
-
-        Session {
+        SessionStarter {
             outputs,
-            input_decoder: Utf8Decoder::new(),
-            output_decoder: Utf8Decoder::new(),
-            tty_size: tty::TtySize::default(),
             record_input,
             keys,
             notifier,
+        }
+    }
+}
+
+impl<N: Notifier> pty::HandlerStarter<Session<N>> for SessionStarter<N> {
+    fn start(mut self, tty_size: tty::TtySize, tty_theme: Option<tty::Theme>) -> Session<N> {
+        let mut outputs = std::mem::take(&mut self.outputs);
+        let time = SystemTime::now();
+        let (sender, receiver) = mpsc::channel::<Event>();
+
+        outputs.retain_mut(|output| output.start(time, tty_size, tty_theme.clone()).is_ok());
+
+        let handle = thread::spawn(move || {
+            for event in receiver {
+                outputs.retain_mut(|output| output.event(event.clone()).is_ok())
+            }
+
+            for mut output in outputs {
+                let _ = output.flush();
+            }
+        });
+
+        Session {
+            notifier: self.notifier,
+            input_decoder: Utf8Decoder::new(),
+            output_decoder: Utf8Decoder::new(),
+            record_input: self.record_input,
+            keys: self.keys,
+            tty_size,
             sender,
-            receiver: Some(receiver),
-            handle: None,
             time_offset: 0,
             pause_time: None,
             prefix_mode: false,
+            _handle: JoinHandle::new(handle),
         }
     }
+}
 
+pub struct Session<N> {
+    notifier: N,
+    input_decoder: Utf8Decoder,
+    output_decoder: Utf8Decoder,
+    tty_size: tty::TtySize,
+    record_input: bool,
+    keys: KeyBindings,
+    sender: mpsc::Sender<Event>,
+    time_offset: u64,
+    pause_time: Option<u64>,
+    prefix_mode: bool,
+    _handle: JoinHandle,
+}
+
+impl<N: Notifier> Session<N> {
     fn elapsed_time(&self, time: Duration) -> u64 {
         if let Some(pause_time) = self.pause_time {
             pause_time
@@ -86,26 +116,6 @@ impl<N: Notifier> Session<N> {
 }
 
 impl<N: Notifier> pty::Handler for Session<N> {
-    fn start(&mut self, tty_size: tty::TtySize, tty_theme: Option<tty::Theme>) {
-        let mut outputs = std::mem::take(&mut self.outputs);
-        let time = SystemTime::now();
-        let receiver = self.receiver.take().unwrap();
-
-        let handle = thread::spawn(move || {
-            outputs.retain_mut(|output| output.start(time, tty_size, tty_theme.clone()).is_ok());
-
-            for event in receiver {
-                outputs.retain_mut(|output| output.event(event.clone()).is_ok())
-            }
-
-            for mut output in outputs {
-                let _ = output.flush();
-            }
-        });
-
-        self.handle = Some(JoinHandle::new(handle));
-    }
-
     fn output(&mut self, time: Duration, data: &[u8]) -> bool {
         if self.pause_time.is_none() {
             let text = self.output_decoder.feed(data);
@@ -173,6 +183,10 @@ impl<N: Notifier> pty::Handler for Session<N> {
 
         true
     }
+
+    fn stop(self) -> Self {
+        self
+    }
 }
 
 pub struct KeyBindings {