3.2.0 add cgroupPath option (#41)

index.d.ts

@@ -10,6 +10,7 @@ export interface PtyOptions {
   envs?: Record<string, string>
   dir?: string
   size?: Size
+  cgroupPath?: string
   interactive?: boolean
   onExit: (err: null | Error, exitCode: number) => void
 }

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@replit/ruspty",
-  "version": "3.1.3",
+  "version": "3.2.0",
   "main": "dist/wrapper.js",
   "types": "dist/wrapper.d.ts",
   "author": "Szymon Kaliski <[email protected]>",

src/lib.rs

@@ -1,6 +1,8 @@
 use std::collections::HashMap;
+use std::fs::File;
 use std::io::Error;
 use std::io::ErrorKind;
+use std::io::Write;
 use std::os::fd::{AsRawFd, OwnedFd};
 use std::os::fd::{BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 use std::os::unix::process::CommandExt;
@@ -40,6 +42,7 @@ struct PtyOptions {
   pub envs: Option<HashMap<String, String>>,
   pub dir: Option<String>,
   pub size: Option<Size>,
+  pub cgroup_path: Option<String>,
   pub interactive: Option<bool>,
   #[napi(ts_type = "(err: null | Error, exitCode: number) => void")]
   pub on_exit: JsFunction,
@@ -105,6 +108,14 @@ impl Pty {
   #[napi(constructor)]
   #[allow(dead_code)]
   pub fn new(_env: Env, opts: PtyOptions) -> Result<Self, napi::Error> {
+    #[cfg(not(target_os = "linux"))]
+    if opts.cgroup_path.is_some() {
+      return Err(napi::Error::new(
+        napi::Status::GenericFailure,
+        "cgroup_path is only supported on Linux",
+      ));
+    }
+
     let size = opts.size.unwrap_or(Size { cols: 80, rows: 24 });
     let window_size = Winsize {
       ws_col: size.cols,
@@ -160,6 +171,15 @@ impl Pty {
           return Err(Error::new(ErrorKind::Other, "setsid"));
         }
 
+        // set the cgroup if specified
+        #[cfg(target_os = "linux")]
+        if let Some(cgroup_path) = &opts.cgroup_path {
+          let pid = libc::getpid();
+          let cgroup_path = format!("{}/cgroup.procs", cgroup_path);
+          let mut cgroup_file = File::create(cgroup_path)?;
+          cgroup_file.write_all(format!("{}", pid).as_bytes())?;
+        }
+
         // become the controlling tty for the program
         let err = libc::ioctl(raw_user_fd, libc::TIOCSCTTY.into(), 0);
         if err == -1 {

tests/index.test.ts

@@ -1,13 +1,17 @@
 import { Pty, getCloseOnExec, setCloseOnExec } from '../wrapper';
 import { type Writable } from 'stream';
 import { readdirSync, readlinkSync } from 'fs';
-import { describe, test, expect } from 'vitest';
+import { describe, test, expect, beforeEach, afterEach } from 'vitest';
+import { exec as execAsync } from 'child_process';
+import { promisify } from 'util';
+const exec = promisify(execAsync);
 
 const EOT = '\x04';
 const procSelfFd = '/proc/self/fd/';
 const IS_DARWIN = process.platform === 'darwin';
 
 const testSkipOnDarwin = IS_DARWIN ? test.skip : test;
+const testOnlyOnDarwin = IS_DARWIN ? test : test.skip;
 
 type FdRecord = Record<string, string>;
 function getOpenFds(): FdRecord {
@@ -414,6 +418,60 @@ describe(
   { repeats: 50 },
 );
 
+describe('cgroup opts', () => {
+  beforeEach(async () => {
+    if (!IS_DARWIN) {
+      // create a new cgroup with the right permissions
+      await exec("sudo cgcreate -g 'cpu:/test.slice'")
+      await exec("sudo chown -R $(id -u):$(id -g) /sys/fs/cgroup/cpu/test.slice")
+    }
+  });
+
+  afterEach(async () => {
+    if (!IS_DARWIN) {
+      // remove the cgroup
+      await exec("sudo cgdelete cpu:/test.slice")
+    }
+  });
+
+  testSkipOnDarwin('basic cgroup', () => new Promise<void>((done) => {
+    const oldFds = getOpenFds();
+    let buffer = '';
+    const pty = new Pty({
+      command: '/bin/cat',
+      args: ['/proc/self/cgroup'],
+      cgroupPath: '/sys/fs/cgroup/cpu/test.slice',
+      onExit: (err, exitCode) => {
+        expect(err).toBeNull();
+        expect(exitCode).toBe(0);
+        expect(buffer).toContain('/test.slice');
+        expect(getOpenFds()).toStrictEqual(oldFds);
+        done();
+      },
+    });
+
+    const readStream = pty.read;
+    readStream.on('data', (data) => {
+      buffer = data.toString();
+    });
+  })
+  );
+
+  testOnlyOnDarwin('cgroup is not supported on darwin', () => {
+    expect(() => {
+      new Pty({
+        command: '/bin/cat',
+        args: ['/proc/self/cgroup'],
+        cgroupPath: '/sys/fs/cgroup/cpu/test.slice',
+        onExit: (err, exitCode) => {
+          expect(err).toBeNull();
+          expect(exitCode).toBe(0);
+        },
+      })
+    }).toThrowError();
+  });
+});
+
 describe('setCloseOnExec', () => {
   test('setCloseOnExec', () => {
     // stdio typically never has the close-on-exec flag since it's always expected to be