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. A plugin's `log` and
//! `progress` messages (protocol 1) go to the same sink.

use std::collections::VecDeque;
use std::io::{BufRead, BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStdin, Command, ExitStatus, Stdio};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
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, Envelope, LogLevel, Request, Response, ResultSetMeta};
use crate::options::OptionField;
use crate::{
    CORE_MIN_VERSION, Kind, MAX_VERSION, PluginId, parse_executable_name, parse_package_executable_name,
};

/// Receives each line a plugin logs: `(plugin, line)`. A `log` message arrives prefixed by its
/// level (`info: `, `warning: ` for warn and error, `debug: ` for debug and trace), as does
/// `progress` (`info: `); a stderr line arrives as written.
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}"))
}

/// Waits for a plugin's stderr to be logged. Lines reach the [`LogSink`] from a background thread,
/// so a line the plugin wrote before answering a request can be logged after core has the answer.
/// A host that logs a request's lines differently (hiding them while the request carries a secret)
/// syncs before it changes how it logs.
#[derive(Clone)]
pub struct StderrSync {
    progress: Arc<StderrProgress>,
    pipe: RawPipe,
}

impl StderrSync {
    /// Wait up to `timeout` until every complete line the plugin has written to stderr so far has
    /// been logged. False when the plugin kept writing (or the pipe can't be inspected) until then.
    pub fn wait(&self, timeout: Duration) -> bool {
        let deadline = std::time::Instant::now() + timeout;
        loop {
            let p = &self.progress;
            if p.done.load(Ordering::SeqCst) {
                return true;
            }
            // Blocked on an empty pipe, and still blocked in the same read after looking: every
            // line read before that read has been logged, and nothing more is waiting.
            let reads = p.reads.load(Ordering::SeqCst);
            if p.reading.load(Ordering::SeqCst)
                && pending(self.pipe) == Some(0)
                && p.reading.load(Ordering::SeqCst)
                && p.reads.load(Ordering::SeqCst) == reads
            {
                return true;
            }
            if std::time::Instant::now() >= deadline {
                return false;
            }
            std::thread::sleep(Duration::from_millis(1));
        }
    }
}

/// What the stderr forwarder is doing, for [`StderrSync`].
#[derive(Default)]
struct StderrProgress {
    /// In a read of the pipe: every complete line read before it has been logged.
    reading: AtomicBool,
    /// Reads that returned.
    reads: AtomicU64,
    /// At end of stream: everything has been logged.
    done: AtomicBool,
}

/// The pipe, read through [`StderrProgress`] so [`StderrSync`] knows where the forwarder is.
struct TrackedPipe<R> {
    pipe: R,
    progress: Arc<StderrProgress>,
}

impl<R: Read> Read for TrackedPipe<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.progress.reading.store(true, Ordering::SeqCst);
        let n = self.pipe.read(buf);
        self.progress.reading.store(false, Ordering::SeqCst);
        self.progress.reads.fetch_add(1, Ordering::SeqCst);
        n
    }
}

#[cfg(unix)]
type RawPipe = std::os::fd::RawFd;
#[cfg(windows)]
type RawPipe = usize;
#[cfg(not(any(unix, windows)))]
type RawPipe = ();

#[cfg(unix)]
fn raw_pipe(pipe: &std::process::ChildStderr) -> RawPipe {
    std::os::fd::AsRawFd::as_raw_fd(pipe)
}
#[cfg(windows)]
fn raw_pipe(pipe: &std::process::ChildStderr) -> RawPipe {
    std::os::windows::io::AsRawHandle::as_raw_handle(pipe) as usize
}
#[cfg(not(any(unix, windows)))]
fn raw_pipe(_: &std::process::ChildStderr) -> RawPipe {}

/// Bytes waiting in the pipe, unread.
#[cfg(unix)]
fn pending(fd: RawPipe) -> Option<usize> {
    let mut n: libc::c_int = 0;
    // SAFETY: FIONREAD writes one c_int; the forwarder keeps the descriptor open until it is done.
    let r = unsafe { libc::ioctl(fd, libc::FIONREAD, &mut n) };
    (r == 0).then_some(n.max(0) as usize)
}
#[cfg(windows)]
fn pending(handle: RawPipe) -> Option<usize> {
    let mut n: u32 = 0;
    // SAFETY: peeks without reading; the forwarder keeps the handle open until it is done.
    let ok = unsafe {
        windows_sys::Win32::System::Pipes::PeekNamedPipe(
            handle as windows_sys::Win32::Foundation::HANDLE,
            std::ptr::null_mut(),
            0,
            std::ptr::null_mut(),
            &mut n,
            std::ptr::null_mut(),
        )
    };
    (ok != 0).then_some(n as usize)
}
#[cfg(not(any(unix, windows)))]
fn pending(_: RawPipe) -> Option<usize> {
    None
}

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>,
}

/// A plugin's `describe` reply.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct Description {
    pub connection_fields: Vec<ConnectionField>,
    pub option_fields: Vec<OptionField>,
    /// A source's identifier quote character.
    pub identifier_quote: Option<String>,
    /// How a source wants string literals spelled (`chr` when absent; `backslash`).
    pub string_literal: Option<String>,
    /// A message destination's length limit, in characters.
    pub message_limit: Option<u64>,
}

#[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, with its kind and code when it gave them.
    Plugin {
        plugin: String,
        message: String,
        kind: Option<String>,
        code: Option<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 {}..={}; ",
                    plugin_range.0, plugin_range.1, core.0, core.1,
                )?;
                if plugin_range.1 >= core.0 {
                    return f.write_str("update DRE");
                }
                let package = parse_package_executable_name(plugin)
                    .or_else(|| parse_executable_name(plugin).map(|(_, n)| n));
                match package {
                    Some(p) => write!(f, "update the plugin with `dre plugin update {p}`"),
                    None => f.write_str("update the plugin"),
                }
            }
            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 {}

impl HostError {
    fn plugin(plugin: &str, message: impl Into<String>) -> HostError {
        HostError::Plugin {
            plugin: plugin.to_string(),
            message: message.into(),
            kind: None,
            code: None,
        }
    }
}

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

/// Asks a running request to stop, from any thread (see [`PluginProcess::canceller`]), and
/// can kill the plugin.
#[derive(Clone)]
pub struct Canceller {
    stdin: Arc<Mutex<Option<BufWriter<ChildStdin>>>>,
    current: Arc<AtomicU64>,
    child: Arc<Mutex<Child>>,
}

impl Canceller {
    /// Send `cancel` for the request running now, if any. Returns whether one was sent.
    pub fn cancel(&self) -> bool {
        let id = self.current.load(Ordering::SeqCst);
        if id == 0 {
            return false;
        }
        // A poisoned lock means a writer panicked mid-frame: nothing more can be sent.
        self.stdin.lock().is_ok_and(|mut stdin| {
            stdin
                .as_mut()
                .is_some_and(|w| frame::write_json(w, &Envelope::new(Some(id), Request::Cancel {})).is_ok())
        })
    }

    /// Kill the plugin process.
    pub fn kill(&self) {
        if let Ok(mut c) = self.child.lock() {
            let _ = c.kill();
        }
    }
}

/// Every plugin process this program has running, by a number of its own, so a cancelled run
/// can reach all of them wherever they were started.
static RUNNING: Mutex<Vec<(u64, Canceller)>> = Mutex::new(Vec::new());
static NEXT_PROCESS: AtomicU64 = AtomicU64::new(1);

/// Send `cancel` to every running plugin's current request (protocol 1). Returns how many had one.
///
/// # Panics
///
/// Panics if the list of running plugins is poisoned (a thread panicked while changing it).
pub fn cancel_all() -> usize {
    let all: Vec<Canceller> = RUNNING.lock().unwrap().iter().map(|(_, c)| c.clone()).collect();
    all.iter().filter(|c| c.cancel()).count()
}

/// Kill every running plugin process.
///
/// # Panics
///
/// Panics if the list of running plugins is poisoned (a thread panicked while changing it).
pub fn kill_all() {
    let all: Vec<Canceller> = RUNNING.lock().unwrap().iter().map(|(_, c)| c.clone()).collect();
    for c in all {
        c.kill();
    }
}

/// How many plugin processes are running.
///
/// # Panics
///
/// Panics if the list of running plugins is poisoned (a thread panicked while changing it).
pub fn running() -> usize {
    RUNNING.lock().unwrap().len()
}
/// 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,
    /// Shared with [`Canceller`]s, which can kill it.
    child: Arc<Mutex<Child>>,
    /// This process's entry in [`RUNNING`].
    registered: u64,
    /// Shared with [`Canceller`]s; `None` once closed.
    stdin: Arc<Mutex<Option<BufWriter<ChildStdin>>>>,
    rx: Receiver<std::result::Result<Frame, FrameError>>,
    log: LogSink,
    /// The last request id given out.
    last_id: u64,
    /// The request running now (0: none).
    current: Arc<AtomicU64>,
    stderr: Arc<Mutex<VecDeque<String>>>,
    stderr_sync: StderrSync,
    /// 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>>,
    /// How many tries the last delivery took (1 without a retry).
    last_attempts: u32,
    /// 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((CORE_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.
    ///
    /// # Panics
    ///
    /// Panics if the list of running plugins is poisoned (a thread panicked while changing it).
    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" });
        }
        // Its own process group: Ctrl-C at the terminal reaches core, which asks the plugin to
        // stop with `cancel`, instead of killing the plugin mid-request.
        #[cfg(unix)]
        {
            use std::os::unix::process::CommandExt;
            cmd.process_group(0);
        }
        // Files on a FUSE mount (a Databricks Volume as the target path) close before `execve`,
        // or the plugin may never start (see `fuse_fds`).
        #[cfg(target_os = "linux")]
        crate::fuse_fds::close_in_child(&mut cmd);
        #[cfg(windows)]
        {
            use std::os::windows::process::CommandExt;
            const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200;
            cmd.creation_flags(CREATE_NEW_PROCESS_GROUP);
        }
        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 (Some(stdout), Some(stderr_pipe)) = (child.stdout.take(), child.stderr.take()) else {
            let _ = child.kill();
            return Err(HostError::Spawn {
                plugin: label,
                error: std::io::Error::other("the plugin's stdout or stderr isn't piped"),
            });
        };
        let stdin = Arc::new(Mutex::new(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 sink = log.clone();
        let stderr_sync = StderrSync {
            progress: Arc::default(),
            pipe: raw_pipe(&stderr_pipe),
        };
        let progress = stderr_sync.progress.clone();
        let stderr_thread = std::thread::spawn(move || {
            let mut reader = BufReader::new(TrackedPipe {
                pipe: stderr_pipe,
                progress: progress.clone(),
            });
            for line in reader.by_ref().lines() {
                let Ok(line) = line else { break };
                sink(&who, &line);
                let mut t = tail.lock().unwrap();
                if t.len() == STDERR_TAIL {
                    t.pop_front();
                }
                t.push_back(line);
            }
            // Before the pipe closes, so a sync never inspects a closed descriptor.
            progress.done.store(true, Ordering::SeqCst);
            drop(reader);
        });
        let serve = parse_executable_name(&label).map(|(k, n)| PluginId::new(k, n));
        let child = Arc::new(Mutex::new(child));
        let current = Arc::new(AtomicU64::new(0));
        let registered = NEXT_PROCESS.fetch_add(1, Ordering::SeqCst);
        RUNNING.lock().unwrap().push((
            registered,
            Canceller {
                stdin: stdin.clone(),
                current: current.clone(),
                child: child.clone(),
            },
        ));
        Ok(PluginProcess {
            label,
            path: path.to_path_buf(),
            child,
            registered,
            stdin,
            rx,
            log,
            last_id: 0,
            current,
            stderr,
            stderr_sync,
            stderr_thread: Some(stderr_thread),
            info: None,
            early: None,
            last_attempts: 1,
            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(
                        &self.label,
                        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, kind, code }) => Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
                kind,
                code,
            }),
            other => Err(self.unexpected("a hello reply", &other)),
        }
    }

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

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

    /// What the plugin said in its handshake.
    ///
    /// # Panics
    ///
    /// Panics if the process was spawned without a handshake ([`Self::spawn_env`],
    /// [`Self::spawn_in`]).
    pub fn info(&self) -> &PluginInfo {
        self.info.as_ref().expect("handshake completed")
    }

    /// The process's exit status if it has exited (the lock is held only for the call).
    fn try_wait(&self) -> std::io::Result<Option<ExitStatus>> {
        self.child.lock().unwrap().try_wait()
    }

    /// A handle that can cancel the running request from another thread. Only a plugin on
    /// protocol 1 or later understands `cancel`.
    pub fn canceller(&self) -> Canceller {
        Canceller {
            stdin: self.stdin.clone(),
            current: self.current.clone(),
            child: self.child.clone(),
        }
    }

    /// A handle that waits for the plugin's stderr to be logged; see [`StderrSync`].
    pub fn stderr_sync(&self) -> StderrSync {
        self.stderr_sync.clone()
    }

    /// The protocol version settled on in the handshake (0 before it).
    fn version(&self) -> u32 {
        self.info.as_ref().map_or(0, |i| i.protocol_version)
    }

    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.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,
        }
    }

    /// Send a request. From protocol 1 it carries an id: a new one, or for the messages that
    /// continue a request (`result_set_end`, `finish`, `cancel`) the running request's.
    pub fn send(&mut self, req: &Request) -> Result<()> {
        let id = match req {
            _ if self.version() == 0 => None,
            Request::Hello { .. } => None,
            Request::ResultSetEnd {} | Request::Finish {} | Request::Cancel {} => {
                Some(self.current.load(Ordering::SeqCst)).filter(|id| *id != 0)
            }
            _ => {
                self.last_id += 1;
                self.current.store(self.last_id, Ordering::SeqCst);
                Some(self.last_id)
            }
        };
        // A poisoned lock means a writer panicked mid-frame: report a failed write.
        let ok = self.stdin.lock().is_ok_and(|mut stdin| {
            stdin
                .as_mut()
                .is_some_and(|w| frame::write_json(w, &Envelope::new(id, req)).is_ok())
        });
        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, kind, code }) = serde_json::from_value::<Response>(v.clone())
        {
            return Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
                kind,
                code,
            });
        }
        let ipc = frame::encode_batch(batch).map_err(|e| HostError::Arrow {
            plugin: self.label.clone(),
            message: e.to_string(),
        })?;
        let ok = self.stdin.lock().is_ok_and(|mut stdin| {
            stdin
                .as_mut()
                .is_some_and(|w| frame::write_arrow(w, &ipc).is_ok())
        });
        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, kind, code }) => HostError::Plugin {
                    plugin: self.label.clone(),
                    message,
                    kind,
                    code,
                },
                _ => self.crashed(),
            },
            _ => self.crashed(),
        }
    }

    /// Receive the next message. `timeout` of `None` waits as long as the plugin is alive. `log`
    /// and `progress` messages go to the log sink on the way.
    pub fn recv(&mut self, timeout: Option<Duration>, waiting_for: &'static str) -> Result<Incoming> {
        loop {
            match self.recv_any(timeout, waiting_for)? {
                Incoming::Json(Response::Log {
                    level,
                    message,
                    fields,
                }) => {
                    (self.log)(&self.label, &log_line(level, &message, &fields));
                }
                Incoming::Json(Response::Progress { message, done, total }) => {
                    (self.log)(&self.label, &progress_line(message.as_deref(), done, total));
                }
                other => return Ok(other),
            }
        }
    }

    /// The next frame's JSON as sent, with nothing skipped or checked (conformance tests).
    pub fn recv_raw(&mut self, timeout: Duration) -> Result<Value> {
        match self.rx.recv_timeout(timeout) {
            Ok(Ok(Frame::Json(v))) => Ok(v),
            Ok(Ok(Frame::Arrow(_))) => Err(HostError::Malformed {
                plugin: self.label.clone(),
                message: "Arrow data where a JSON message was expected".into(),
            }),
            Ok(Err(FrameError::Malformed(m))) => Err(HostError::Malformed {
                plugin: self.label.clone(),
                message: m,
            }),
            Err(RecvTimeoutError::Timeout) => Err(HostError::Timeout {
                plugin: self.label.clone(),
                waiting_for: "a message",
            }),
            _ => Err(self.crashed()),
        }
    }

    fn recv_any(&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::<Envelope<Response>>(v.clone()) {
                Ok(Envelope { id: Some(id), body })
                    if !matches!(body, Response::Log { .. } | Response::Progress { .. }) =>
                {
                    let current = self.current.load(Ordering::SeqCst);
                    if id != current {
                        return Err(HostError::Malformed {
                            plugin: self.label.clone(),
                            message: format!(
                                "a reply to request {id} arrived while request {current} was running"
                            ),
                        });
                    }
                    Ok(Incoming::Json(body))
                }
                Ok(e) => Ok(Incoming::Json(e.body)),
                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, kind, code }) => Err(HostError::Plugin {
                plugin: self.label.clone(),
                message,
                kind,
                code,
            }),
            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>)> {
        let d = self.description()?;
        Ok((d.connection_fields, d.option_fields))
    }

    /// The whole `describe` reply.
    pub fn description(&mut self) -> Result<Description> {
        self.send(&Request::Describe {})?;
        match self.recv_json("a describe reply")? {
            Response::Describe {
                connection_fields,
                option_fields,
                identifier_quote,
                string_literal,
                message_limit,
            } => Ok(Description {
                connection_fields,
                option_fields,
                identifier_quote,
                string_literal,
                message_limit,
            }),
            other => Err(self.unexpected("a describe reply", &Incoming::Json(other))),
        }
    }

    /// Check a profile entry's connection settings without connecting (protocol 1): the errors
    /// and the warnings. `unresolved` keys (an unset `env_var()`) are sent as `null`.
    pub fn validate_connection(
        &mut self,
        mut connection: Map<String, Value>,
        unresolved: Vec<String>,
    ) -> Result<(Vec<String>, Vec<String>)> {
        for k in &unresolved {
            connection.insert(k.clone(), Value::Null);
        }
        self.send(&Request::ValidateConnection {
            connection,
            unresolved,
        })?;
        match self.recv_json("a validated reply")? {
            Response::Validated { errors, warnings } => Ok((errors, warnings)),
            other => Err(self.unexpected("a validated 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>> {
        Ok(self.validate_with(options, None)?.0)
    }

    /// Check a format's config block together with its `template` (as `write` would send it);
    /// returns the errors, then the warnings.
    pub fn validate_with(
        &mut self,
        options: Map<String, Value>,
        template: Option<Value>,
    ) -> Result<(Vec<String>, Vec<String>)> {
        self.send(&Request::Validate { options, template })?;
        match self.recv_json("a validate reply")? {
            Response::Validated { errors, warnings } => Ok((errors, warnings)),
            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,
                sandbox: false,
            },
            "an open reply",
        )
    }

    /// Open a read-only session that may not reach files, programs or the network beyond the
    /// database, where the plugin supports it (`dre query`).
    pub fn open_sandboxed(&mut self, connection: Map<String, Value>) -> Result<()> {
        self.call_ok(
            &Request::Open {
                connection,
                read_only: true,
                sandbox: true,
            },
            "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, kind, code }) => {
                    return Err(HostError::Plugin {
                        plugin: self.label.clone(),
                        message,
                        kind,
                        code,
                    });
                }
                other => return Err(self.unexpected("result data", &other)),
            }
        }
        if let Some(e) = sink_error {
            return Err(HostError::plugin(&self.label, 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,
                message: None,
            },
            many => {
                if !self.has(crate::CAP_MULTI_FILE) {
                    return Err(HostError::plugin(
                        &self.label,
                        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,
                    message: None,
                }
            }
        };
        self.send(&req)?;
        match self.recv_json("a delivered reply")? {
            Response::Delivered { location, attempts } => {
                self.last_attempts = attempts.unwrap_or(1);
                Ok(location)
            }
            other => Err(self.unexpected("a delivered reply", &Incoming::Json(other))),
        }
    }

    /// Deliver a message, with `attach` files, to a plugin advertising `message`.
    pub fn deliver_message(
        &mut self,
        message: &crate::msg::Message,
        attach: &[DeliveryFile],
        connection: Map<String, Value>,
        options: Map<String, Value>,
    ) -> Result<String> {
        if !self.has(crate::CAP_MESSAGE) {
            return Err(HostError::plugin(
                &self.label,
                "a message, but the plugin doesn't advertise `message`",
            ));
        }
        self.send(&Request::Deliver {
            local_path: None,
            remote_path: None,
            files: attach.to_vec(),
            connection,
            options,
            message: Some(message.clone()),
        })?;
        match self.recv_json("a delivered reply")? {
            Response::Delivered { location, attempts } => {
                self.last_attempts = attempts.unwrap_or(1);
                Ok(location)
            }
            other => Err(self.unexpected("a delivered reply", &Incoming::Json(other))),
        }
    }

    /// How many tries the last delivery took: 1, or more when the plugin retried a temporary
    /// error.
    pub fn last_attempts(&self) -> u32 {
        self.last_attempts
    }

    /// 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");
        if let Ok(mut stdin) = self.stdin.lock() {
            stdin.take();
        }
        for _ in 0..250 {
            if let Ok(Some(_)) = self.try_wait() {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        if let Ok(mut child) = self.child.lock() {
            let _ = 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 = self.stdin.lock().is_ok_and(|mut stdin| {
            stdin
                .as_mut()
                .is_some_and(|w| w.write_all(bytes).and_then(|_| w.flush()).is_ok())
        });
        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> {
        if let Ok(mut stdin) = self.stdin.lock() {
            stdin.take();
        }
        let deadline = std::time::Instant::now() + timeout;
        while std::time::Instant::now() < deadline {
            if let Ok(Some(s)) = self.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) {
        if let Ok(mut r) = RUNNING.lock() {
            r.retain(|(id, _)| *id != self.registered);
        }
        if let Ok(mut stdin) = self.stdin.lock() {
            stdin.take();
        }
        if let Ok(None) = self.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.try_wait() {
                    break;
                }
                std::thread::sleep(Duration::from_millis(10));
            }
            let _ = self.child.lock().unwrap().kill();
            let _ = self.child.lock().unwrap().wait();
        }
        self.drain_stderr();
    }
}

/// A `log` message as a log sink line: prefixed by its level, then its fields as `key=value`.
fn log_line(level: LogLevel, message: &str, fields: &Map<String, Value>) -> String {
    let prefix = match level {
        LogLevel::Error | LogLevel::Warn => "warning: ",
        LogLevel::Info => "info: ",
        LogLevel::Debug | LogLevel::Trace => "debug: ",
    };
    let mut line = format!("{prefix}{message}");
    for (k, v) in fields {
        match v {
            Value::String(s) => line.push_str(&format!(" {k}={s}")),
            v => line.push_str(&format!(" {k}={v}")),
        }
    }
    line
}

/// A `progress` message as a log sink line.
fn progress_line(message: Option<&str>, done: Option<u64>, total: Option<u64>) -> String {
    let count = match (done, total) {
        (Some(d), Some(t)) => format!("{d}/{t}"),
        (Some(d), None) => d.to_string(),
        _ => String::new(),
    };
    match (message, count.is_empty()) {
        (Some(m), true) => format!("info: {m}"),
        (Some(m), false) => format!("info: {m} ({count})"),
        (None, _) => format!("info: {count}"),
    }
}

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())
}