sausage/sdk/src/transport.rs

//! Stdin routing, stdout framing, and process lifetime.
#[cfg(target_os = "wasi")]
use crate::lines::Lines;
use crate::*;
use futures::{channel::oneshot, future::LocalBoxFuture};
use serde_json::{Value, json};
use std::{
    cell::RefCell,
    collections::HashMap,
    future::{Future, IntoFuture},
    io::Write,
    path::Path,
};
#[path = "calls.rs"]
mod calls;
#[path = "locks.rs"]
mod locks;
pub use calls::*;
pub use locks::*;
/// Generated main-interface dispatch.
pub trait Interface: 'static {
    /// Describe this interface without constructing its value.
    fn describe() -> Value
    where
        Self: Sized;
    /// Add this type's methods to shared document generators.
    #[doc(hidden)]
    fn add_methods(doc: &mut crate::Discovery)
    where
        Self: Sized;
    /// Describe only this type's methods, without endpoint declarations.
    fn methods_document() -> Value
    where
        Self: Sized;
    /// List the declared endpoint types and their names.
    fn endpoint_types(&self) -> Vec<(std::any::TypeId, String)> {
        vec![]
    }
    /// Describe this interface as OpenRPC.
    fn discover(&self) -> Value;
    /// Dispatch a request to a generated handler.
    fn dispatch<'a>(
        &'a self,
        method: &'a str,
        params: Value,
        ctx: &'a Ctx,
    ) -> LocalBoxFuture<'a, Result<Value>>;
    /// Run the EOF hook.
    fn stop(&self) -> LocalBoxFuture<'_, ()> {
        Box::pin(async {})
    }
}

/// Context of an incoming call.
#[derive(Default)]
pub struct Ctx {
    pub(crate) exit: std::cell::Cell<bool>,
    pub(crate) sink: Option<futures::channel::mpsc::UnboundedSender<Value>>,
}
impl Ctx {
    /// Send progress to the caller of this request.
    pub fn notify(&self, name: &str, params: impl serde::Serialize) -> Result<()> {
        let params = json!({"name":name,"params":params});
        if let Some(sink) = &self.sink {
            sink.unbounded_send(json!({"jsonrpc":"2.0","method":"core.notify","params":params}))
                .map_err(|_| Error::new(-32000, "endpoint closed"))
        } else {
            notify("core.notify", params)
        }
    }
    /// End the process after this request receives its response.
    pub fn exit_after_response(&self) {
        self.exit.set(true);
    }
}

type Notify = Box<dyn FnMut(&str, Value)>;
struct Pending {
    tx: oneshot::Sender<Result<Value>>,
    notify: Option<Notify>,
    unlock: Option<std::rc::Rc<Lease>>,
}
struct Lease {
    path: String,
    held: std::cell::Cell<bool>,
}
impl Drop for Lease {
    fn drop(&mut self) {
        if self.held.get() {
            let _ = notify("core.unlock", json!({"path":self.path}));
        }
    }
}
#[derive(Default)]
struct State {
    active: bool,
    endpoints: HashMap<std::any::TypeId, String>,
    stopping: bool,
    draining: bool,
    next: u64,
    pending: HashMap<u64, Pending>,
    // Only leases, never callbacks or response senders, survive dropped lock acquisitions.
    abandoned_locks: HashMap<u64, std::rc::Rc<Lease>>,
    #[cfg(target_os = "wasi")]
    tasks: usize,
    #[cfg(target_os = "wasi")]
    wake: Option<std::task::Waker>,
    #[cfg(target_os = "wasi")]
    exit: bool,
    #[cfg(target_os = "wasi")]
    listeners: HashMap<u64, futures::future::AbortHandle>,
}
thread_local! { static STATE: RefCell<State> = RefCell::new(State::default()); }
#[cfg(any(target_os = "wasi", test))]
fn clear_pending() {
    let pending = STATE.with(|s| {
        let mut s = s.borrow_mut();
        (
            std::mem::take(&mut s.pending),
            std::mem::take(&mut s.abandoned_locks),
        )
    });
    // Callback captures may own other call futures whose drop guards use STATE.
    drop(pending);
}
pub(crate) fn endpoint_name<T: 'static>() -> Result<String> {
    STATE
        .with(|s| {
            s.borrow()
                .endpoints
                .get(&std::any::TypeId::of::<T>())
                .cloned()
        })
        .ok_or_else(|| {
            Error::new(
                -32602,
                "endpoint type is not declared by the main interface",
            )
        })
}
#[cfg(target_os = "wasi")]
pub(crate) fn endpoint_task(future: impl Future<Output = ()> + 'static) {
    background(async move {
        let _ = until_drained(future).await;
    });
}
/// Cancel an accept/read wait when draining, but never an active handler.
#[cfg(target_os = "wasi")]
pub(crate) async fn until_drained<F: Future>(future: F) -> Option<F::Output> {
    let (handle, registration) = futures::future::AbortHandle::new_pair();
    let id = STATE.with(|s| {
        let mut s = s.borrow_mut();
        if s.draining {
            return None;
        }
        s.next += 1;
        let id = s.next;
        s.listeners.insert(id, handle);
        Some(id)
    })?;
    let result = futures::future::Abortable::new(future, registration)
        .await
        .ok();
    STATE.with(|s| s.borrow_mut().listeners.remove(&id));
    result
}
/// Stop accepting SDK endpoint work, then exit when handlers and background tasks finish.
///
/// Calls and notifications still work. This does not run the EOF stop hook or cancel
/// background tasks: services with their own listeners/tasks must arrange for those
/// to finish too. Repeated calls are harmless. Core's upgrade grace remains the bound.
pub fn drain() {
    STATE.with(|s| {
        let mut s = s.borrow_mut();
        s.draining = true;
        #[cfg(target_os = "wasi")]
        {
            for (_, listener) in s.listeners.drain() {
                listener.abort();
            }
            if let Some(w) = s.wake.take() {
                w.wake();
            }
        }
    });
}
#[cfg(target_os = "wasi")]
pub(crate) fn exit() {
    STATE.with(|s| {
        let mut s = s.borrow_mut();
        s.exit = true;
        if let Some(w) = s.wake.take() {
            w.wake();
        }
    });
}
fn send(value: Value) -> Result<()> {
    let bytes = format!("\n{value}\n").into_bytes();
    let mut stdout = std::io::stdout().lock();
    stdout.write_all(&bytes)?;
    stdout.flush()?;
    Ok(())
}
fn notify(method: &str, params: Value) -> Result<()> {
    if !STATE.with(|s| s.borrow().active) {
        return Err(Error::new(-32000, "no core transport"));
    }
    send(json!({"jsonrpc":"2.0","method":method,"params":params}))
}

#[cfg(target_os = "wasi")]
/// Execute a future with the SDK runtime.
pub fn block_on<F: Future>(f: F) -> F::Output {
    wstd::runtime::block_on(f)
}
#[cfg(not(target_os = "wasi"))]
/// Execute a future with the SDK runtime.
pub fn block_on<F: Future>(f: F) -> F::Output {
    futures::executor::block_on(f)
}

/// Run an SDK test with a core transport when the WASI runner supplies one.
#[cfg(target_os = "wasi")]
/// Execute an async test with the available core transport.
pub fn test_run<F: Future>(f: F) -> F::Output {
    block_on(async {
        STATE.with(|s| s.borrow_mut().active = true);
        let reader = wstd::runtime::spawn(async {
            let mut input = wstd::io::stdin();
            let mut lines = Lines::default();
            while let Ok(Some(line)) = lines.next(&mut input).await {
                if let Ok(v) = serde_json::from_slice::<Value>(&line)
                    && v["jsonrpc"] == "2.0"
                {
                    receive(v);
                }
            }
            STATE.with(|s| {
                let mut s = s.borrow_mut();
                s.stopping = true;
            });
            clear_pending();
        });
        let result = f.await;
        drop(reader);
        STATE.with(|s| {
            let mut s = s.borrow_mut();
            s.active = false;
        });
        clear_pending();
        result
    })
}
#[cfg(not(target_os = "wasi"))]
/// Execute an async test with the available core transport.
pub fn test_run<F: Future>(f: F) -> F::Output {
    block_on(f)
}

/// Register a task that keeps the process alive.
#[cfg(target_os = "wasi")]
/// Keep this process alive while a future runs.
pub fn background(f: impl Future<Output = ()> + 'static) {
    STATE.with(|s| s.borrow_mut().tasks += 1);
    wstd::runtime::spawn(async move {
        f.await;
        STATE.with(|s| {
            let mut s = s.borrow_mut();
            s.tasks -= 1;
            if let Some(w) = s.wake.take() {
                w.wake();
            }
        });
    })
    .detach();
}
#[cfg(not(target_os = "wasi"))]
/// Keep this process alive while a future runs.
pub fn background(_f: impl Future<Output = ()> + 'static) {
    panic!("background requires the WASI runtime");
}

#[doc(hidden)]
pub fn receive(value: Value) {
    if value["method"] == "core.notify" {
        if let Some(id) = value["params"]["call"].as_u64() {
            let callback = STATE.with(|s| {
                s.borrow_mut()
                    .pending
                    .get_mut(&id)
                    .and_then(|p| p.notify.take())
            });
            if let Some(mut f) = callback {
                f(
                    value["params"]["name"].as_str().unwrap_or_default(),
                    value["params"]["params"].clone(),
                );
                STATE.with(|s| {
                    if let Some(p) = s.borrow_mut().pending.get_mut(&id) {
                        p.notify = Some(f);
                    }
                });
            }
        }
        return;
    }
    let pending = value["id"]
        .as_u64()
        .and_then(|id| STATE.with(|s| s.borrow_mut().pending.remove(&id)));
    if let Some(p) = pending {
        let result = if value.get("error").is_some() {
            serde_json::from_value(value["error"].clone()).map_or_else(|e| Err(e.into()), Err)
        } else {
            Ok(value["result"].clone())
        };
        if let Some(lease) = &p.unlock {
            lease
                .held
                .set(result.as_ref().is_ok_and(|v| v["held"] == true));
        }
        let _ = p.tx.send(result);
    } else {
        let lease = value["id"]
            .as_u64()
            .and_then(|id| STATE.with(|s| s.borrow_mut().abandoned_locks.remove(&id)));
        if let Some(lease) = lease {
            lease
                .held
                .set(value.get("error").is_none() && value["result"]["held"] == true);
        } else {
            eprintln!("note: response with unknown ID {}", value["id"]);
        }
    }
}

/// Serve the main interface until its work ends or stdin reaches EOF.
#[cfg(target_os = "wasi")]
/// Serve the main interface over stdin and stdout.
pub fn run<T: Interface>(value: T) {
    block_on(serve(value));
}
#[cfg(not(target_os = "wasi"))]
/// Serve the main interface over stdin and stdout.
pub fn run<T: Interface>(_value: T) {
    panic!("run requires WASI stdin/stdout");
}

#[cfg(target_os = "wasi")]
async fn serve<T: Interface>(value: T) {
    use std::{rc::Rc, task::Poll};
    STATE.with(|s| {
        *s.borrow_mut() = State {
            active: true,
            endpoints: value.endpoint_types().into_iter().collect(),
            ..State::default()
        }
    });
    let value = Rc::new(value);
    let mut input = wstd::io::stdin();
    let mut lines = Lines::default();
    let mut started = false;
    loop {
        let read = lines.next(&mut input);
        let idle = futures::future::poll_fn(|cx| {
            STATE.with(|s| {
                let mut s = s.borrow_mut();
                s.wake = Some(cx.waker().clone());
                if started && (s.tasks == 0 || s.exit) {
                    Poll::Ready(())
                } else {
                    Poll::Pending
                }
            })
        });
        futures::pin_mut!(read, idle);
        let bytes = match futures::future::select(read, idle).await {
            futures::future::Either::Left((Ok(Some(b)), _)) => b,
            futures::future::Either::Left((_, _)) => {
                STATE.with(|s| {
                    let mut s = s.borrow_mut();
                    s.stopping = true;
                    for (_, listener) in s.listeners.drain() {
                        listener.abort();
                    }
                });
                clear_pending();
                value.stop().await;
                break;
            }
            futures::future::Either::Right(_) => break,
        };
        let Ok(msg) = serde_json::from_slice::<Value>(&bytes) else {
            continue;
        };
        if msg["jsonrpc"] != "2.0" {
            continue;
        }
        if let Some(method) = msg["method"].as_str() {
            if method == "core.notify" {
                receive(msg);
                continue;
            }
            started = true;
            let value = value.clone();
            let method = method.to_owned();
            background(async move {
                let ctx = Ctx::default();
                let result = if method == "rpc.discover" {
                    Ok(value.discover())
                } else {
                    value
                        .dispatch(
                            &method,
                            msg.get("params").cloned().unwrap_or(json!({})),
                            &ctx,
                        )
                        .await
                };
                if let Some(id) = msg.get("id") {
                    let response = match result {
                        Ok(v) => json!({"jsonrpc":"2.0","id":id,"result":v}),
                        Err(e) => json!({"jsonrpc":"2.0","id":id,"error":e}),
                    };
                    let _ = send(response);
                }
                if ctx.exit.get() {
                    STATE.with(|s| s.borrow_mut().exit = true);
                }
            });
        } else {
            receive(msg);
        }
    }
    STATE.with(|s| {
        let mut s = s.borrow_mut();
        s.active = false;
    });
    clear_pending();
    let _ = std::io::stdout().flush();
    std::process::exit(0);
}