Notify superseded plugin processes and support SDK draining

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.<method>` 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": "<sha256 of the new file>"}`
+   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

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<OldProcess>,
     services: Vec<(Service, Option<watch::Receiver<bool>>)>,
@@ -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) {

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.

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::<Value>(&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;
         }

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<dyn ::std::future::Future<Output=()> + '_>> {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<T: ::sausage_plugin::Params>() {} 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 {

sdk/src/endpoints.rs

@@ -65,7 +65,7 @@ impl<T: serde::de::DeserializeOwned + 'static> IntoFuture for EndpointCall<T> {
     }
 }
 
-/// 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<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> Result<()> {
     use wstd::iter::AsyncIterator;
@@ -80,7 +80,7 @@ pub async fn listen<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> 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<T: Interface>(
     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 {

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};

sdk/src/transport.rs

@@ -96,6 +96,7 @@ struct State {
     active: bool,
     endpoints: HashMap<std::any::TypeId, String>,
     stopping: bool,
+    draining: bool,
     next: u64,
     pending: HashMap<u64, Pending>,
     #[cfg(target_os = "wasi")]
@@ -125,17 +126,48 @@ pub(crate) fn endpoint_name<T: 'static>() -> Result<String> {
 }
 #[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);
-        id
-    });
-    background(async move {
-        let _ = futures::future::Abortable::new(future, registration).await;
+        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")]

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<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?;
@@ -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<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();

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() {}

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.
+  |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^

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::<Value>(&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::<Value>(&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();