use std::collections::BTreeSet;
use std::env;
use std::fs::{self, DirBuilder};
use std::os::unix::fs::{DirBuilderExt, MetadataExt, PermissionsExt};
use std::path::{Path, PathBuf};
use std::process::Command;
use super::ProbeResult;
pub struct Options {
pub records: u32,
pub dims: usize,
pub repeat: usize,
pub dir: Option<PathBuf>,
pub keep: bool,
}
impl Options {
pub fn read() -> ProbeResult<Self> {
let mut options = Self {
records: 10_000,
dims: 384,
repeat: 5,
dir: None,
keep: false,
};
let mut args = env::args_os().skip(1);
while let Some(argument) = args.next() {
if argument == "--keep" {
options.keep = true;
continue;
}
let value = args.next().ok_or_else(|| {
format!(
"missing value for {}; use --records N --dims D --repeat R --dir PATH --keep",
argument.to_string_lossy()
)
})?;
if argument == "--dir" {
options.dir = Some(PathBuf::from(value));
continue;
}
let text = value.to_str().ok_or("numeric argument is not UTF-8")?;
if argument == "--records" {
options.records = text.parse()?;
} else if argument == "--dims" {
options.dims = text.parse()?;
} else if argument == "--repeat" {
options.repeat = text.parse()?;
} else {
return Err(format!("unknown argument: {}", argument.to_string_lossy()).into());
}
}
if options.records < 10 || options.dims == 0 || options.repeat == 0 {
return Err("records must be at least 10; dims and repeat must be positive".into());
}
options
.records
.checked_add(options.records / 10)
.ok_or("record and relation counts exceed u32")?;
Ok(options)
}
}
pub struct Binaries {
pub serve: PathBuf,
pub query: PathBuf,
}
impl Binaries {
pub fn find() -> ProbeResult<Self> {
let executable = env::current_exe()?;
let examples = executable
.parent()
.ok_or("example has no parent directory")?;
if examples.file_name().is_none_or(|name| name != "examples") {
return Err("run the example from Cargo's profile/examples directory".into());
}
let profile = examples
.parent()
.ok_or("example has no build profile directory")?;
Ok(Self {
serve: built_binary(profile, "haematite")?,
query: built_binary(profile, "haematite-jq")?,
})
}
}
fn built_binary(profile: &Path, name: &str) -> ProbeResult<PathBuf> {
let candidate = profile.join(name);
let path = fs::canonicalize(&candidate).map_err(|error| {
format!(
"built binary {} is unavailable: {error}; build it in the example's profile first",
candidate.display()
)
})?;
let metadata = fs::metadata(&path)?;
if !metadata.is_file() || metadata.permissions().mode() & 0o111 == 0 {
return Err(format!("built binary {} is not an executable file", path.display()).into());
}
Ok(path)
}
pub fn commit() -> ProbeResult<String> {
command_text(
Command::new("git")
.arg("-C")
.arg(env!("CARGO_MANIFEST_DIR"))
.args(["rev-parse", "--short", "HEAD"]),
)
}
pub fn machine() -> ProbeResult<String> {
let host = command_text(&mut Command::new("hostname"))?;
Ok(format!(
"{} {} {} host={host}",
env::consts::OS,
env::consts::ARCH,
env::consts::FAMILY
))
}
fn command_text(command: &mut Command) -> ProbeResult<String> {
let output = command.output()?;
if !output.status.success() {
return Err(format!(
"metadata command ended with {}: {}",
output.status,
String::from_utf8(output.stderr)?
)
.into());
}
let text = String::from_utf8(output.stdout)?.trim().to_owned();
if text.is_empty() {
return Err("metadata command returned empty text".into());
}
Ok(text)
}
pub struct Directory {
pub path: PathBuf,
keep: bool,
active: bool,
}
impl Directory {
pub fn create(requested: Option<&Path>, keep: bool) -> ProbeResult<Self> {
let path = match requested {
Some(path) => {
let path = env::current_dir()?.join(path);
DirBuilder::new().mode(0o700).create(&path)?;
path
}
None => tempfile::Builder::new()
.prefix("hq")
.permissions(std::fs::Permissions::from_mode(0o700))
.tempdir_in("/tmp")?
.keep(),
};
let mut directory = Self {
path,
keep,
active: true,
};
directory.path = fs::canonicalize(&directory.path)?;
Ok(directory)
}
pub fn finish(&mut self) -> ProbeResult<()> {
if self.active {
self.active = false;
if !self.keep {
fs::remove_dir_all(&self.path)?;
}
}
Ok(())
}
}
impl Drop for Directory {
fn drop(&mut self) {
if let Err(error) = self.finish() {
eprintln!("probe directory cleanup failed: {error}");
}
}
}
pub struct DirectoryBytes {
pub logical: u64,
pub allocated: u64,
}
pub fn directory_bytes(path: &Path) -> ProbeResult<DirectoryBytes> {
let root = fs::symlink_metadata(path)?;
if !root.is_dir() {
return Err("store size walk needs a non-symlink directory".into());
}
let mut bytes = DirectoryBytes {
logical: 0,
allocated: root
.blocks()
.checked_mul(512)
.ok_or("allocated size exceeds u64")?,
};
let mut visited = BTreeSet::from([(root.dev(), root.ino())]);
let mut pending = vec![path.to_owned()];
while let Some(path) = pending.pop() {
for entry in fs::read_dir(path)? {
let entry = entry?;
let metadata = fs::symlink_metadata(entry.path())?;
if metadata.file_type().is_symlink() {
return Err(
format!("store size walk refuses symlink {}", entry.path().display()).into(),
);
}
if !visited.insert((metadata.dev(), metadata.ino())) {
continue;
}
let allocated = metadata
.blocks()
.checked_mul(512)
.ok_or("allocated size exceeds u64")?;
bytes.allocated = bytes
.allocated
.checked_add(allocated)
.ok_or("allocated size exceeds u64")?;
if metadata.is_dir() {
pending.push(entry.path());
} else if metadata.is_file() {
bytes.logical = bytes
.logical
.checked_add(metadata.len())
.ok_or("store size exceeds u64")?;
} else {
return Err(format!(
"store size walk refuses non-file {}",
entry.path().display()
)
.into());
}
}
}
Ok(bytes)
}