329 lines
11 KiB
Rust
329 lines
11 KiB
Rust
//! Probe of SDK transport and lifecycle behavior.
|
|
use sausage_plugin::serde_json::{Value, json};
|
|
use sausage_plugin::{Ctx, Result, call, lock, rpc, schema};
|
|
struct Probe;
|
|
/// A virtual path to observe.
|
|
#[schema]
|
|
struct WatchPath {
|
|
/// Directory or file, not necessarily present yet.
|
|
path: String,
|
|
}
|
|
/// The values to add.
|
|
#[schema]
|
|
struct Add {
|
|
/// The first number.
|
|
a: i32,
|
|
/// The second number.
|
|
b: i32,
|
|
}
|
|
/// The retired version.
|
|
#[schema]
|
|
struct Retired {
|
|
/// The previous hash.
|
|
previous: String,
|
|
}
|
|
/// The replacement version.
|
|
#[schema]
|
|
struct Superseded {
|
|
/// The current hash.
|
|
current: String,
|
|
}
|
|
/// A large payload.
|
|
#[schema]
|
|
struct Echo {
|
|
/// The text to return unchanged.
|
|
text: String,
|
|
}
|
|
#[rpc(grace_seconds = 12,endpoints(control=Control))]
|
|
impl Probe {
|
|
/// Scan only after the watch is armed; keep the call open until process death.
|
|
async fn watch(&self, p: WatchPath) -> Result<()> {
|
|
let path = p.path.clone();
|
|
call("core.watch", json!({"path":p.path}))
|
|
.on_notify(move |name, _| {
|
|
if name == "armed" && path == "/self" {
|
|
std::fs::write("/self/owned", "visible via proc").unwrap();
|
|
}
|
|
let mut entries: Vec<_> = std::fs::read_dir(&path)
|
|
.into_iter()
|
|
.flatten()
|
|
.filter_map(std::result::Result::ok)
|
|
.map(|e| e.file_name().into_string().unwrap())
|
|
.collect();
|
|
entries.sort();
|
|
let text = std::fs::read_to_string(format!("{path}/before")).ok();
|
|
Ctx::default()
|
|
.notify(
|
|
name,
|
|
json!({"pid":sausage_plugin::pid(), "entries":entries, "text":text}),
|
|
)
|
|
.unwrap();
|
|
})
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Poll the watch in the handler, then return idle without a background task.
|
|
async fn watch_once(&self, p: WatchPath) -> Result<()> {
|
|
use std::{
|
|
cell::RefCell,
|
|
rc::Rc,
|
|
task::{Poll, Waker},
|
|
};
|
|
let state = Rc::new(RefCell::new((false, None::<Waker>)));
|
|
let received = state.clone();
|
|
let mut watch = std::pin::pin!(std::future::IntoFuture::into_future(
|
|
call("core.watch", json!({"path":p.path})).on_notify(move |name, _| {
|
|
if name == "armed" {
|
|
let mut state = received.borrow_mut();
|
|
state.0 = true;
|
|
if let Some(waker) = state.1.take() {
|
|
waker.wake();
|
|
}
|
|
}
|
|
})
|
|
));
|
|
std::future::poll_fn(|cx| {
|
|
if let Poll::Ready(result) = watch.as_mut().poll(cx) {
|
|
return Poll::Ready(result.map(|_| ()));
|
|
}
|
|
let mut state = state.borrow_mut();
|
|
state.1 = Some(cx.waker().clone());
|
|
if state.0 {
|
|
Poll::Ready(Ok(()))
|
|
} else {
|
|
Poll::Pending
|
|
}
|
|
})
|
|
.await
|
|
}
|
|
|
|
/// Drop one watch while another still delivers, in a process that remains alive.
|
|
async fn drop_watch(&self) -> Result<Vec<String>> {
|
|
use std::{cell::RefCell, future::IntoFuture, rc::Rc, task::Poll};
|
|
let events = Rc::new(RefCell::new(Vec::new()));
|
|
let captured = events.clone();
|
|
let mut dropped = call("core.watch", json!({"path":"/dropped"}))
|
|
.on_notify(move |_, _| captured.borrow_mut().push("dropped".to_owned()))
|
|
.into_future();
|
|
std::future::poll_fn(|cx| {
|
|
assert!(dropped.as_mut().poll(cx).is_pending());
|
|
Poll::Ready(())
|
|
})
|
|
.await;
|
|
drop(dropped);
|
|
let captured = events.clone();
|
|
let mut kept = call("core.watch", json!({"path":"/kept"}))
|
|
.on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
|
|
.into_future();
|
|
std::future::poll_fn(|cx| {
|
|
assert!(kept.as_mut().poll(cx).is_pending());
|
|
Poll::Ready(())
|
|
})
|
|
.await;
|
|
// The harness sends late traffic and kept-watch events before answering.
|
|
call("other.barrier", json!({})).await?;
|
|
drop(kept);
|
|
let result = events.borrow().clone();
|
|
Ok(result)
|
|
}
|
|
|
|
/// Publish a control endpoint until EOF.
|
|
async fn serve(&self) -> Result<()> {
|
|
sausage_plugin::listen("/endpoint", std::rc::Rc::new(Control)).await
|
|
}
|
|
/// Forward detached progress through the real core.
|
|
async fn relay(&self) -> Result<i32> {
|
|
call("probe.progress", json!({}))
|
|
.detach()
|
|
.on_notify(move |name, params| {
|
|
Ctx::default().notify(name, params).unwrap();
|
|
})
|
|
.await
|
|
.map(|value| value.as_i64().unwrap() as i32)
|
|
}
|
|
/// Send progress before the final result.
|
|
fn progress(&self, ctx: &Ctx) -> Result<i32> {
|
|
ctx.notify("step", json!({"done":true}))?;
|
|
Ok(42)
|
|
}
|
|
/// Start a detached service that survives this caller.
|
|
async fn launch(&self) -> Result<String> {
|
|
Ok(sausage_plugin::serde_json::from_value(
|
|
call("probe.service", json!({})).detach().await?,
|
|
)?)
|
|
}
|
|
/// Keep the process alive until the core closes stdin.
|
|
async fn service(&self) -> Result<String> {
|
|
sausage_plugin::listen("/state/control", std::rc::Rc::new(Control)).await?;
|
|
Ok(sausage_plugin::pid())
|
|
}
|
|
/// Keep an idle listener and ownership lock until opt-in drain ends this process.
|
|
async fn drain_service(&self) -> Result<String> {
|
|
lock("/drain-owner").await?.keep();
|
|
self.service().await
|
|
}
|
|
/// Add two numbers.
|
|
///
|
|
/// The result is their sum.
|
|
fn add(&self, p: Add) -> Result<i32> {
|
|
Ok(p.a + p.b)
|
|
}
|
|
/// Echo a payload larger than a host I/O buffer.
|
|
fn echo(&self, p: Echo) -> Result<String> {
|
|
Ok(p.text)
|
|
}
|
|
/// Return an internal serialization error.
|
|
fn serde_error(&self) -> Result<String> {
|
|
Ok(sausage_plugin::serde_json::from_str::<String>("invalid")?)
|
|
}
|
|
/// Return a deliberate error.
|
|
fn fail(&self) -> Result<String> {
|
|
Err(sausage_plugin::Error {
|
|
code: -32010,
|
|
message: "deliberate".into(),
|
|
data: Some(json!({"diagnostic":"probe error details"})),
|
|
})
|
|
}
|
|
/// Exercise calls, notifications, and advisory locks.
|
|
async fn exercise(&self, ctx: &Ctx) -> Result<String> {
|
|
let a =
|
|
call("other.first", json!({})).on_notify(|name, _| eprintln!("notification: {name}"));
|
|
let b = call("other.second", json!({})).detach();
|
|
let (a, b) = join(a, b).await;
|
|
a?;
|
|
b?;
|
|
ctx.notify("progress", json!({"done":true}))?;
|
|
{
|
|
let guard = lock("/state/probe/file").await?;
|
|
assert!(guard.held());
|
|
}
|
|
assert!(!lock("/state/probe/file").try_now().await?.held());
|
|
Ok("done".into())
|
|
}
|
|
/// Drop an acquisition before the core grants its lock.
|
|
async fn late(&self) -> Result<String> {
|
|
{
|
|
let mut future = std::pin::pin!(std::future::IntoFuture::into_future(lock("/late")));
|
|
std::future::poll_fn(|cx| {
|
|
assert!(future.as_mut().poll(cx).is_pending());
|
|
std::task::Poll::Ready(())
|
|
})
|
|
.await;
|
|
}
|
|
call("other.barrier", json!({})).await?;
|
|
lock("/kept").await?.keep();
|
|
Ok("late done".into())
|
|
}
|
|
/// Mark a retired version.
|
|
#[sausage_plugin::on_retired]
|
|
fn retired(&self, p: Retired) -> Result<String> {
|
|
Ok(p.previous)
|
|
}
|
|
/// Drain on replacement while calls are still available.
|
|
#[sausage_plugin::on_superseded]
|
|
async fn superseded(&self, p: Superseded) -> Result<()> {
|
|
call("other.superseded", json!({"current":p.current})).await?;
|
|
sausage_plugin::drain();
|
|
sausage_plugin::drain();
|
|
Ctx::default().notify("draining", json!({}))?;
|
|
Ok(())
|
|
}
|
|
/// Wait until EOF, then run the stop hook.
|
|
async fn wait(&self) -> Result<String> {
|
|
call("other.wait", json!({})).await?;
|
|
Ok("done".into())
|
|
}
|
|
/// Keep the process alive with a background call.
|
|
fn start(&self) -> Result<String> {
|
|
sausage_plugin::background(async {
|
|
let _ = call("other.background", json!({})).await;
|
|
});
|
|
Ok("started".into())
|
|
}
|
|
/// Leave an unfinished diagnostic before the SDK writes its framed response.
|
|
fn framed_output(&self) -> Result<String> {
|
|
print!("partial diagnostic");
|
|
Ok("intact".into())
|
|
}
|
|
/// Keep the process alive until the test releases its background task.
|
|
fn background_hold(&self) -> Result<String> {
|
|
sausage_plugin::background(async {
|
|
while !std::path::Path::new("/release-background").exists() {
|
|
wstd::task::sleep(wstd::time::Duration::from_millis(10)).await;
|
|
}
|
|
std::fs::write("/background-ended", "done").unwrap();
|
|
});
|
|
Ok(sausage_plugin::pid())
|
|
}
|
|
/// Keep a pollable alive for one hour.
|
|
fn sleepy(&self) -> Result<()> {
|
|
sausage_plugin::background(async {
|
|
wstd::task::sleep(wstd::time::Duration::from_secs(3600)).await;
|
|
});
|
|
Ok(())
|
|
}
|
|
/// Exit after the response despite a pending background pollable.
|
|
fn sleepy_exit(&self, ctx: &Ctx) -> Result<()> {
|
|
self.sleepy()?;
|
|
ctx.exit_after_response();
|
|
Ok(())
|
|
}
|
|
/// Save an EOF marker.
|
|
#[sausage_plugin::on_stop]
|
|
async fn stop(&self) {
|
|
std::fs::write("/stopped", b"yes").unwrap();
|
|
std::fs::write(format!("/stopped-{}", sausage_plugin::pid()), b"yes").unwrap();
|
|
Ctx::default().notify("stopped", json!({})).unwrap();
|
|
assert!(call("other.no", json!({})).await.is_err());
|
|
}
|
|
}
|
|
async fn join(
|
|
a: impl std::future::IntoFuture<Output = Result<Value>>,
|
|
b: impl std::future::IntoFuture<Output = Result<Value>>,
|
|
) -> (Result<Value>, Result<Value>) {
|
|
let mut a = std::pin::pin!(a.into_future());
|
|
let mut b = std::pin::pin!(b.into_future());
|
|
let (mut av, mut bv) = (None, None);
|
|
std::future::poll_fn(move |cx| {
|
|
if av.is_none() {
|
|
if let std::task::Poll::Ready(v) = a.as_mut().poll(cx) {
|
|
av = Some(v);
|
|
}
|
|
}
|
|
if bv.is_none() {
|
|
if let std::task::Poll::Ready(v) = b.as_mut().poll(cx) {
|
|
bv = Some(v);
|
|
}
|
|
}
|
|
if av.is_some() && bv.is_some() {
|
|
std::task::Poll::Ready((av.take().unwrap(), bv.take().unwrap()))
|
|
} else {
|
|
std::task::Poll::Pending
|
|
}
|
|
})
|
|
.await
|
|
}
|
|
struct Control;
|
|
#[rpc]
|
|
impl Control {
|
|
/// Report progress and return a stable counter.
|
|
fn stats(&self, ctx: &Ctx) -> Result<u64> {
|
|
ctx.notify("endpoint-progress", json!({"value":7}))?;
|
|
Ok(7)
|
|
}
|
|
/// Finish an in-flight call after draining begins.
|
|
async fn wait(&self) -> Result<String> {
|
|
call("other.inflight", json!({})).await?;
|
|
Ok("finished".into())
|
|
}
|
|
/// End the endpoint process after the response.
|
|
fn stop(&self, ctx: &Ctx) -> Result<()> {
|
|
ctx.exit_after_response();
|
|
Ok(())
|
|
}
|
|
}
|
|
fn main() {
|
|
sausage_plugin::run(Probe)
|
|
}
|