Make pty::exec measure time

src/pty.rs

@@ -15,15 +15,16 @@ use std::io::{self, ErrorKind, Read, Write};
 use std::os::fd::{AsFd, RawFd};
 use std::os::fd::{BorrowedFd, OwnedFd};
 use std::os::unix::io::{AsRawFd, FromRawFd};
+use std::time::{Duration, Instant};
 use std::{env, fs};
 
 type ExtraEnv = HashMap<String, String>;
 
 pub trait Handler {
-    fn start(&mut self, size: TtySize);
-    fn output(&mut self, data: &[u8]) -> bool;
-    fn input(&mut self, data: &[u8]) -> bool;
-    fn resize(&mut self, size: TtySize) -> bool;
+    fn start(&mut self, epoch: Instant, tty_size: TtySize);
+    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;
 }
 
 pub fn exec<S: AsRef<str>, T: Tty + ?Sized, H: Handler>(
@@ -33,12 +34,13 @@ pub fn exec<S: AsRef<str>, T: Tty + ?Sized, H: Handler>(
     handler: &mut H,
 ) -> Result<i32> {
     let winsize = tty.get_size();
-    handler.start(winsize.into());
+    let epoch = Instant::now();
+    handler.start(epoch, winsize.into());
     let result = unsafe { pty::forkpty(Some(&winsize), None) }?;
 
     match result.fork_result {
         ForkResult::Parent { child } => {
-            handle_parent(result.master.as_raw_fd(), child, tty, handler)
+            handle_parent(result.master.as_raw_fd(), child, tty, handler, epoch)
         }
 
         ForkResult::Child => {
@@ -53,8 +55,9 @@ fn handle_parent<T: Tty + ?Sized, H: Handler>(
     child: unistd::Pid,
     tty: &mut T,
     handler: &mut H,
+    epoch: Instant,
 ) -> Result<i32> {
-    let wait_result = match copy(master_fd, child, tty, handler) {
+    let wait_result = match copy(master_fd, child, tty, handler, epoch) {
         Ok(Some(status)) => Ok(status),
         Ok(None) => wait::waitpid(child, None),
 
@@ -79,6 +82,7 @@ fn copy<T: Tty + ?Sized, H: Handler>(
     child: unistd::Pid,
     tty: &mut T,
     handler: &mut H,
+    epoch: Instant,
 ) -> Result<Option<WaitStatus>> {
     let mut master = unsafe { fs::File::from_raw_fd(master_raw_fd) };
     let mut buf = [0u8; BUF_SIZE];
@@ -146,7 +150,7 @@ fn copy<T: Tty + ?Sized, H: Handler>(
         if master_read {
             while let Some(n) = read_non_blocking(&mut master, &mut buf)? {
                 if n > 0 {
-                    if handler.output(&buf[0..n]) {
+                    if handler.output(epoch.elapsed(), &buf[0..n]) {
                         output.extend_from_slice(&buf[0..n]);
                     }
                 } else if output.is_empty() {
@@ -205,7 +209,7 @@ fn copy<T: Tty + ?Sized, H: Handler>(
         if tty_read {
             while let Some(n) = read_non_blocking(tty, &mut buf)? {
                 if n > 0 {
-                    if handler.input(&buf[0..n]) {
+                    if handler.input(epoch.elapsed(), &buf[0..n]) {
                         input.extend_from_slice(&buf[0..n]);
                     }
                 } else {
@@ -218,7 +222,7 @@ fn copy<T: Tty + ?Sized, H: Handler>(
             sigwinch_fd.flush();
             let winsize = tty.get_size();
 
-            if handler.resize(winsize.into()) {
+            if handler.resize(epoch.elapsed(), winsize.into()) {
                 set_pty_size(master_raw_fd, &winsize);
             }
         }
@@ -372,6 +376,7 @@ mod tests {
     use super::Handler;
     use crate::pty::ExtraEnv;
     use crate::tty::{FixedSizeTty, NullTty, TtySize};
+    use std::time::{Duration, Instant};
 
     #[derive(Default)]
     struct TestHandler {
@@ -380,21 +385,21 @@ mod tests {
     }
 
     impl Handler for TestHandler {
-        fn start(&mut self, tty_size: TtySize) {
+        fn start(&mut self, _epoch: Instant, tty_size: TtySize) {
             self.tty_size = Some(tty_size);
         }
 
-        fn output(&mut self, data: &[u8]) -> bool {
+        fn output(&mut self, _time: Duration, data: &[u8]) -> bool {
             self.output.push(data.into());
 
             true
         }
 
-        fn input(&mut self, _data: &[u8]) -> bool {
+        fn input(&mut self, _time: Duration, _data: &[u8]) -> bool {
             true
         }
 
-        fn resize(&mut self, _size: TtySize) -> bool {
+        fn resize(&mut self, _time: Duration, _size: TtySize) -> bool {
             true
         }
     }

src/recorder.rs

@@ -16,7 +16,7 @@ pub struct Recorder {
     sender: mpsc::Sender<Message>,
     receiver: Option<mpsc::Receiver<Message>>,
     handle: Option<util::JoinHandle>,
-    start_time: Instant,
+    time_offset: u64,
     pause_time: Option<u64>,
     prefix_mode: bool,
 }
@@ -58,17 +58,17 @@ impl Recorder {
             sender,
             receiver: Some(receiver),
             handle: None,
-            start_time: Instant::now(),
+            time_offset: 0,
             pause_time: None,
             prefix_mode: false,
         }
     }
 
-    fn elapsed_time(&self) -> u64 {
+    fn elapsed_time(&self, time: Duration) -> u64 {
         if let Some(pause_time) = self.pause_time {
             pause_time
         } else {
-            self.start_time.elapsed().as_micros() as u64
+            time.as_micros() as u64 - self.time_offset
         }
     }
 
@@ -82,7 +82,7 @@ impl Recorder {
 }
 
 impl pty::Handler for Recorder {
-    fn start(&mut self, tty_size: tty::TtySize) {
+    fn start(&mut self, _epoch: Instant, tty_size: tty::TtySize) {
         let mut output = self.output.take().unwrap();
         let _ = output.start(&tty_size);
         let receiver = self.receiver.take().unwrap();
@@ -133,19 +133,18 @@ impl pty::Handler for Recorder {
         });
 
         self.handle = Some(util::JoinHandle::new(handle));
-        self.start_time = Instant::now();
     }
 
-    fn output(&mut self, data: &[u8]) -> bool {
+    fn output(&mut self, time: Duration, data: &[u8]) -> bool {
         if self.pause_time.is_none() {
-            let msg = Message::Output(self.elapsed_time(), data.into());
+            let msg = Message::Output(self.elapsed_time(time), data.into());
             self.sender.send(msg).expect("output send should succeed");
         }
 
         true
     }
 
-    fn input(&mut self, data: &[u8]) -> bool {
+    fn input(&mut self, time: Duration, data: &[u8]) -> bool {
         let prefix_key = self.keys.prefix.as_ref();
         let pause_key = self.keys.pause.as_ref();
         let add_marker_key = self.keys.add_marker.as_ref();
@@ -160,17 +159,17 @@ impl pty::Handler for Recorder {
 
             if pause_key.is_some_and(|key| data == key) {
                 if let Some(pt) = self.pause_time {
-                    self.start_time = Instant::now() - Duration::from_micros(pt);
                     self.pause_time = None;
+                    self.time_offset += self.elapsed_time(time) - pt;
                     self.notify("Resumed recording");
                 } else {
-                    self.pause_time = Some(self.elapsed_time());
+                    self.pause_time = Some(self.elapsed_time(time));
                     self.notify("Paused recording");
                 }
 
                 return false;
             } else if add_marker_key.is_some_and(|key| data == key) {
-                let msg = Message::Marker(self.elapsed_time());
+                let msg = Message::Marker(self.elapsed_time(time));
                 self.sender.send(msg).expect("marker send should succeed");
                 self.notify("Marker added");
                 return false;
@@ -178,15 +177,15 @@ impl pty::Handler for Recorder {
         }
 
         if self.record_input && self.pause_time.is_none() {
-            let msg = Message::Input(self.elapsed_time(), data.into());
+            let msg = Message::Input(self.elapsed_time(time), data.into());
             self.sender.send(msg).expect("input send should succeed");
         }
 
         true
     }
 
-    fn resize(&mut self, size: tty::TtySize) -> bool {
-        let msg = Message::Resize(self.elapsed_time(), size);
+    fn resize(&mut self, time: Duration, tty_size: tty::TtySize) -> bool {
+        let msg = Message::Resize(self.elapsed_time(time), tty_size);
         self.sender.send(msg).expect("resize send should succeed");
 
         true

src/streamer/mod.rs

@@ -20,7 +20,6 @@ pub struct Streamer {
     notifier: Option<Box<dyn Notifier>>,
     notifier_rx: Option<std::sync::mpsc::Receiver<String>>,
     pty_rx: Option<mpsc::UnboundedReceiver<Event>>,
-    start_time: Instant,
     paused: bool,
     prefix_mode: bool,
     listener: Option<net::TcpListener>,
@@ -61,7 +60,6 @@ impl Streamer {
             pty_tx,
             pty_rx: Some(pty_rx),
             event_loop_handle: None,
-            start_time: Instant::now(),
             paused: false,
             prefix_mode: false,
             listener,
@@ -70,8 +68,8 @@ impl Streamer {
         }
     }
 
-    fn elapsed_time(&self) -> u64 {
-        self.start_time.elapsed().as_micros() as u64
+    fn elapsed_time(&self, time: Duration) -> u64 {
+        time.as_micros() as u64
     }
 
     fn notify<S: ToString>(&self, message: S) {
@@ -85,7 +83,7 @@ impl Streamer {
 }
 
 impl pty::Handler for Streamer {
-    fn start(&mut self, tty_size: tty::TtySize) {
+    fn start(&mut self, _epoch: Instant, tty_size: tty::TtySize) {
         let pty_rx = self.pty_rx.take().unwrap();
         let (clients_tx, mut clients_rx) = mpsc::channel(1);
         let shutdown_token = tokio_util::sync::CancellationToken::new();
@@ -136,20 +134,18 @@ impl pty::Handler for Streamer {
                 let _ = notifier.notify(message);
             }
         }));
-
-        self.start_time = Instant::now();
     }
 
-    fn output(&mut self, raw: &[u8]) -> bool {
+    fn output(&mut self, time: Duration, raw: &[u8]) -> bool {
         if !self.paused {
-            let event = Event::Output(self.elapsed_time(), raw.into());
+            let event = Event::Output(self.elapsed_time(time), raw.into());
             let _ = self.pty_tx.send(event);
         }
 
         true
     }
 
-    fn input(&mut self, raw: &[u8]) -> bool {
+    fn input(&mut self, time: Duration, raw: &[u8]) -> bool {
         let prefix_key = self.keys.prefix.as_ref();
         let pause_key = self.keys.pause.as_ref();
 
@@ -175,15 +171,15 @@ impl pty::Handler for Streamer {
         }
 
         if self.record_input && !self.paused {
-            let event = Event::Input(self.elapsed_time(), raw.into());
+            let event = Event::Input(self.elapsed_time(time), raw.into());
             let _ = self.pty_tx.send(event);
         }
 
         true
     }
 
-    fn resize(&mut self, size: crate::tty::TtySize) -> bool {
-        let event = Event::Resize(self.elapsed_time(), size);
+    fn resize(&mut self, time: Duration, tty_size: tty::TtySize) -> bool {
+        let event = Event::Resize(self.elapsed_time(time), tty_size);
         let _ = self.pty_tx.send(event);
 
         true