use super::StdioProvider;
use crate::provider::{ErrorKind, ProviderError, ProviderResult};
use parking_lot::{Condvar, Mutex};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::sync::mpsc;
use std::time::Duration;
#[derive(Clone, Copy)]
pub(super) enum StderrMode {
Null,
Inherit,
}
pub(super) fn backoff_for(restart_attempt: u32) -> Duration {
match restart_attempt {
1 => Duration::from_secs(1),
2 => Duration::from_secs(2),
3 => Duration::from_secs(5),
_ => Duration::from_secs(30),
}
}
pub(super) struct Process {
pub(super) child: Child,
pub(super) stdin: ChildStdin,
}
pub(super) struct Shared {
pub(super) routing: Mutex<Routing>,
pub(super) changed: Condvar,
}
impl Default for Shared {
fn default() -> Self {
Shared {
routing: Mutex::new(Routing::default()),
changed: Condvar::new(),
}
}
}
#[derive(Default)]
pub(super) struct Routing {
pub(super) next_id: u64,
pub(super) pending: HashMap<u64, mpsc::Sender<Value>>,
pub(super) lifecycle: Lifecycle,
pub(super) restarts: u32,
pub(super) restart_error: Option<String>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(super) enum Lifecycle {
#[default]
Alive,
Respawning,
Dead,
}
impl StdioProvider {
pub(super) fn exchange(
&self,
method: &str,
params: Value,
gated: bool,
) -> ProviderResult<Value> {
let (id, rx) = {
let mut routing = self.shared.routing.lock();
if gated && routing.lifecycle != Lifecycle::Alive {
return Err(ProviderError::new(
ErrorKind::Provider,
"provider restarting — try again",
));
}
routing.next_id += 1;
let id = routing.next_id;
let (tx, rx) = mpsc::channel::<Value>();
routing.pending.insert(id, tx);
(id, rx)
};
let line = json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params,
})
.to_string();
{
let mut process = self.process.lock();
let written = writeln!(process.stdin, "{line}").and_then(|()| process.stdin.flush());
if let Err(e) = written {
let mut routing = self.shared.routing.lock();
routing.pending.remove(&id);
if routing.lifecycle == Lifecycle::Alive {
routing.lifecycle = Lifecycle::Dead;
}
return Err(ProviderError::new(
ErrorKind::Provider,
format!("provider write: {e}"),
));
}
}
let _cleared = CurrentIdGuard::new(self, id);
let msg = match rx.recv_timeout(self.timeout) {
Ok(msg) => msg,
Err(mpsc::RecvTimeoutError::Timeout) => {
self.shared.routing.lock().pending.remove(&id);
return Err(ProviderError::new(
ErrorKind::Timeout,
format!("provider timeout after {}s", self.timeout.as_secs_f64()),
));
}
Err(mpsc::RecvTimeoutError::Disconnected) => {
return Err(ProviderError::new(
ErrorKind::Provider,
"provider closed its output",
));
}
};
if let Some(err) = msg.get("error") {
return Err(error_from_reply(err));
}
msg.get("result")
.cloned()
.ok_or_else(|| ProviderError::other("provider reply without result"))
}
}
fn error_from_reply(err: &Value) -> ProviderError {
let message = err
.get("message")
.and_then(Value::as_str)
.unwrap_or("provider error");
let data = err.get("data");
let kind = data
.and_then(|d| d.get("kind"))
.and_then(Value::as_str)
.map(|k| match k {
"auth" => ErrorKind::Auth,
"rate_limited" => ErrorKind::RateLimited,
"not_found" => ErrorKind::NotFound,
"network" => ErrorKind::Network,
"timeout" => ErrorKind::Timeout,
"provider" => ErrorKind::Provider,
_ => ErrorKind::Other,
})
.unwrap_or(ErrorKind::Other);
let error = ProviderError::new(kind, message);
match data
.and_then(|d| d.get("retry_after_s"))
.and_then(Value::as_u64)
{
Some(s) => error.with_retry_after(Duration::from_secs(s)),
None => error,
}
}
pub(super) fn spawn_process(
command: &[String],
env: &[(&str, &str)],
stderr_mode: StderrMode,
) -> ProviderResult<(Process, ChildStdout)> {
let (program, args) = command
.split_first()
.ok_or_else(|| ProviderError::other("empty provider command"))?;
let mut cmd = Command::new(program);
cmd.args(args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(match stderr_mode {
StderrMode::Null => Stdio::null(),
StderrMode::Inherit => Stdio::inherit(),
});
for (key, value) in env {
cmd.env(key, value);
}
let mut child = cmd
.spawn()
.map_err(|e| ProviderError::other(format!("spawn {program}: {e}")))?;
let stdin = child.stdin.take().expect("piped stdin");
let stdout = child.stdout.take().expect("piped stdout");
Ok((Process { child, stdin }, stdout))
}
pub(super) fn reader_loop(stdout: ChildStdout, shared: Arc<Shared>) {
let mut stdout = BufReader::new(stdout);
let mut buf = String::new();
loop {
buf.clear();
let n = match stdout.read_line(&mut buf) {
Ok(n) => n,
Err(_) => break,
};
if n == 0 {
break; }
let Ok(msg) = serde_json::from_str::<Value>(buf.trim()) else {
continue; };
let Some(id) = msg.get("id").and_then(Value::as_u64) else {
continue; };
let tx = shared.routing.lock().pending.remove(&id);
if let Some(tx) = tx {
let _ = tx.send(msg);
}
}
let mut routing = shared.routing.lock();
routing.pending.clear();
routing.lifecycle = Lifecycle::Dead;
}
struct CurrentIdGuard<'a> {
provider: &'a StdioProvider,
}
impl<'a> CurrentIdGuard<'a> {
fn new(provider: &'a StdioProvider, id: u64) -> Self {
provider.current_id.store(id, Ordering::Release);
CurrentIdGuard { provider }
}
}
impl Drop for CurrentIdGuard<'_> {
fn drop(&mut self) {
self.provider.current_id.store(0, Ordering::Release);
}
}
pub(super) fn cancel_notification(id: u64) -> String {
json!({
"jsonrpc": "2.0",
"method": "$/cancelRequest",
"params": { "id": id },
})
.to_string()
}
pub(super) fn de<T: serde::de::DeserializeOwned>(v: Value) -> ProviderResult<T> {
serde_json::from_value(v)
.map_err(|e| ProviderError::other(format!("provider reply shape: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cancel_notification_shape() {
let line = cancel_notification(7);
let v: Value = serde_json::from_str(&line).unwrap();
assert_eq!(v["jsonrpc"], "2.0");
assert_eq!(v["method"], "$/cancelRequest");
assert_eq!(v["params"]["id"], 7);
assert!(v.get("id").is_none()); }
#[test]
fn backoff_ladder() {
assert_eq!(backoff_for(1), Duration::from_secs(1));
assert_eq!(backoff_for(2), Duration::from_secs(2));
assert_eq!(backoff_for(3), Duration::from_secs(5));
assert_eq!(backoff_for(9), Duration::from_secs(30));
}
}