use crate::lockfree::OnceSlot;
use std::ffi::OsStr;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::process::{Command, ExitStatus, Output, Stdio};
use std::sync::{Arc, OnceLock, mpsc};
use std::task::{Context, Poll};
type Job = Box<dyn FnOnce() + Send + 'static>;
struct ProcessPool {
sender: mpsc::Sender<Job>,
}
static PROCESS_POOL: OnceLock<ProcessPool> = OnceLock::new();
pub fn init(workers: usize) {
PROCESS_POOL.get_or_init(|| {
let (tx, rx) = mpsc::channel::<Job>();
let rx = Arc::new(std::sync::Mutex::new(rx));
for _ in 0..workers.max(1) {
let rx = Arc::clone(&rx);
std::thread::Builder::new()
.name("dtact-process-worker".into())
.spawn(move || {
loop {
let job = { rx.lock().unwrap().recv() };
match job {
Ok(job) => job(),
Err(_) => break,
}
}
})
.expect("failed to spawn dtact-process worker thread");
}
ProcessPool { sender: tx }
});
}
pub fn init_process(
workers: usize,
_ring_depth: u32,
_buffer_pool_size: usize,
_chunk_size: usize,
_pin_cpus: &[usize],
) {
init(workers);
}
pub struct BlockingOp<T> {
slot: Arc<OnceSlot<T>>,
}
impl<T: Send + 'static> Future for BlockingOp<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
self.slot.poll(cx)
}
}
fn spawn_blocking<T, F>(f: F) -> BlockingOp<T>
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
if PROCESS_POOL.get().is_none() {
init(4);
}
let slot = Arc::new(OnceSlot::new());
let slot2 = Arc::clone(&slot);
let job: Job = Box::new(move || {
let result = f();
slot2.set(result);
});
let _ = PROCESS_POOL.get().unwrap().sender.send(job);
BlockingOp { slot }
}
pub struct DtactCommand(Command);
impl DtactCommand {
pub fn new(program: impl AsRef<OsStr>) -> Self {
Self(Command::new(program))
}
pub fn arg(&mut self, arg: impl AsRef<OsStr>) -> &mut Self {
self.0.arg(arg);
self
}
pub fn args<I, S>(&mut self, args: I) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<OsStr>,
{
self.0.args(args);
self
}
pub fn env(&mut self, key: impl AsRef<OsStr>, val: impl AsRef<OsStr>) -> &mut Self {
self.0.env(key, val);
self
}
pub fn current_dir(&mut self, dir: impl AsRef<std::path::Path>) -> &mut Self {
self.0.current_dir(dir);
self
}
pub fn stdin(&mut self, cfg: Stdio) -> &mut Self {
self.0.stdin(cfg);
self
}
pub fn stdout(&mut self, cfg: Stdio) -> &mut Self {
self.0.stdout(cfg);
self
}
pub fn stderr(&mut self, cfg: Stdio) -> &mut Self {
self.0.stderr(cfg);
self
}
pub fn spawn(&mut self) -> io::Result<DtactChild> {
self.0.spawn().map(DtactChild::new)
}
}
pub struct DtactChild {
inner: std::process::Child,
}
impl DtactChild {
const fn new(inner: std::process::Child) -> Self {
Self { inner }
}
#[must_use]
pub fn id(&self) -> u32 {
self.inner.id()
}
pub fn kill(&mut self) -> io::Result<()> {
self.inner.kill()
}
pub fn take_stdin(&mut self) -> Option<DtactChildStdin> {
self.inner.stdin.take().map(DtactChildStdin::new)
}
pub fn take_stdout(&mut self) -> Option<DtactChildStdout> {
self.inner.stdout.take().map(DtactChildStdout::new)
}
pub fn take_stderr(&mut self) -> Option<DtactChildStderr> {
self.inner.stderr.take().map(DtactChildStderr::new)
}
pub async fn wait(mut self) -> io::Result<ExitStatus> {
spawn_blocking(move || self.inner.wait()).await
}
pub async fn wait_with_output(self) -> io::Result<Output> {
spawn_blocking(move || self.inner.wait_with_output()).await
}
}
pub struct DtactChildStdin(Option<std::process::ChildStdin>);
impl DtactChildStdin {
const fn new(inner: std::process::ChildStdin) -> Self {
Self(Some(inner))
}
}
pub struct DtactChildStdout(Option<std::process::ChildStdout>);
impl DtactChildStdout {
const fn new(inner: std::process::ChildStdout) -> Self {
Self(Some(inner))
}
}
pub struct DtactChildStderr(Option<std::process::ChildStderr>);
impl DtactChildStderr {
const fn new(inner: std::process::ChildStderr) -> Self {
Self(Some(inner))
}
}
impl DtactChildStdin {
pub async fn write(&mut self, buf: Vec<u8>) -> io::Result<(usize, Vec<u8>)> {
use std::io::Write;
let mut handle = self
.0
.take()
.expect("dtact-process: concurrent write on the same DtactChildStdin");
let (result, handle) = spawn_blocking(move || {
let r = handle.write(&buf).map(|n| (n, buf));
(r, handle)
})
.await;
self.0 = Some(handle);
result
}
pub fn close(mut self) {
self.0 = None;
}
}
impl DtactChildStdout {
pub async fn read(&mut self, mut buf: Vec<u8>) -> io::Result<(usize, Vec<u8>)> {
use std::io::Read;
let mut handle = self
.0
.take()
.expect("dtact-process: concurrent read on the same DtactChildStdout");
let (result, handle) = spawn_blocking(move || {
let r = handle.read(&mut buf).map(|n| (n, buf));
(r, handle)
})
.await;
self.0 = Some(handle);
result
}
}
impl DtactChildStderr {
pub async fn read(&mut self, mut buf: Vec<u8>) -> io::Result<(usize, Vec<u8>)> {
use std::io::Read;
let mut handle = self
.0
.take()
.expect("dtact-process: concurrent read on the same DtactChildStderr");
let (result, handle) = spawn_blocking(move || {
let r = handle.read(&mut buf).map(|n| (n, buf));
(r, handle)
})
.await;
self.0 = Some(handle);
result
}
}