frc/src/chromedriver.rs

110 lines
3.6 KiB
Rust

use futures::channel::oneshot::Sender;
use std::{
io::ErrorKind,
process::{Child, Command, Output, Stdio},
thread::{self, JoinHandle},
time::Duration,
};
use sysinfo::{Pid, ProcessRefreshKind, RefreshKind, System, SystemExt};
pub struct Driver {
child: Child,
}
impl Driver {
pub fn new() -> Result<Self, anyhow::Error> {
let mut child = Command::new("chromedriver")
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.arg("--port=6444")
.spawn()?;
let child_id = child.id();
let mut sys_info = System::new_with_specifics(
RefreshKind::new().with_processes(ProcessRefreshKind::new()),
);
// Wait for it to say Chromedriver started successfully
print!("chromedriver health check...");
loop {
// Check if chromedriver has exited
sys_info.refresh_all();
if let None = sys_info.process(Pid::from(to_pid(child_id.clone()))) {
let status_code = child.wait()?;
println!("{}", status_code.code().unwrap());
return Err(anyhow::anyhow!(
"chromedriver exited with status code {}",
status_code.code().unwrap_or(1337),
));
} else {
break;
}
}
println!(" done.");
// Check if chrome/chromium is already running. Die if so.
#[cfg(target_os = "linux")]
let chrome_names = ("chrome", "chromium");
#[cfg(target_os = "windows")]
let chrome_names = ("chrome.exe", "chromium.exe");
if sys_info.processes_by_exact_name(chrome_names.0).count() > 0
|| sys_info.processes_by_exact_name(chrome_names.1).count() > 0
{
return Err(anyhow::anyhow!("chrome/chromium is already running"));
}
println!("chromedriver started");
Ok(Self { child })
}
pub fn exit_waiter(&mut self, notify: Sender<()>) -> JoinHandle<()> {
// Check if chromedriver has exited
let mut sys_info = System::new_with_specifics(
RefreshKind::new().with_processes(ProcessRefreshKind::new()),
);
let pid = Pid::from(to_pid(self.child.id().clone()));
thread::spawn(move || loop {
sys_info.refresh_all();
let proc = sys_info.process(pid);
// Consider two cases of chromedriver being dead:
// 1) the process itself is gone (might be on windows)
// 2) the process is alive but not reaped (zombie) on linux
// in the second case, we check for if it has zero subtasks.
// For our purposes, that means it's dead.
if !proc.is_none() {
#[cfg(target_os = "linux")]
if proc.unwrap().tasks.len() != 0 {
thread::sleep(Duration::from_secs(3));
continue;
}
}
println!("chromedriver has exited");
notify.send(()).unwrap();
break;
})
}
pub fn exit_with_output(mut self) -> Result<Output, anyhow::Error> {
match self.child.kill() {
Err(e) => {
if e.kind() != ErrorKind::InvalidInput {
return Err(anyhow::anyhow!(e));
} else {
println!("killing child: {}", e);
}
}
_ => {}
}
Ok(self.child.wait_with_output()?)
}
}
#[cfg(target_os = "windows")]
fn to_pid(v: u32) -> usize {
v as usize
}
#[cfg(target_os = "linux")]
fn to_pid(v: u32) -> i32 {
v as i32
}