Notify superseded plugin processes and support SDK draining

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