use super::arg_str;
use crate::host::{with_host, IoTask, JsObj};
use fusevm::Value;
use indexmap::IndexMap;
use std::cell::RefCell;
use std::collections::HashMap;
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicU64, Ordering};
pub const METHODS: &[&str] = &[
"fork",
"setupPrimary",
"setupMaster",
"disconnect",
"on",
"addListener",
"prependListener",
"once",
"prependOnceListener",
"emit",
"removeListener",
"off",
"removeAllListeners",
"listenerCount",
"listeners",
"eventNames",
"setMaxListeners",
"getMaxListeners",
];
pub const WORKER_METHODS: &[&str] = &[
"send",
"kill",
"destroy",
"disconnect",
"isConnected",
"isDead",
];
const EMITTER_METHODS: &[&str] = &[
"on",
"addListener",
"prependListener",
"once",
"prependOnceListener",
"emit",
"removeListener",
"off",
"removeAllListeners",
"listenerCount",
"listeners",
"eventNames",
"setMaxListeners",
"getMaxListeners",
];
static NEXT_WORKER_ID: AtomicU64 = AtomicU64::new(1);
#[derive(Default, Clone)]
struct Settings {
exec: Option<String>,
args: Option<Vec<String>>,
exec_argv: Option<Vec<String>>,
silent: bool,
}
thread_local! {
static WORKERS: RefCell<HashMap<u64, Value>> = RefCell::new(HashMap::new());
static CLUSTER_EMITTER: RefCell<Option<Value>> = const { RefCell::new(None) };
static SELF_WORKER: RefCell<Option<Value>> = const { RefCell::new(None) };
static SETTINGS: RefCell<Settings> = RefCell::new(Settings::default());
}
fn worker_id_from_env() -> Option<u64> {
std::env::var("CLUSTER_WORKER")
.ok()
.or_else(|| std::env::var("NODE_UNIQUE_ID").ok())
.and_then(|s| s.trim().parse::<u64>().ok())
}
fn is_primary() -> bool {
worker_id_from_env().is_none()
}
pub fn call(method: &str, args: &[Value]) -> Option<Result<Value, String>> {
if EMITTER_METHODS.contains(&method) {
let em = cluster_emitter();
return Some(super::events::instance_call(&em, method, args.to_vec()));
}
Some(match method {
"fork" => fork(args),
"setupPrimary" | "setupMaster" => setup_primary(args),
"disconnect" => disconnect(args),
_ => return None,
})
}
pub fn constant(name: &str) -> Option<Value> {
Some(match name {
"isPrimary" | "isMaster" => Value::Bool(is_primary()),
"isWorker" => Value::Bool(!is_primary()),
"workers" => workers_object(),
"worker" => {
if is_primary() {
with_host(|h| h.null())
} else {
self_worker()
}
}
"settings" => settings_object(),
"SCHED_NONE" => Value::Float(1.0),
"SCHED_RR" => Value::Float(2.0),
"schedulingPolicy" => Value::Float(1.0),
_ => return None,
})
}
fn fork(args: &[Value]) -> Result<Value, String> {
if !is_primary() {
return Err("Error: cluster.fork() can only be called from the primary process".into());
}
let s = SETTINGS.with(|s| s.borrow().clone());
let exec = s
.exec
.clone()
.or_else(|| std::env::args().nth(1))
.unwrap_or_default();
if exec.is_empty() {
return Err(
"Error: cluster.fork() requires a main script (process.argv[1]); none was found".into(),
);
}
let fwd_args: Vec<String> = s
.args
.clone()
.unwrap_or_else(|| std::env::args().skip(2).collect());
let exe = std::env::current_exe().map_err(|e| format!("Error: cluster.fork(): {e}"))?;
let overrides = args.first().map(env_overrides).unwrap_or_default();
let id = NEXT_WORKER_ID.fetch_add(1, Ordering::SeqCst);
let mut cmd = Command::new(exe);
cmd.arg(&exec);
cmd.args(&fwd_args);
cmd.env("CLUSTER_WORKER", id.to_string());
cmd.env("NODE_UNIQUE_ID", id.to_string());
for (k, v) in overrides {
cmd.env(k, v);
}
if s.silent {
cmd.stdout(Stdio::null()).stderr(Stdio::null());
} else {
cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
}
let child = cmd
.spawn()
.map_err(|e| format!("Error: cluster.fork(): {e}"))?;
let pid = child.id();
let worker = new_worker(id, pid);
WORKERS.with(|w| {
w.borrow_mut().insert(id, worker.clone());
});
with_host(|h| h.incr_handle());
let _ = emit_on(&cluster_emitter(), "fork", vec![worker.clone()]);
let io_online = with_host(|h| h.io_sender());
let _ = io_online.send(Box::new(move || dispatch_online(id)));
let io_exit: std::sync::mpsc::Sender<IoTask> = with_host(|h| h.io_sender());
std::thread::spawn(move || {
let mut child = child;
let code = child.wait().ok().and_then(|st| st.code()).unwrap_or(0);
let _ = io_exit.send(Box::new(move || dispatch_exit(id, code)));
});
Ok(worker)
}
fn env_overrides(v: &Value) -> Vec<(String, String)> {
with_host(|h| match h.get(v) {
Some(JsObj::Object(p)) => p
.iter()
.filter(|(k, _)| !k.starts_with("@@"))
.map(|(k, val)| (k.clone(), h.str_of(val)))
.collect(),
_ => Vec::new(),
})
}
fn setup_primary(args: &[Value]) -> Result<Value, String> {
if let Some(opts) = args.first() {
let exec = str_prop(opts, "exec");
let arr = arr_prop(opts, "args");
let ea = arr_prop(opts, "execArgv");
let silent = bool_prop(opts, "silent");
SETTINGS.with(|s| {
let mut s = s.borrow_mut();
if exec.is_some() {
s.exec = exec;
}
if arr.is_some() {
s.args = arr;
}
if ea.is_some() {
s.exec_argv = ea;
}
if let Some(b) = silent {
s.silent = b;
}
});
}
Ok(Value::Undef)
}
fn settings_object() -> Value {
let s = SETTINGS.with(|s| s.borrow().clone());
let exec = s
.exec
.clone()
.unwrap_or_else(|| std::env::args().nth(1).unwrap_or_default());
let args_vec = s
.args
.clone()
.unwrap_or_else(|| std::env::args().skip(2).collect());
let exec_argv = s.exec_argv.clone().unwrap_or_default();
with_host(|h| {
let arg_items: Vec<Value> = args_vec.into_iter().map(|a| h.new_str(a)).collect();
let args_arr = h.new_array(arg_items);
let ea_items: Vec<Value> = exec_argv.into_iter().map(|a| h.new_str(a)).collect();
let ea_arr = h.new_array(ea_items);
let exec_v = h.new_str(exec);
let mut m = IndexMap::new();
m.insert("exec".into(), exec_v);
m.insert("args".into(), args_arr);
m.insert("execArgv".into(), ea_arr);
m.insert("silent".into(), Value::Bool(s.silent));
h.new_object(m)
})
}
fn disconnect(args: &[Value]) -> Result<Value, String> {
let workers: Vec<Value> = WORKERS.with(|w| w.borrow().values().cloned().collect());
for wk in workers {
mark_disconnected(&wk);
}
if let Some(cb) = args.first() {
if with_host(|h| h.type_of(cb)) == "function" {
crate::host::invoke(cb, vec![], None)?;
}
}
Ok(Value::Undef)
}
fn new_worker(id: u64, pid: u32) -> Value {
let proc_obj = with_host(|h| {
let mut p = IndexMap::new();
p.insert("pid".into(), Value::Float(pid as f64));
p.insert("connected".into(), Value::Bool(true));
h.new_object(p)
});
let mut extra = IndexMap::new();
extra.insert("id".into(), Value::Float(id as f64));
extra.insert("process".into(), proc_obj);
extra.insert("@@cwid".into(), Value::Float(id as f64));
extra.insert("@@pid".into(), Value::Float(pid as f64));
extra.insert("@@connected".into(), Value::Bool(true));
super::net::new_emitter_object("ClusterWorker", extra)
}
fn self_worker() -> Value {
if let Some(v) = SELF_WORKER.with(|c| c.borrow().clone()) {
return v;
}
let id = worker_id_from_env().unwrap_or(0);
let w = new_worker(id, std::process::id());
SELF_WORKER.with(|c| *c.borrow_mut() = Some(w.clone()));
w
}
fn workers_object() -> Value {
let entries: Vec<(String, Value)> = WORKERS.with(|w| {
w.borrow()
.iter()
.map(|(id, wk)| (id.to_string(), wk.clone()))
.collect()
});
with_host(|h| {
let mut m = IndexMap::new();
for (k, v) in entries {
m.insert(k, v);
}
h.new_object(m)
})
}
pub fn instance_call(recv: &Value, method: &str, args: Vec<Value>) -> Result<Value, String> {
if EMITTER_METHODS.contains(&method) {
return super::events::instance_call(recv, method, args);
}
match method {
"send" => Ok(Value::Bool(false)),
"kill" | "destroy" => {
let sig = signal_number(args.first());
if let Some(pid) = pid_of(recv) {
unsafe {
libc::kill(pid as libc::pid_t, sig);
}
}
mark_disconnected(recv);
Ok(Value::Undef)
}
"disconnect" => {
mark_disconnected(recv);
Ok(recv.clone())
}
"isConnected" => Ok(Value::Bool(bool_prop(recv, "@@connected").unwrap_or(false))),
"isDead" => {
let id = pid_or_id(recv, "@@cwid");
let alive = id
.map(|i| WORKERS.with(|w| w.borrow().contains_key(&i)))
.unwrap_or(false);
Ok(Value::Bool(!alive))
}
_ => Err(crate::host::type_error(&format!(
"worker.{method} is not a function"
))),
}
}
fn mark_disconnected(worker: &Value) {
with_host(|h| {
if let Some(JsObj::Object(p)) = h.get_mut(worker) {
p.insert("@@connected".into(), Value::Bool(false));
}
});
let proc = with_host(|h| match h.get(worker) {
Some(JsObj::Object(p)) => p.get("process").cloned(),
_ => None,
});
if let Some(proc) = proc {
with_host(|h| {
if let Some(JsObj::Object(p)) = h.get_mut(&proc) {
p.insert("connected".into(), Value::Bool(false));
}
});
}
let _ = emit_on(worker, "disconnect", vec![]);
let _ = emit_on(&cluster_emitter(), "disconnect", vec![worker.clone()]);
}
fn dispatch_online(id: u64) -> Result<(), String> {
let Some(worker) = WORKERS.with(|w| w.borrow().get(&id).cloned()) else {
return Ok(());
};
emit_on(&worker, "online", vec![])?;
emit_on(&cluster_emitter(), "online", vec![worker])
}
fn dispatch_exit(id: u64, code: i32) -> Result<(), String> {
let Some(worker) = WORKERS.with(|w| w.borrow().get(&id).cloned()) else {
return Ok(());
};
let null_sig = with_host(|h| h.null());
emit_on(
&worker,
"exit",
vec![Value::Float(code as f64), null_sig.clone()],
)?;
emit_on(
&cluster_emitter(),
"exit",
vec![worker, Value::Float(code as f64), null_sig],
)?;
WORKERS.with(|w| {
w.borrow_mut().remove(&id);
});
with_host(|h| h.decr_handle());
Ok(())
}
fn cluster_emitter() -> Value {
if let Some(v) = CLUSTER_EMITTER.with(|c| c.borrow().clone()) {
return v;
}
let e = super::events::new_emitter();
CLUSTER_EMITTER.with(|c| *c.borrow_mut() = Some(e.clone()));
e
}
fn emit_on(emitter: &Value, name: &str, mut args: Vec<Value>) -> Result<(), String> {
let mut a = vec![with_host(|h| h.new_str(name))];
a.append(&mut args);
super::events::instance_call(emitter, "emit", a).map(|_| ())
}
fn pid_of(worker: &Value) -> Option<u32> {
pid_or_id(worker, "@@pid").map(|n| n as u32)
}
fn pid_or_id(worker: &Value, key: &str) -> Option<u64> {
with_host(|h| match h.get(worker) {
Some(JsObj::Object(p)) => p.get(key).map(|v| h.to_number(v) as u64),
_ => None,
})
}
fn signal_number(arg: Option<&Value>) -> libc::c_int {
let Some(v) = arg else { return libc::SIGTERM };
let n = with_host(|h| h.to_number(v));
if n.is_finite() && n != 0.0 {
return n as libc::c_int;
}
match arg_str(std::slice::from_ref(v), 0).to_uppercase().as_str() {
"SIGKILL" => libc::SIGKILL,
"SIGINT" => libc::SIGINT,
"SIGHUP" => libc::SIGHUP,
"SIGQUIT" => libc::SIGQUIT,
"SIGUSR1" => libc::SIGUSR1,
"SIGUSR2" => libc::SIGUSR2,
_ => libc::SIGTERM,
}
}
fn str_prop(obj: &Value, key: &str) -> Option<String> {
with_host(|h| match h.get(obj) {
Some(JsObj::Object(p)) => p
.get(key)
.map(|v| h.str_of(v))
.filter(|s| !s.is_empty() && s != "undefined"),
_ => None,
})
}
fn bool_prop(obj: &Value, key: &str) -> Option<bool> {
with_host(|h| match h.get(obj) {
Some(JsObj::Object(p)) => p.get(key).map(|v| h.truthy(v)),
_ => None,
})
}
fn arr_prop(obj: &Value, key: &str) -> Option<Vec<String>> {
with_host(|h| match h.get(obj) {
Some(JsObj::Object(p)) => match p.get(key).and_then(|v| h.get(v)) {
Some(JsObj::Array(items)) => Some(items.iter().map(|v| h.str_of(v)).collect()),
_ => None,
},
_ => None,
})
}