Notify superseded plugin processes and support SDK draining

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