use std::fs;
use std::io::{self, BufRead, BufReader, Write};
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStderr, Command, Stdio};
use std::sync::mpsc;
use std::thread::{self, JoinHandle};
use std::time::Duration;
use haematite_cli::client::{Client, Reply, refused};
use haematite_cli::serve::cli::Config;
use haematite_cli::serve::protocol::Version;
use haematite_cli::serve::storage::Receipt;
use haematite_cli::serve::{registry, runtime};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use super::ProbeResult;
use super::options::{Binaries, Options};
pub const IO_BOUND: Duration = Duration::from_secs(300);
pub fn configure(directory: &Path, binaries: &Binaries, options: &Options) -> ProbeResult<PathBuf> {
let config = directory.join("serve.json");
fs::write(
&config,
serde_json::to_vec_pretty(&json!({
"data_dir": directory.join("database"),
"socket": directory.join("serve.sock"),
"name": directory.file_name().and_then(|name| name.to_str()).ok_or("directory name is not UTF-8")?,
"branch_cache_bytes_per_shard": 16_777_216,
"max_frame": 33_554_432, "max_clients": 4,
"max_page": options.records.max(1000), "timeout_ms": 300_000,
"query_worker": binaries.query,
"query_time_ms": 120_000, "query_memory_bytes": 268_435_456,
"query_outputs_per_record": 32, "query_max_running": 1,
"query_scan_records": u64::from(options.records) + u64::from(options.records / 10) + 21,
"create": {"shards": 2, "cache_bytes_total": 33_554_432, "executor_threads": 2}
}))?,
)?;
Ok(config)
}
pub struct Running {
child: Option<Child>,
stderr: Option<JoinHandle<io::Result<Vec<u8>>>>,
registry: Option<OwnedRegistry>,
}
impl Running {
pub fn start(binary: &Path, config_path: &Path) -> ProbeResult<Self> {
let config: Config = serde_json::from_slice(&fs::read(config_path)?)?;
let endpoint = runtime::resolve(&config.data_dir, config.name.as_deref(), None, true)?;
let digest = Sha256::digest(endpoint.name.as_bytes());
let entry_path = endpoint
.socket
.parent()
.ok_or("runtime socket has no parent")?
.join(format!("{digest:x}.json"));
let child = Command::new(binary)
.args(["serve", "--config"])
.arg(config_path)
.arg("--create")
.stdout(Stdio::null())
.stderr(Stdio::piped())
.spawn()?;
let mut running = Self {
child: Some(child),
stderr: None,
registry: None,
};
let pipe = running
.child
.as_mut()
.ok_or("missing serve child")?
.stderr
.take()
.ok_or("missing serve stderr")?;
let (ready, receiver) = mpsc::channel();
running.stderr = Some(
thread::Builder::new()
.name("probe-serve-stderr".to_owned())
.spawn(move || capture(pipe, &ready))?,
);
receiver
.recv()
.map_err(|error| format!("serve exited before its ready signal: {error}"))?;
running.registry = Some(OwnedRegistry::capture(
&entry_path,
&config,
running.pid()?,
)?);
Ok(running)
}
pub fn version(&self) -> ProbeResult<&str> {
Ok(&self
.registry
.as_ref()
.ok_or("serve registry entry was not captured")?
.entry
.version)
}
pub fn pid(&self) -> ProbeResult<u32> {
Ok(self.child.as_ref().ok_or("serve is already stopped")?.id())
}
pub fn stop(&mut self) -> ProbeResult<()> {
let mut errors = Vec::new();
let mut reaped = self.child.is_none();
if let Some(mut child) = self.child.take() {
let exited = match child.try_wait() {
Ok(Some(status)) => {
errors.push(format!(
"serve exited before the probe stopped it: {status}"
));
reaped = true;
true
}
Ok(None) => false,
Err(error) => {
errors.push(format!("serve status: {error}"));
false
}
};
if !exited {
if let Err(error) = child.kill() {
errors.push(format!(
"serve pid {} could not be stopped: {error}",
child.id()
));
}
match child.wait() {
Ok(_) => reaped = true,
Err(error) => errors.push(format!(
"serve pid {} could not be reaped: {error}",
child.id()
)),
}
}
}
if let Some(reader) = self.stderr.take() {
match reader.join() {
Ok(Ok(bytes)) => {
if let Err(error) = io::stderr().lock().write_all(&bytes) {
errors.push(format!("serve stderr output: {error}"));
}
}
Ok(Err(error)) => errors.push(format!("serve stderr read: {error}")),
Err(_) => errors.push("serve stderr reader panicked".to_owned()),
}
}
if let Some(entry) = self.registry.take() {
if reaped {
if let Err(error) = entry.remove() {
errors.push(format!("serve registry cleanup: {error}"));
}
} else {
errors.push("serve exit not confirmed; registry entry was retained".to_owned());
}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors.join("; ").into())
}
}
}
impl Drop for Running {
fn drop(&mut self) {
if let Err(error) = self.stop() {
eprintln!("probe serve cleanup failed: {error}");
}
}
}
fn capture(pipe: ChildStderr, ready: &mpsc::Sender<()>) -> io::Result<Vec<u8>> {
let mut reader = BufReader::new(pipe);
let mut bytes = Vec::new();
let mut announced = false;
loop {
let mut line = Vec::new();
if reader.read_until(b'\n', &mut line)? == 0 {
return Ok(bytes);
}
let is_ready = line.starts_with(b"haem/1 ready on ");
bytes.extend(line);
if is_ready && !announced {
ready.send(()).map_err(io::Error::other)?;
announced = true;
}
}
}
pub fn answered(client: &mut Client, method: &str, params: &Value) -> ProbeResult<Value> {
let raw = serde_json::value::to_raw_value(params)?;
match client.call(method, Some(&raw))? {
Reply::Answered(result) => Ok(serde_json::from_str(result.get())?),
Reply::Refused(refusal) => Err(refused(&refusal).into()),
}
}
pub struct Writer {
pub client: Client,
pub head: Version,
requests: u64,
}
impl Writer {
pub fn new(mut client: Client) -> ProbeResult<Self> {
let branches = answered(&mut client, "branches.list", &json!({"limit": 1}))?;
let main = branches
.get("entries")
.and_then(Value::as_array)
.and_then(|rows| {
rows.iter()
.find(|row| row.get("branch").and_then(Value::as_str) == Some("main"))
})
.ok_or("new store has no main branch")?;
let version = main.get("version").ok_or("main branch has no version")?;
Ok(Self {
client,
head: serde_json::from_value(version.clone())?,
requests: 0,
})
}
pub fn apply(&mut self, kind: &str, value: &Value) -> ProbeResult<()> {
self.requests = self
.requests
.checked_add(1)
.ok_or("request sequence exhausted")?;
let request_id = format!("probe-{}", self.requests);
let receipt: Receipt = serde_json::from_value(answered(
&mut self.client,
"mutation.apply",
&json!({
"request_id": request_id, "branch": "main", "expected": self.head,
"change": {"kind": kind, "value": value}
}),
)?)?;
if receipt.request_id != request_id || receipt.branch != "main" {
return Err("mutation receipt belongs to another request or branch".into());
}
self.head = receipt.version;
Ok(())
}
}
struct OwnedRegistry {
path: PathBuf,
identity: (u64, u64),
entry: registry::Entry,
}
impl OwnedRegistry {
fn capture(path: &Path, config: &Config, pid: u32) -> ProbeResult<Self> {
let entry = registry::read(path)?;
if entry.pid.get() != pid
|| entry.data_dir != config.data_dir
|| config.socket.as_ref() != Some(&entry.socket)
|| config.name.as_deref() != Some(entry.name.as_str())
{
return Err("serve registry entry does not belong to this probe child".into());
}
let metadata = fs::symlink_metadata(path)?;
Ok(Self {
path: path.to_owned(),
identity: (metadata.dev(), metadata.ino()),
entry,
})
}
fn remove(self) -> ProbeResult<()> {
let actual = registry::read(&self.path)?;
let metadata = fs::symlink_metadata(&self.path)?;
if actual != self.entry || (metadata.dev(), metadata.ino()) != self.identity {
return Err("serve registry entry was replaced; replacement was retained".into());
}
fs::remove_file(&self.path)?;
Ok(())
}
}