try keeping user fd alive longer

Author: jackyzha0

index.d.ts

@@ -67,4 +67,5 @@ export declare class Pty {
    * descriptor.
    */
   takeFd(): c_int;
+  closeUserFd(): void;
 }

src/lib.rs

@@ -17,7 +17,7 @@ use napi::Status::GenericFailure;
 use napi::{self, Env};
 use nix::errno::Errno;
 use nix::fcntl::{fcntl, FcntlArg, FdFlag, OFlag};
-use nix::libc::{self, c_int, ioctl, FIONREAD, TIOCOUTQ, TIOCSCTTY, TIOCSWINSZ};
+use nix::libc::{self, c_int, ioctl, FIONREAD, TIOCSCTTY, TIOCSWINSZ};
 use nix::pty::{openpty, Winsize};
 use nix::sys::termios::{self, SetArg};
 
@@ -31,6 +31,7 @@ mod sandbox;
 #[allow(dead_code)]
 struct Pty {
   controller_fd: Option<OwnedFd>,
+  user_fd: Option<OwnedFd>,
   /// The pid of the forked process.
   pub pid: u32,
 }
@@ -107,9 +108,7 @@ fn poll_pty_fds_until_read(controller_fd: RawFd, user_fd: RawFd) {
   loop {
     // check both input and output queues for both FDs
     let mut controller_inq: i32 = 0;
-    let mut controller_outq: i32 = 0;
     let mut user_inq: i32 = 0;
-    let mut user_outq: i32 = 0;
 
     // safe because we're passing valid file descriptors and properly sized integers
     unsafe {
@@ -120,18 +119,10 @@ fn poll_pty_fds_until_read(controller_fd: RawFd, user_fd: RawFd) {
         // break if we can't read
         break;
       }
-
-      // check bytes waiting to be written (TIOCOUTQ)
-      if ioctl(controller_fd, TIOCOUTQ, &mut controller_outq) == -1
-        || ioctl(user_fd, TIOCOUTQ, &mut user_outq) == -1
-      {
-        // break if we can't read
-        break;
-      }
     }
 
     // if all queues are empty, we're done
-    if controller_inq == 0 && controller_outq == 0 && user_inq == 0 && user_outq == 0 {
+    if controller_inq == 0 && user_inq == 0 {
       break;
     }
 
@@ -349,7 +340,6 @@ impl Pty {
 
       // try to wait for the controller fd to be fully read
       poll_pty_fds_until_read(raw_controller_fd, raw_user_fd);
-      drop(user_fd);
 
       match wait_result {
         Ok(status) => {
@@ -379,6 +369,7 @@ impl Pty {
 
     Ok(Pty {
       controller_fd: Some(controller_fd),
+      user_fd: Some(user_fd),
       pid,
     })
   }
@@ -398,6 +389,13 @@ impl Pty {
       ))
     }
   }
+
+  #[napi]
+  #[allow(dead_code)]
+  pub fn close_user_fd(&mut self) -> Result<(), napi::Error> {
+    self.user_fd.take();
+    Ok(())
+  }
 }
 
 /// Resize the terminal.

tests/index.test.ts

@@ -354,6 +354,27 @@ describe('PTY', { repeats: 500 }, () => {
     expect(buffer.toString().length).toBe(payload.length);
   });
 
+  test.only('doesnt miss lots of lines from bash', async () => {
+    const payload = Array.from({ length: 4096 * 5 }, (_, i) => i).join('\n');
+    let buffer = Buffer.from('');
+    const onExit = vi.fn();
+
+    const pty = new Pty({
+      command: 'bash',
+      args: ['-c', `echo -n "${payload}"`],
+      onExit,
+    });
+
+    const readStream = pty.read;
+    readStream.on('data', (data) => {
+      buffer = Buffer.concat([buffer, data]);
+    });
+
+    await vi.waitFor(() => expect(onExit).toHaveBeenCalledTimes(1));
+    expect(onExit).toHaveBeenCalledWith(null, 0);
+    expect(buffer.toString().trim().replace(/\r/g, '')).toBe(payload.trim());
+  });
+
   testSkipOnDarwin('does not leak files', async () => {
     const oldFds = getOpenFds();
     const promises = [];

wrapper.ts

@@ -78,7 +78,12 @@ export class Pty {
       markReadFinished = resolve;
     });
     const mockedExit = (error: NodeJS.ErrnoException | null, code: number) => {
+      console.log('mockedExit', { error, code });
       markExited({ error, code });
+
+      setImmediate(() => {
+        this.#pty.closeUserFd();
+      });
     };
 
     // when pty exits, we should wait until the fd actually ends (end OR error)
@@ -105,8 +110,14 @@ export class Pty {
       realExit(result.error, result.code);
     };
 
-    this.read.once('end', markReadFinished);
-    this.read.once('close', handleClose);
+    this.read.once('end', () => {
+      console.log('read end');
+      markReadFinished();
+    });
+    this.read.once('close', () => {
+      console.log('read close');
+      handleClose();
+    });
 
     // PTYs signal their done-ness with an EIO error. we therefore need to filter them out (as well as
     // cleaning up other spurious errors) so that the user doesn't need to handle them and be in
@@ -135,7 +146,10 @@ export class Pty {
       throw err;
     };
 
-    this.read.on('error', handleError);
+    this.read.on('error', (err) => {
+      console.error('error:', err);
+      handleError(err);
+    });
   }
 
   close() {