use super::{Capabilities, ErrorKind, ProviderError, ProviderResult};
use parking_lot::Mutex;
use serde_json::{Value, json};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::thread::JoinHandle;
use std::time::Duration;
mod transport;
mod wire;
#[cfg(test)]
mod tests;
use transport::{Lifecycle, Process, Shared, StderrMode, backoff_for, reader_loop, spawn_process};
pub const DEFAULT_TIMEOUT: Duration = Duration::from_millis(30_000);
pub struct StdioProvider {
name: String,
capabilities: Capabilities,
process: Mutex<Process>,
shared: Arc<Shared>,
current_id: AtomicU64,
timeout: Duration,
command: Vec<String>,
env: Vec<(String, String)>,
stderr_mode: StderrMode,
reader: Mutex<Option<JoinHandle<()>>>,
notice: Mutex<Option<String>>,
closed: AtomicBool,
}
impl Drop for StdioProvider {
fn drop(&mut self) {
self.closed.store(true, Ordering::Release);
let process = self.process.get_mut();
let _ = process.child.kill();
let _ = process.child.wait();
if let Some(handle) = self.reader.lock().take() {
let _ = handle.join();
}
}
}
impl StdioProvider {
pub fn spawn(command: &[String], timeout: Duration) -> ProviderResult<Self> {
Self::spawn_inner(command, timeout, &[], StderrMode::Null)
}
pub fn spawn_with_stderr(
command: &[String],
timeout: Duration,
inherit_stderr: bool,
) -> ProviderResult<Self> {
let mode = if inherit_stderr {
StderrMode::Inherit
} else {
StderrMode::Null
};
Self::spawn_inner(command, timeout, &[], mode)
}
#[cfg(test)]
fn spawn_with_env(
command: &[String],
timeout: Duration,
env: &[(&str, &str)],
) -> ProviderResult<Self> {
Self::spawn_inner(command, timeout, env, StderrMode::Null)
}
fn spawn_inner(
command: &[String],
timeout: Duration,
env: &[(&str, &str)],
stderr_mode: StderrMode,
) -> ProviderResult<Self> {
let (process, stdout) = spawn_process(command, env, stderr_mode)?;
let shared = Arc::new(Shared::default());
let reader = std::thread::spawn({
let shared = Arc::clone(&shared);
move || reader_loop(stdout, shared)
});
let provider = StdioProvider {
name: "stdio".into(),
capabilities: Capabilities {
orgs: true,
code_search: true,
},
process: Mutex::new(process),
shared,
current_id: AtomicU64::new(0),
timeout,
command: command.to_vec(),
env: env
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
stderr_mode,
reader: Mutex::new(Some(reader)),
notice: Mutex::new(None),
closed: AtomicBool::new(false),
};
let reply = provider.handshake()?;
provider.initialize_from(reply)
}
fn handshake(&self) -> ProviderResult<Value> {
let reply = self.exchange("initialize", json!({ "protocol": 1 }), false)?;
let protocol = reply.get("protocol").and_then(Value::as_u64).unwrap_or(1);
if protocol != 1 {
return Err(ProviderError::new(
ErrorKind::Provider,
format!("unsupported provider protocol {protocol}"),
));
}
Ok(reply)
}
fn initialize_from(mut self, reply: Value) -> ProviderResult<Self> {
if let Some(name) = reply.get("name").and_then(Value::as_str) {
self.name = format!("stdio:{name}");
}
if let Some(caps) = reply.get("capabilities") {
self.capabilities = Capabilities {
orgs: caps.get("orgs").and_then(Value::as_bool).unwrap_or(true),
code_search: caps
.get("code_search")
.and_then(Value::as_bool)
.unwrap_or(true),
};
}
Ok(self)
}
fn request(&self, method: &str, params: Value) -> ProviderResult<Value> {
self.ensure_alive()?;
self.exchange(method, params, true)
}
fn ensure_alive(&self) -> ProviderResult<()> {
let mut routing = self.shared.routing.lock();
let mut waited = false;
loop {
match routing.lifecycle {
Lifecycle::Alive => return Ok(()),
Lifecycle::Respawning => {
waited = true;
self.shared.changed.wait(&mut routing);
}
Lifecycle::Dead if waited => {
let reason = routing.restart_error.clone();
drop(routing);
return Err(match reason {
Some(why) => ProviderError::new(
ErrorKind::Provider,
format!("provider restart failed: {why}"),
),
None => ProviderError::new(
ErrorKind::Provider,
"provider restarting — try again",
),
});
}
Lifecycle::Dead => {
routing.lifecycle = Lifecycle::Respawning;
routing.restart_error = None;
let attempt = routing.restarts + 1;
drop(routing);
return self.rebuild(attempt);
}
}
}
}
fn rebuild(&self, attempt: u32) -> ProviderResult<()> {
std::thread::sleep(backoff_for(attempt));
if self.closed.load(Ordering::Acquire) {
let err = ProviderError::other("provider dropped during restart");
self.finish_rebuild(attempt, Err(&err));
return Err(err);
}
if let Some(handle) = self.reader.lock().take() {
let _ = handle.join();
}
let env: Vec<(&str, &str)> = self
.env
.iter()
.map(|(k, v)| (k.as_str(), v.as_str()))
.collect();
let (process, stdout) = spawn_process(&self.command, &env, self.stderr_mode)?;
let reader = std::thread::spawn({
let shared = Arc::clone(&self.shared);
move || reader_loop(stdout, shared)
});
*self.process.lock() = process;
*self.reader.lock() = Some(reader);
let proof = self.handshake();
if proof.is_ok() {
*self.notice.lock() = Some(format!(
"provider restarted (attempt {attempt}, backoff {}s)",
backoff_for(attempt).as_secs()
));
}
let outcome: Result<(), &ProviderError> = proof.as_ref().map(|_| ());
self.finish_rebuild(attempt, outcome);
proof.map(|_| ())
}
fn finish_rebuild(&self, attempt: u32, outcome: Result<(), &ProviderError>) {
let mut routing = self.shared.routing.lock();
match outcome {
Ok(()) => {
routing.lifecycle = Lifecycle::Alive;
routing.restarts = attempt;
routing.restart_error = None;
}
Err(e) => {
routing.lifecycle = Lifecycle::Dead;
routing.restart_error = Some(e.message.clone());
}
}
drop(routing);
self.shared.changed.notify_all();
}
}