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