diff --git a/PROTOCOL.md b/PROTOCOL.md index a0f1154..e9e6b6b 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -70,7 +70,7 @@ Rust plugins use the SDK `sausage-plugin`. Its source is at `notes.rpc.discover` sends it `rpc.discover`. - `core.` is a method of the core (section 9). - Method names that start with `rpc.` are reserved: `rpc.discover`, - `rpc.retired`, and `rpc.endpoint`. + `rpc.superseded`, `rpc.retired`, and `rpc.endpoint`. ### 3.3 What a call does @@ -153,9 +153,14 @@ A process is an instance of one plugin version. When a new version becomes active (section 10.3): 1. New calls get the new version at once. -2. Each process of an older version gets a grace period to finish its work. - The default is 60 seconds. A plugin can declare its own value in - `rpc.discover` (`x-sausage-grace`). Then the core stops the process. +2. At grace start, the core sends each process of an older version the + JSON-RPC notification `rpc.superseded {"current": ""}` + on stdin, without an ID. The process is still running: calls and responses + continue to work. It can stop accepting work, finish its in-flight calls, + and exit early. Ignoring the notice changes nothing. The grace period is + the upper bound before the core stops the process (section 3.6); its + default is 60 seconds. A plugin can declare its own value in its old + `rpc.discover` document (`x-sausage-grace`). 3. When the old process of a service has ended, the core starts the service again with the new version. 4. When the last process of an older version has ended, the core starts a diff --git a/cli/src/upgrades.rs b/cli/src/upgrades.rs index 0861982..bfa1783 100644 --- a/cli/src/upgrades.rs +++ b/cli/src/upgrades.rs @@ -12,6 +12,7 @@ pub(crate) struct OldProcess { pub(crate) struct Upgrade { plugin: String, previous: String, + current: String, grace: Duration, processes: Vec, services: Vec<(Service, Option>)>, @@ -67,6 +68,7 @@ impl Host { Upgrade { plugin: plugin.into(), previous, + current: current.into(), grace: Duration::from_secs_f64(grace), processes, services, @@ -74,6 +76,13 @@ impl Host { } /// Lets old calls finish, stops them after grace, and starts each lifecycle hook once. pub(crate) fn begin_upgrade(&self, upgrade: Upgrade) { + // Queue the notice before returning from activation, without waiting on guest code. + for process in &upgrade.processes { + self.send( + &process.pid, + json!({"jsonrpc":"2.0","method":"rpc.superseded","params":{"current":upgrade.current}}), + ); + } let host = self.clone(); tokio::spawn(async move { if host.0.shutdown.load(Ordering::SeqCst) { diff --git a/cli/tests/host/activation.rs b/cli/tests/host/activation.rs index 04769ee..dd20dab 100644 --- a/cli/tests/host/activation.rs +++ b/cli/tests/host/activation.rs @@ -839,6 +839,76 @@ async fn old_process_grace_and_new_calls() { "2" ); marker(world.0.join("state/notes/ended-1"), "ended").await; + let notice: Value = serde_json::from_slice( + &std::fs::read(world.0.join(format!("state/notes/superseded-{pid}"))).unwrap(), + ) + .unwrap(); + assert_eq!(notice["params"]["current"], fixture_hash(2)); + host.stop().await; +} +/// PROTOCOL 3.7: every old process gets the new hash while still running, before grace ends. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn superseded_notifies_each_old_process() { + let world = TestWorld::new(); + let mut old = document("hold"); + old["x-sausage-grace"] = json!(60); + world.plugin("notes", old, document("hold")); + let host = world.start().await; + call(&host, "notes.hold", json!({})).await; + call(&host, "notes.hold", json!({})).await; + let pids: Vec<_> = host + .processes() + .into_iter() + .filter(|(_, process)| process.method == "hold") + .map(|(pid, _)| pid) + .collect(); + assert_eq!(pids.len(), 2); + activate(&host, "notes").await.unwrap(); + let expected = + json!({"jsonrpc":"2.0","method":"rpc.superseded","params":{"current":fixture_hash(2)}}) + .to_string(); + for pid in pids { + // marker's existing 30-second bound is shorter than the old 60-second grace. + marker( + world.0.join(format!("state/notes/superseded-{pid}")), + &expected, + ) + .await; + assert_eq!(host.processes()[&pid].state, "running"); + } + assert_eq!( + call(&host, "notes.version", json!({})).await["version"], + "2" + ); + host.stop().await; +} +/// PROTOCOL 3.7: a service that exits on the notice restarts before its grace deadline. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn superseded_service_exits_early() { + let world = TestWorld::new(); + let mut old = document("serve"); + old["x-sausage-grace"] = json!(60); + world.plugin("notes", old, document("serve")); + std::fs::write( + world.0.join("config/boot.yaml"), + json!([{"method":"notes.serve","params":{"exit_on_superseded":true}}]).to_string(), + ) + .unwrap(); + let host = world.start().await; + marker(world.0.join("state/notes/service"), "1").await; + let pid = host + .processes() + .into_iter() + .find(|(_, p)| p.method == "serve") + .unwrap() + .0; + activate(&host, "notes").await.unwrap(); + marker(world.0.join("state/notes/service"), "2").await; + assert!(!host.processes().contains_key(&pid)); + assert!( + !world.0.join("state/notes/ended-1").exists(), + "old service must exit on the notice, not EOF" + ); host.stop().await; } /// PROTOCOL 3.7: retirement receives the previous hash only after old processes end. diff --git a/cli/tests/plugins/activation/src/main.rs b/cli/tests/plugins/activation/src/main.rs index 8695025..60c0c05 100644 --- a/cli/tests/plugins/activation/src/main.rs +++ b/cli/tests/plugins/activation/src/main.rs @@ -89,7 +89,24 @@ fn main() { if request["method"] == "serve" { std::fs::write(format!("/state/{plugin}/service"), version).unwrap(); } - for _ in lines {} + for line in lines { + let Ok(message) = serde_json::from_str::(&line.unwrap()) else { + continue; + }; + if message["method"] == "rpc.superseded" { + std::fs::write( + format!( + "/state/{plugin}/superseded-{}", + std::env::var("SAUSAGE_PID").unwrap() + ), + message.to_string(), + ) + .unwrap(); + if request["params"]["exit_on_superseded"] == true { + return; + } + } + } std::fs::write(format!("/state/{plugin}/ended-{version}"), "ended").unwrap(); return; } diff --git a/sdk/macros/src/lib.rs b/sdk/macros/src/lib.rs index a74a0f6..a998865 100644 --- a/sdk/macros/src/lib.rs +++ b/sdk/macros/src/lib.rs @@ -153,6 +153,7 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream { let mut stop = quote!(); let mut has_stop = false; let mut has_retired = false; + let mut has_superseded = false; for member in &mut item.items { let ImplItem::Fn(f) = member else { continue }; let retired = f.attrs.iter().any(|a| { @@ -161,6 +162,12 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream { .last() .is_some_and(|s| s.ident == "on_retired") }); + let superseded = f.attrs.iter().any(|a| { + a.path() + .segments + .last() + .is_some_and(|s| s.ident == "on_superseded") + }); let stopping = f.attrs.iter().any(|a| { a.path() .segments @@ -179,11 +186,16 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream { "only one on_retired method is allowed", )); } + if superseded && std::mem::replace(&mut has_superseded, true) { + return Err(syn::Error::new( + f.span(), + "only one on_superseded method is allowed", + )); + } f.attrs.retain(|a| { - !a.path() - .segments - .last() - .is_some_and(|s| s.ident == "on_retired" || s.ident == "on_stop") + !a.path().segments.last().is_some_and(|s| { + s.ident == "on_retired" || s.ident == "on_stop" || s.ident == "on_superseded" + }) }); documented(&f.attrs, f.span())?; if f.sig.inputs.is_empty() @@ -224,7 +236,9 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream { stop = quote!(fn stop(&self)->::std::pin::Pin + '_>> {Box::pin(async move {let _=self.#name() #await_;})}); continue; } - let wire = if retired { + let wire = if superseded { + "rpc.superseded".to_string() + } else if retired { "rpc.retired".to_string() } else { name.unraw().to_string() @@ -253,7 +267,7 @@ pub fn rpc(args: TokenStream, input: TokenStream) -> TokenStream { }; let doc = docs(&f.attrs); arms.push(quote!(#wire => { fn params_type() {} params_type::<#p>(); let result=self.#name(#(#args),*) #await_ ?; Ok(::sausage_plugin::serde_json::to_value(result)?) })); - if !retired { + if !retired && !superseded { entries.push(quote!(::sausage_plugin::method::<#p,#output>(&mut doc,#wire,#doc);)); } } @@ -328,6 +342,12 @@ pub fn import(input: TokenStream) -> TokenStream { pub fn on_retired(_: TokenStream, input: TokenStream) -> TokenStream { input } +/// Mark the method that handles `rpc.superseded` at the start of upgrade grace. +/// Calls still work here. Use `sausage_plugin::drain` to finish active work and exit. +#[proc_macro_attribute] +pub fn on_superseded(_: TokenStream, input: TokenStream) -> TokenStream { + input +} /// Mark the hook that runs at stdin EOF without a live call transport. #[proc_macro_attribute] pub fn on_stop(_: TokenStream, input: TokenStream) -> TokenStream { diff --git a/sdk/src/endpoints.rs b/sdk/src/endpoints.rs index c86842c..37c5b64 100644 --- a/sdk/src/endpoints.rs +++ b/sdk/src/endpoints.rs @@ -65,7 +65,7 @@ impl IntoFuture for EndpointCall { } } -/// Publish a declared interface and keep its listener alive until stdin EOF. +/// Publish a declared interface and keep its listener alive until stdin EOF or [`drain`]. #[cfg(target_os = "wasi")] pub async fn listen(path: impl AsRef, value: Rc) -> Result<()> { use wstd::iter::AsyncIterator; @@ -80,7 +80,7 @@ pub async fn listen(path: impl AsRef, value: Rc) -> Resul while let Some(Ok(stream)) = incoming.next().await { let value = value.clone(); let name = name.clone(); - crate::transport::endpoint_task(async move { + background(async move { let _ = serve_connection(stream, value, name).await; }); } @@ -211,7 +211,11 @@ async fn serve_connection( use futures::{FutureExt, StreamExt}; let (mut reader, mut writer) = stream.split(); let mut lines = crate::lines::Lines::default(); - while let Some(request) = lines.rpc(&mut reader).await? { + while let Some(request) = crate::transport::until_drained(lines.rpc(&mut reader)) + .await + .transpose()? + .flatten() + { let method = request["method"].as_str().unwrap_or_default(); let (tx, mut rx) = futures::channel::mpsc::unbounded(); let ctx = Ctx { diff --git a/sdk/src/lib.rs b/sdk/src/lib.rs index c8e2d44..d8865d9 100644 --- a/sdk/src/lib.rs +++ b/sdk/src/lib.rs @@ -358,11 +358,16 @@ //! `core.revert` restores code, not state. Preserve compatibility or migrate the state. //! `#[rpc(grace_seconds = N)]` declares the grace period for old processes. The default is 60 seconds. //! New calls use the new version immediately. Old processes can still use old imports during their grace period. +//! [`macro@on_superseded`] handles the notice at grace start, with normal calls still available. +//! It is opt-in: without a hook the process continues as before until core stops it. +//! Call [`drain`] in this hook to close SDK endpoint listeners and idle connections, +//! let active handlers and background tasks finish, and exit without waiting for grace. +//! Custom listeners and long-lived background tasks must be ended by the plugin itself. //! [`macro@on_retired`] handles `rpc.retired` after the last process of an old version ends. //! This hook is best effort. If the host stops first, the hook does not run. //! New processes must handle unmigrated state. Check a migration marker and migrate if it is absent. //! See section 3.7 of `/proc/core/PROTOCOL.md`. -//! Declare at most one stop hook and one retirement hook per interface. +//! Declare at most one hook of each kind per interface. //! //! ```no_run //! use sausage_plugin::{rpc, schema, state_dir, write_atomic, Result}; @@ -374,8 +379,21 @@ //! /// The SHA-256 of the previous WASM file. //! previous: String, //! } +//! /// The version now receiving new calls. +//! #[schema] +//! struct SupersededParams { +//! /// The SHA-256 of the newly active WASM file. +//! current: String, +//! } //! #[rpc] //! impl Migrator { +//! /// Finish current work and exit early during upgrade grace. +//! #[sausage_plugin::on_superseded] +//! fn superseded(&self, p: SupersededParams) -> Result<()> { +//! eprintln!("replaced by {}", p.current); +//! sausage_plugin::drain(); +//! Ok(()) +//! } //! /// Record a retired version. Normal methods must also handle missing markers. //! #[sausage_plugin::on_retired] //! fn retired(&self, p: RetiredParams) -> Result<()> { @@ -401,6 +419,6 @@ pub use helpers::*; pub use imports::*; #[doc(hidden)] pub use inventory; -pub use sausage_plugin_macros::{import, on_retired, on_stop, rpc, schema, test}; +pub use sausage_plugin_macros::{import, on_retired, on_stop, on_superseded, rpc, schema, test}; pub use transport::*; pub use {schemars, serde, serde_json}; diff --git a/sdk/src/transport.rs b/sdk/src/transport.rs index 97e7aa7..c560b79 100644 --- a/sdk/src/transport.rs +++ b/sdk/src/transport.rs @@ -96,6 +96,7 @@ struct State { active: bool, endpoints: HashMap, stopping: bool, + draining: bool, next: u64, pending: HashMap, #[cfg(target_os = "wasi")] @@ -125,17 +126,48 @@ pub(crate) fn endpoint_name() -> Result { } #[cfg(target_os = "wasi")] pub(crate) fn endpoint_task(future: impl Future + '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(future: F) -> Option { 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); - id - }); - background(async move { - let _ = futures::future::Abortable::new(future, registration).await; - STATE.with(|s| s.borrow_mut().listeners.remove(&id)); + 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")] diff --git a/sdk/tests/plugins/probe/src/main.rs b/sdk/tests/plugins/probe/src/main.rs index e6aae1a..a4b013b 100644 --- a/sdk/tests/plugins/probe/src/main.rs +++ b/sdk/tests/plugins/probe/src/main.rs @@ -16,6 +16,12 @@ struct Retired { /// The previous hash. previous: String, } +/// The replacement version. +#[schema] +struct Superseded { + /// The current hash. + current: String, +} /// A large payload. #[schema] struct Echo { @@ -111,6 +117,15 @@ impl Probe { fn retired(&self, p: Retired) -> Result { 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 { call("other.wait", json!({})).await?; @@ -193,6 +208,11 @@ impl Control { ctx.notify("endpoint-progress", json!({"value":7}))?; Ok(7) } + /// Finish an in-flight call after draining begins. + async fn wait(&self) -> Result { + 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(); diff --git a/sdk/tests/ui/superseded.rs b/sdk/tests/ui/superseded.rs new file mode 100644 index 0000000..ad7826f --- /dev/null +++ b/sdk/tests/ui/superseded.rs @@ -0,0 +1,11 @@ +struct Supersedes; +#[sausage_plugin::rpc] +impl Supersedes { + /// First superseded hook. + #[sausage_plugin::on_superseded] + fn a(&self) -> sausage_plugin::Result<()> { Ok(()) } + /// Second superseded hook. + #[sausage_plugin::on_superseded] + fn b(&self) -> sausage_plugin::Result<()> { Ok(()) } +} +fn main() {} diff --git a/sdk/tests/ui/superseded.stderr b/sdk/tests/ui/superseded.stderr new file mode 100644 index 0000000..3383a8a --- /dev/null +++ b/sdk/tests/ui/superseded.stderr @@ -0,0 +1,5 @@ +error: only one on_superseded method is allowed + --> tests/ui/superseded.rs:7:5 + | +7 | /// Second superseded hook. + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/sdk/tests/wasi.rs b/sdk/tests/wasi.rs index b7bfcf2..251a3de 100644 --- a/sdk/tests/wasi.rs +++ b/sdk/tests/wasi.rs @@ -201,6 +201,88 @@ fn protocol() { assert_eq!(outgoing["method"], "other.background"); p.answer(&outgoing, json!({})); assert!(p.child.wait().unwrap().success()); + // Replacement is a notification, not EOF. Drain keeps both endpoint handlers + // and background calls alive, while closing idle connections and the listener. + assert!( + !d["result"]["methods"] + .as_array() + .unwrap() + .iter() + .any(|m| m["name"] == "rpc.superseded" || m["name"] == "superseded") + ); + std::fs::remove_file(dir.join("stopped")).unwrap(); + let mut provider = start(&wasm, &dir); + provider.request("serve", json!({})); + assert_eq!(provider.read()["result"], Value::Null); + provider.request("start", json!({})); + let first = provider.read(); + let second = provider.read(); + let background = if first["method"] == "other.background" { + first + } else { + second + }; + assert_eq!(background["method"], "other.background"); + let address = std::fs::read_to_string(dir.join("endpoint")) + .unwrap() + .split_whitespace() + .next() + .unwrap() + .trim_start_matches("tcp:") + .to_owned(); + let mut idle = std::net::TcpStream::connect(&address).unwrap(); + idle.set_read_timeout(Some(std::time::Duration::from_secs(4))) + .unwrap(); + writeln!( + idle, + "{}", + json!({"jsonrpc":"2.0","id":0,"method":"rpc.endpoint","params":{}}) + ) + .unwrap(); + let mut idle = BufReader::new(idle); + let mut line = String::new(); + idle.read_line(&mut line).unwrap(); + assert_eq!( + sausage_plugin::serde_json::from_str::(&line).unwrap()["id"], + 0 + ); + let mut active = std::net::TcpStream::connect(&address).unwrap(); + active + .set_read_timeout(Some(std::time::Duration::from_secs(4))) + .unwrap(); + writeln!( + active, + "{}", + json!({"jsonrpc":"2.0","id":1,"method":"wait","params":{}}) + ) + .unwrap(); + let inflight = provider.read(); + assert_eq!(inflight["method"], "other.inflight"); + provider + .send(json!({"jsonrpc":"2.0","method":"rpc.superseded","params":{"current":"new-hash"}})); + let hook = provider.read(); + assert_eq!(hook["method"], "other.superseded"); + assert_eq!(hook["params"]["current"], "new-hash"); + provider.answer(&hook, json!({})); + assert_eq!(provider.read()["params"]["name"], "draining"); + line.clear(); + assert_eq!(idle.read_line(&mut line).unwrap(), 0); + assert!(std::net::TcpStream::connect(&address).is_err()); + assert!(provider.child.try_wait().unwrap().is_none()); + provider.answer(&inflight, json!({})); + line.clear(); + BufReader::new(active).read_line(&mut line).unwrap(); + assert_eq!( + sausage_plugin::serde_json::from_str::(&line).unwrap()["result"], + "finished" + ); + assert!(provider.child.try_wait().unwrap().is_none()); + provider.answer(&background, json!({})); + provider.finishes(); + assert!( + !dir.join("stopped").exists(), + "drain is not the EOF stop hook" + ); let provider_doc = &d["result"]; let imports = target.join("imports"); std::fs::create_dir_all(&imports).unwrap();