sausage/sdk/tests/plugins/probe/src/main.rs

//! 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;
/// 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,
}
/// A large payload.
#[schema]
struct Echo {
    /// The text to return unchanged.
    text: String,
}
#[rpc(grace_seconds = 12,endpoints(control=Control))]
impl Probe {
    /// 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())
    }
    /// 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::new(-32010, "deliberate"))
    }
    /// 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)
    }
    /// 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())
    }
    /// Save an EOF marker.
    #[sausage_plugin::on_stop]
    async fn stop(&self) {
        std::fs::write("/stopped", 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)
    }
    /// End the endpoint process after the response.
    fn stop(&self, ctx: &Ctx) -> Result<()> {
        ctx.exit_after_response();
        Ok(())
    }
}
fn main() {
    sausage_plugin::run(Probe)
}