sausage-plugin part 2: import! with typed clients (sausage_interface::client_rs), TCP endpoints with rpc.endpoint, examples
Part 1 fixes: complete stdout writes, buffered stdin, Rust type names in schemas, enum variant docs, -32603 for serde errors in handlers, crate docs and modules. End-to-end check of a generated consumer through a fake core.
This commit is contained in:
@@ -0,0 +1,379 @@
|
||||
//! 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;
|
||||
/// 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(-32002, "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,
|
||||
next: u64,
|
||||
pending: HashMap<u64, Pending>,
|
||||
#[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()); }
|
||||
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) {
|
||||
let (handle, registration) = futures::future::AbortHandle::new_pair();
|
||||
let id = STATE.with(|s| {
|
||||
let mut s = s.borrow_mut();
|
||||
s.next += 1;
|
||||
let id = s.next;
|
||||
s.listeners.insert(id, handle);
|
||||
id
|
||||
});
|
||||
background(async move {
|
||||
let _ = futures::future::Abortable::new(future, registration).await;
|
||||
STATE.with(|s| s.borrow_mut().listeners.remove(&id));
|
||||
});
|
||||
}
|
||||
#[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(-32002, "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;
|
||||
s.pending.clear();
|
||||
});
|
||||
});
|
||||
let result = f.await;
|
||||
drop(reader);
|
||||
STATE.with(|s| {
|
||||
let mut s = s.borrow_mut();
|
||||
s.active = false;
|
||||
s.pending.clear();
|
||||
});
|
||||
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 {
|
||||
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;
|
||||
s.pending.clear();
|
||||
for (_, listener) in s.listeners.drain() {
|
||||
listener.abort();
|
||||
}
|
||||
});
|
||||
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;
|
||||
s.pending.clear();
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user