dre_protocol/

host.rs

//! Core's side of the protocol: spawn a plugin, negotiate a version, send requests.
//!
//! Frames are read on a background thread so core can time out a silent plugin during the
//! handshake and never blocks forever on a plugin that has died. stderr is drained on another
//! thread, forwarded to a log sink and kept (last lines) for error messages.

use std::collections::VecDeque;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, Command, ExitStatus, Stdio};
use std::sync::mpsc::{Receiver, RecvTimeoutError, channel};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use serde_json::{Map, Value};

use crate::frame::{self, Frame, FrameError};
use crate::msg::{ConnectionField, DeliveryFile, Request, Response, ResultSetMeta};
use crate::options::OptionField;
use crate::{Kind, MAX_VERSION, MIN_VERSION, PluginId, parse_executable_name};

/// Receives each stderr line a plugin writes.
pub type LogSink = Arc<dyn Fn(&str, &str) + Send + Sync>;

/// A log sink that prefixes each plugin line with the plugin's file name, on core's stderr.
pub fn stderr_log() -> LogSink {
    Arc::new(|plugin, line| eprintln!("[{plugin}] {line}"))
}

pub const DEFAULT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);

/// Set to `1` in a plugin's environment when a person is at core's terminal, else `0`.
pub const INTERACTIVE_ENV: &str = "DRE_INTERACTIVE";
const STDERR_TAIL: usize = 20;

#[derive(Debug, Clone, PartialEq)]
pub struct PluginInfo {
    pub protocol_version: u32,
    pub kind: Kind,
    pub name: String,
    pub version: String,
    pub capabilities: Vec<String>,
    /// Every plugin the executable serves (at least the one it's serving now).
    pub provides: Vec<PluginId>,
}

#[derive(Debug)]
pub enum HostError {
    Spawn {
        plugin: String,
        error: std::io::Error,
    },
    Crashed {
        plugin: String,
        status: Option<ExitStatus>,
        stderr: Vec<String>,
    },
    Malformed {
        plugin: String,
        message: String,
    },
    Incompatible {
        plugin: String,
        core: (u32, u32),
        plugin_range: (u32, u32),
    },
    Timeout {
        plugin: String,
        waiting_for: &'static str,
    },
    Unexpected {
        plugin: String,
        expected: &'static str,
        got: String,
    },
    /// The plugin reported an error for a request.
    Plugin {
        plugin: String,
        message: String,
    },
    Arrow {
        plugin: String,
        message: String,
    },
}

impl std::fmt::Display for HostError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HostError::Spawn { plugin, error } => write!(f, "can't start plugin `{plugin}`: {error}"),
            HostError::Crashed {
                plugin,
                status,
                stderr,
            } => {
                match status {
                    Some(s) => write!(f, "plugin `{plugin}` exited unexpectedly ({s})")?,
                    None => write!(f, "plugin `{plugin}` closed its output unexpectedly")?,
                }
                if !stderr.is_empty() {
                    write!(f, "; its last log lines:\n  {}", stderr.join("\n  "))?;
                }
                Ok(())
            }
            HostError::Malformed { plugin, message } => {
                write!(f, "plugin `{plugin}` broke the protocol: {message}")
            }
            HostError::Incompatible {
                plugin,
                core,
                plugin_range,
            } => write!(
                f,
                "plugin `{plugin}` speaks protocol versions {}..={}, but this DRE core speaks {}..={}; update {}",
                plugin_range.0,
                plugin_range.1,
                core.0,
                core.1,
                if plugin_range.1 < core.0 {
                    "the plugin"
                } else {
                    "DRE"
                }
            ),
            HostError::Timeout { plugin, waiting_for } => {
                write!(
                    f,
                    "plugin `{plugin}` didn't answer in time (waiting for {waiting_for})"
                )
            }
            HostError::Unexpected {
                plugin,
                expected,
                got,
            } => {
                write!(f, "plugin `{plugin}` sent {got} where {expected} was expected")
            }
            HostError::Plugin { message, .. } => f.write_str(message),
            HostError::Arrow { plugin, message } => {
                write!(f, "plugin `{plugin}` sent invalid Arrow data: {message}")
            }
        }
    }
}

impl std::error::Error for HostError {}

pub type Result<T> = std::result::Result<T, HostError>;

/// A message from the plugin.
pub enum Incoming {
    Json(Response),
    Arrow(Vec<u8>),
}

impl std::fmt::Debug for Incoming {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Incoming::Json(r) => write!(f, "Json({r:?})"),
            Incoming::Arrow(b) => write!(f, "Arrow frame ({} bytes)", b.len()),
        }
    }
}

/// The outcome of a source `execute`.
#[derive(Debug)]
pub enum Execution {
    NoResult { rows_affected: Option<u64> },
    Result { schema: SchemaRef, rows: u64 },
}

pub struct PluginProcess {
    label: String,
    path: PathBuf,
    child: Child,
    stdin: Option<BufWriter<ChildStdin>>,
    rx: Receiver<std::result::Result<Frame, FrameError>>,
    stderr: Arc<Mutex<VecDeque<String>>>,
    /// Forwards stderr; joined on close/drop so no log line is lost when the plugin exits.
    stderr_thread: Option<std::thread::JoinHandle<()>>,
    info: Option<PluginInfo>,
    /// A reply that arrived while core was still sending (a format failing mid-stream).
    early: Option<std::result::Result<Frame, FrameError>>,
    /// The plugin asked for in the handshake. By default the one a `dre-<kind>-<name>` file name
    /// names; none for a package executable, which then serves its first plugin.
    serve: Option<PluginId>,
}

impl PluginProcess {
    /// Spawn the plugin at `path` and complete the handshake.
    pub fn start(path: &Path, log: LogSink) -> Result<PluginProcess> {
        Self::start_in(path, log, None)
    }

    /// Like `start`, running the plugin in `cwd` (core uses the project directory).
    pub fn start_in(path: &Path, log: LogSink, cwd: Option<&Path>) -> Result<PluginProcess> {
        Self::start_for(path, None, log, cwd)
    }

    /// Like `start_in`, asking the executable for `plugin` (one of a package's plugins).
    pub fn start_for(
        path: &Path,
        plugin: Option<&PluginId>,
        log: LogSink,
        cwd: Option<&Path>,
    ) -> Result<PluginProcess> {
        let timeout = std::env::var("DRE_PLUGIN_HANDSHAKE_TIMEOUT_MS")
            .ok()
            .and_then(|v| v.parse().ok())
            .map(Duration::from_millis)
            .unwrap_or(DEFAULT_HANDSHAKE_TIMEOUT);
        let mut p = Self::spawn_in(path, log, &[], cwd)?;
        if let Some(id) = plugin {
            p.serve = Some(id.clone());
        }
        p.handshake((MIN_VERSION, MAX_VERSION), timeout)?;
        Ok(p)
    }

    /// Spawn and handshake offering an explicit version range (the conformance suite uses this).
    pub fn start_with(
        path: &Path,
        log: LogSink,
        versions: (u32, u32),
        timeout: Duration,
    ) -> Result<PluginProcess> {
        let mut p = Self::spawn(path, log)?;
        p.handshake(versions, timeout)?;
        Ok(p)
    }

    /// Spawn without a handshake (for tests that probe raw protocol behaviour).
    pub fn spawn(path: &Path, log: LogSink) -> Result<PluginProcess> {
        Self::spawn_env(path, log, &[])
    }

    /// Spawn with extra environment variables, without a handshake.
    pub fn spawn_env(path: &Path, log: LogSink, env: &[(&str, &str)]) -> Result<PluginProcess> {
        Self::spawn_in(path, log, env, None)
    }

    /// Spawn with extra environment variables and a working directory, without a handshake.
    pub fn spawn_in(
        path: &Path,
        log: LogSink,
        env: &[(&str, &str)],
        cwd: Option<&Path>,
    ) -> Result<PluginProcess> {
        let label = path
            .file_name()
            .map(|f| f.to_string_lossy().to_string())
            .unwrap_or_default();
        let mut cmd = Command::new(path);
        if let Some(dir) = cwd {
            cmd.current_dir(dir);
        }
        // Plugins can't see the terminal (their stdio is piped), so tell them whether a person is
        // there, e.g. to finish a browser sign-in. An existing DRE_INTERACTIVE is left alone.
        if std::env::var_os(INTERACTIVE_ENV).is_none() {
            use std::io::IsTerminal;
            let there = std::io::stderr().is_terminal() && std::io::stdin().is_terminal();
            cmd.env(INTERACTIVE_ENV, if there { "1" } else { "0" });
        }
        let mut child = cmd
            .envs(env.iter().copied())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|error| HostError::Spawn {
                plugin: label.clone(),
                error,
            })?;
        let stdout = child.stdout.take().unwrap();
        let stderr_pipe = child.stderr.take().unwrap();
        let stdin = child.stdin.take().map(BufWriter::new);

        let (tx, rx) = channel();
        std::thread::spawn(move || {
            let mut r = BufReader::new(stdout);
            loop {
                let f = frame::read_frame(&mut r);
                let stop = f.is_err();
                if tx.send(f).is_err() || stop {
                    break;
                }
            }
        });
        let stderr = Arc::new(Mutex::new(VecDeque::new()));
        let tail = stderr.clone();
        let who = label.clone();
        let stderr_thread = std::thread::spawn(move || {
            for line in BufReader::new(stderr_pipe).lines() {
                let Ok(line) = line else { break };
                log(&who, &line);
                let mut t = tail.lock().unwrap();
                if t.len() == STDERR_TAIL {
                    t.pop_front();
                }
                t.push_back(line);
            }
        });
        let serve = parse_executable_name(&label).map(|(k, n)| PluginId::new(k, n));
        Ok(PluginProcess {
            label,
            path: path.to_path_buf(),
            child,
            stdin,
            rx,
            stderr,
            stderr_thread: Some(stderr_thread),
            info: None,
            early: None,
            serve,
        })
    }

    /// Name the plugin to ask for in the handshake (`None`: a package's first).
    pub fn ask_for(&mut self, plugin: Option<PluginId>) {
        self.serve = plugin;
    }

    /// Complete the handshake on a process from `spawn`/`spawn_env`.
    pub fn handshake(&mut self, (min, max): (u32, u32), timeout: Duration) -> Result<()> {
        self.send(&Request::Hello {
            min_version: min,
            max_version: max,
            core_version: core_version(),
            plugin: self.serve.clone(),
        })?;
        match self.recv(Some(timeout), "the hello reply")? {
            Incoming::Json(Response::Hello {
                protocol_version,
                kind,
                name,
                version,
                capabilities,
                mut provides,
            }) => {
                if protocol_version < min || protocol_version > max {
                    return Err(HostError::Incompatible {
                        plugin: self.label.clone(),
                        core: (min, max),
                        plugin_range: (protocol_version, protocol_version),
                    });
                }
                if let Some(want) = &self.serve
                    && (want.kind != kind || want.name != name)
                {
                    return Err(HostError::Plugin {
                        plugin: self.label.clone(),
                        message: format!(
                            "`{}` was asked for the {want} plugin but serves {kind}/{name}",
                            self.label
                        ),
                    });
                }
                if provides.is_empty() {
                    provides.push(PluginId::new(kind, name.clone()));
                }
                self.info = Some(PluginInfo {
                    protocol_version,
                    kind,
                    name,
                    version,
                    capabilities,
                    provides,
                });
                Ok(())
            }
            Incoming::Json(Response::VersionMismatch {
                min_version,
                max_version,
            }) => Err(HostError::Incompatible {
                plugin: self.label.clone(),
                core: (min, max),
                plugin_range: (min_version, max_version),
            }),
            Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
            }),
            other => Err(self.unexpected("a hello reply", &other)),
        }
    }

    pub fn label(&self) -> &str {
        &self.label
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn info(&self) -> &PluginInfo {
        self.info.as_ref().expect("handshake completed")
    }

    pub fn has(&self, capability: &str) -> bool {
        self.info
            .as_ref()
            .is_some_and(|i| i.capabilities.iter().any(|c| c == capability))
    }

    fn unexpected(&self, expected: &'static str, got: &Incoming) -> HostError {
        let got = match got {
            Incoming::Json(r) => format!("{:?}", r)
                .split([' ', '{'])
                .next()
                .unwrap_or("?")
                .to_string(),
            Incoming::Arrow(_) => "Arrow data".into(),
        };
        HostError::Unexpected {
            plugin: self.label.clone(),
            expected,
            got: format!("`{got}`"),
        }
    }

    fn crashed(&mut self) -> HostError {
        // Give the process a moment to finish exiting so its status and last stderr are known.
        let mut status = None;
        for _ in 0..50 {
            if let Ok(Some(s)) = self.child.try_wait() {
                status = Some(s);
                break;
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        std::thread::sleep(Duration::from_millis(50));
        let stderr = self.stderr.lock().unwrap().iter().cloned().collect();
        HostError::Crashed {
            plugin: self.label.clone(),
            status,
            stderr,
        }
    }

    pub fn send(&mut self, req: &Request) -> Result<()> {
        let ok = match self.stdin.as_mut() {
            Some(w) => frame::write_json(w, req).is_ok(),
            None => false,
        };
        if ok { Ok(()) } else { Err(self.write_failed()) }
    }

    pub fn send_batch(&mut self, batch: &RecordBatch) -> Result<()> {
        // A plugin that fails part-way replies before reading the rest: stop sending and
        // surface its error rather than streaming everything into a plugin that gave up.
        if self.early.is_none()
            && let Ok(f) = self.rx.try_recv()
        {
            self.early = Some(f);
        }
        if let Some(Ok(Frame::Json(v))) = &self.early
            && let Ok(Response::Error { message }) = serde_json::from_value::<Response>(v.clone())
        {
            return Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
            });
        }
        let ipc = frame::encode_batch(batch).map_err(|e| HostError::Arrow {
            plugin: self.label.clone(),
            message: e.to_string(),
        })?;
        let ok = match self.stdin.as_mut() {
            Some(w) => frame::write_arrow(w, &ipc).is_ok(),
            None => false,
        };
        if ok { Ok(()) } else { Err(self.write_failed()) }
    }

    /// Writing failed: the plugin has gone away. Prefer its own error message if it sent one.
    fn write_failed(&mut self) -> HostError {
        match self.rx.recv_timeout(Duration::from_secs(2)) {
            Ok(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v) {
                Ok(Response::Error { message }) => HostError::Plugin {
                    plugin: self.label.clone(),
                    message,
                },
                _ => self.crashed(),
            },
            _ => self.crashed(),
        }
    }

    /// Receive the next message. `timeout` of `None` waits as long as the plugin is alive.
    pub fn recv(&mut self, timeout: Option<Duration>, waiting_for: &'static str) -> Result<Incoming> {
        let got = match timeout {
            _ if self.early.is_some() => self.early.take(),
            Some(t) => match self.rx.recv_timeout(t) {
                Ok(f) => Some(f),
                Err(RecvTimeoutError::Timeout) => {
                    return Err(HostError::Timeout {
                        plugin: self.label.clone(),
                        waiting_for,
                    });
                }
                Err(RecvTimeoutError::Disconnected) => None,
            },
            None => self.rx.recv().ok(),
        };
        match got {
            Some(Ok(Frame::Json(v))) => match serde_json::from_value::<Response>(v.clone()) {
                Ok(r) => Ok(Incoming::Json(r)),
                Err(e) => Err(HostError::Malformed {
                    plugin: self.label.clone(),
                    message: format!("unknown message {v}: {e}"),
                }),
            },
            Some(Ok(Frame::Arrow(b))) => Ok(Incoming::Arrow(b)),
            Some(Err(FrameError::Malformed(m))) => Err(HostError::Malformed {
                plugin: self.label.clone(),
                message: m,
            }),
            Some(Err(_)) | None => Err(self.crashed()),
        }
    }

    /// Receive a JSON response, turning `error` into `HostError::Plugin`.
    pub fn recv_json(&mut self, waiting_for: &'static str) -> Result<Response> {
        match self.recv(None, waiting_for)? {
            Incoming::Json(Response::Error { message }) => Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
            }),
            Incoming::Json(r) => Ok(r),
            other => Err(self.unexpected(waiting_for, &other)),
        }
    }

    /// Send a request and expect `ok`.
    fn call_ok(&mut self, req: &Request, what: &'static str) -> Result<()> {
        self.send(req)?;
        match self.recv_json(what)? {
            Response::Ok {} => Ok(()),
            other => Err(self.unexpected(what, &Incoming::Json(other))),
        }
    }

    pub fn describe(&mut self) -> Result<Vec<ConnectionField>> {
        Ok(self.describe_all()?.0)
    }

    /// The connection fields and the options the plugin declares.
    pub fn describe_all(&mut self) -> Result<(Vec<ConnectionField>, Vec<OptionField>)> {
        self.send(&Request::Describe {})?;
        match self.recv_json("a describe reply")? {
            Response::Describe {
                connection_fields,
                option_fields,
            } => Ok((connection_fields, option_fields)),
            other => Err(self.unexpected("a describe reply", &Incoming::Json(other))),
        }
    }

    /// Check a config block of options (needs `validate`); returns every problem found.
    pub fn validate(&mut self, options: Map<String, Value>) -> Result<Vec<String>> {
        self.send(&Request::Validate { options })?;
        match self.recv_json("a validate reply")? {
            Response::Validated { errors } => Ok(errors),
            other => Err(self.unexpected("a validate reply", &Incoming::Json(other))),
        }
    }

    pub fn open(&mut self, connection: Map<String, Value>, read_only: bool) -> Result<()> {
        self.call_ok(
            &Request::Open {
                connection,
                read_only,
            },
            "an open reply",
        )
    }

    pub fn check(&mut self, sql: &str) -> Result<()> {
        self.call_ok(&Request::Check { sql: sql.to_string() }, "a check reply")
    }

    /// Run one statement, handing every batch to `on_batch` as it arrives.
    pub fn execute(
        &mut self,
        sql: &str,
        row_limit: Option<u64>,
        mut on_batch: impl FnMut(&SchemaRef, RecordBatch) -> std::result::Result<(), String>,
    ) -> Result<Execution> {
        self.send(&Request::Execute {
            sql: sql.to_string(),
            row_limit,
        })?;
        match self.recv_json("an execute reply")? {
            Response::NoResult { rows_affected } => return Ok(Execution::NoResult { rows_affected }),
            Response::Result { .. } => {}
            other => return Err(self.unexpected("an execute reply", &Incoming::Json(other))),
        }
        let mut schema: Option<SchemaRef> = None;
        let mut rows = 0u64;
        let mut sink_error = None;
        loop {
            match self.recv(None, "result data")? {
                Incoming::Arrow(ipc) => {
                    let (s, batches) = frame::decode_batches(&ipc).map_err(|e| HostError::Arrow {
                        plugin: self.label.clone(),
                        message: e.to_string(),
                    })?;
                    let schema = schema.get_or_insert(s);
                    for b in batches {
                        rows += b.num_rows() as u64;
                        if sink_error.is_none()
                            && let Err(e) = on_batch(schema, b)
                        {
                            // Keep draining so the session stays usable; report after.
                            sink_error = Some(e);
                        }
                    }
                }
                Incoming::Json(Response::ResultEnd { .. }) => break,
                Incoming::Json(Response::Error { message }) => {
                    return Err(HostError::Plugin {
                        plugin: self.label.clone(),
                        message,
                    });
                }
                other => return Err(self.unexpected("result data", &other)),
            }
        }
        if let Some(e) = sink_error {
            return Err(HostError::Plugin {
                plugin: self.label.clone(),
                message: e,
            });
        }
        let schema = schema.ok_or_else(|| HostError::Malformed {
            plugin: self.label.clone(),
            message: "a result ended without any Arrow frame carrying its schema".into(),
        })?;
        Ok(Execution::Result { schema, rows })
    }

    /// Start a format `write`; follow with `write_result_set` per result set, then `write_finish`.
    pub fn write_begin(
        &mut self,
        path: &str,
        format: &str,
        options: Map<String, Value>,
        result_sets: Vec<ResultSetMeta>,
        template: Option<Value>,
    ) -> Result<()> {
        self.send(&Request::Write {
            path: path.to_string(),
            format: format.to_string(),
            options,
            result_sets,
            template,
        })
    }

    /// Stream one result set: at least one batch (carrying the schema), then the end marker.
    pub fn write_result_set(
        &mut self,
        schema: &SchemaRef,
        batches: impl IntoIterator<Item = RecordBatch>,
    ) -> Result<()> {
        let mut any = false;
        for b in batches {
            any = true;
            self.send_batch(&b)?;
        }
        if !any {
            self.send_batch(&RecordBatch::new_empty(schema.clone()))?;
        }
        self.send(&Request::ResultSetEnd {})
    }

    /// Load rows into a temporary table on the source's session. Returns the relation to use in
    /// SQL, and the plugin's warning, if any. Only for plugins advertising `load`.
    pub fn load(
        &mut self,
        name: &str,
        schema: &SchemaRef,
        batches: impl IntoIterator<Item = RecordBatch>,
    ) -> Result<(String, Option<String>)> {
        self.send(&Request::Load {
            name: name.to_string(),
        })?;
        self.write_result_set(schema, batches)?;
        match self.recv_json("a loaded reply")? {
            Response::Loaded {
                relation, warning, ..
            } => Ok((relation, warning)),
            other => Err(self.unexpected("a loaded reply", &Incoming::Json(other))),
        }
    }

    pub fn write_finish(&mut self) -> Result<Vec<String>> {
        Ok(self.write_finish_with_warnings()?.0)
    }

    /// `write_finish`, also returning the format's warnings.
    pub fn write_finish_with_warnings(&mut self) -> Result<(Vec<String>, Vec<String>)> {
        self.send(&Request::Finish {})?;
        match self.recv_json("a written reply")? {
            Response::Written { files, warnings } => Ok((files, warnings)),
            other => Err(self.unexpected("a written reply", &Incoming::Json(other))),
        }
    }

    /// Deliver one file with no plugin options.
    pub fn deliver(
        &mut self,
        local_path: &str,
        remote_path: Option<&str>,
        connection: Map<String, Value>,
    ) -> Result<String> {
        let file = DeliveryFile {
            local_path: local_path.to_string(),
            remote_path: remote_path.map(str::to_string),
        };
        self.deliver_files(&[file], connection, Map::new())
    }

    /// Deliver `files` in one request with the destination entry's `options`. More than one file
    /// needs a plugin advertising `multi_file`; call once per file otherwise.
    pub fn deliver_files(
        &mut self,
        files: &[DeliveryFile],
        connection: Map<String, Value>,
        options: Map<String, Value>,
    ) -> Result<String> {
        let req = match files {
            [] => {
                return Err(HostError::Malformed {
                    plugin: self.label.clone(),
                    message: "nothing to deliver".into(),
                });
            }
            [one] => Request::Deliver {
                local_path: Some(one.local_path.clone()),
                remote_path: one.remote_path.clone(),
                files: Vec::new(),
                connection,
                options,
            },
            many => {
                if !self.has(crate::CAP_MULTI_FILE) {
                    return Err(HostError::Plugin {
                        plugin: self.label.clone(),
                        message: format!(
                            "{} files in one delivery, but the plugin doesn't advertise `multi_file`",
                            many.len()
                        ),
                    });
                }
                Request::Deliver {
                    local_path: None,
                    remote_path: None,
                    files: many.to_vec(),
                    connection,
                    options,
                }
            }
        };
        self.send(&req)?;
        match self.recv_json("a delivered reply")? {
            Response::Delivered { location } => Ok(location),
            other => Err(self.unexpected("a delivered reply", &Incoming::Json(other))),
        }
    }

    /// Ask the plugin to exit, and wait for it.
    pub fn close(mut self) -> Result<()> {
        let _ = self.send(&Request::Close {});
        let _ = self.recv(Some(Duration::from_secs(5)), "the close reply");
        self.stdin.take();
        for _ in 0..250 {
            if let Ok(Some(_)) = self.child.try_wait() {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        let _ = self.child.kill();
        Ok(())
    }

    /// Raw access for protocol tests: write arbitrary bytes to the plugin.
    pub fn write_raw(&mut self, bytes: &[u8]) -> Result<()> {
        let ok = match self.stdin.as_mut() {
            Some(w) => w.write_all(bytes).and_then(|_| w.flush()).is_ok(),
            None => false,
        };
        if ok { Ok(()) } else { Err(self.write_failed()) }
    }

    /// Wait up to `timeout` for the process to exit.
    pub fn wait_exit(&mut self, timeout: Duration) -> Option<ExitStatus> {
        self.stdin.take();
        let deadline = std::time::Instant::now() + timeout;
        while std::time::Instant::now() < deadline {
            if let Ok(Some(s)) = self.child.try_wait() {
                return Some(s);
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        None
    }
}

impl PluginProcess {
    /// Wait (briefly) for the stderr forwarder to reach end of stream.
    fn drain_stderr(&mut self) {
        if let Some(t) = self.stderr_thread.take() {
            let deadline = std::time::Instant::now() + Duration::from_secs(2);
            while !t.is_finished() && std::time::Instant::now() < deadline {
                std::thread::sleep(Duration::from_millis(5));
            }
            if t.is_finished() {
                let _ = t.join();
            }
        }
    }
}

impl Drop for PluginProcess {
    fn drop(&mut self) {
        self.stdin.take();
        if let Ok(None) = self.child.try_wait() {
            // Give it a moment to exit on its own (stdin is closed), then insist.
            for _ in 0..25 {
                if let Ok(Some(_)) = self.child.try_wait() {
                    break;
                }
                std::thread::sleep(Duration::from_millis(10));
            }
            let _ = self.child.kill();
            let _ = self.child.wait();
        }
        self.drain_stderr();
    }
}

static CORE_VERSION: OnceLock<String> = OnceLock::new();

/// Set the DRE version sent to plugins in the handshake (`core_version`). This crate is versioned
/// on its own, so the host names the version of DRE it's part of; unset, it's "unreleased".
pub fn set_core_version(version: &str) {
    let _ = CORE_VERSION.set(version.to_string());
}

fn core_version() -> String {
    CORE_VERSION.get().cloned().unwrap_or_else(|| "unreleased".into())
}