Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

3.2.0 add cgroupPath option #41

Merged
merged 14 commits into from
Aug 26, 2024
1 change: 1 addition & 0 deletions index.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
4 changes: 2 additions & 2 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -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]>",
Expand Down
20 changes: 20 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 {
Expand Down
60 changes: 59 additions & 1 deletion tests/index.test.ts
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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
Expand Down