Notify superseded plugin processes and support SDK draining

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;
-        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")]