Clean up dropped SDK calls and make drain assertions order-independent

sdk/src/transport.rs

@@ -99,6 +99,8 @@ struct State {
     draining: bool,
     next: u64,
     pending: HashMap<u64, Pending>,
+    // Only leases, never callbacks or response senders, survive dropped lock acquisitions.
+    abandoned_locks: HashMap<u64, std::rc::Rc<Lease>>,
     #[cfg(target_os = "wasi")]
     tasks: usize,
     #[cfg(target_os = "wasi")]
@@ -109,6 +111,18 @@ struct State {
     listeners: HashMap<u64, futures::future::AbortHandle>,
 }
 thread_local! { static STATE: RefCell<State> = RefCell::new(State::default()); }
+#[cfg(any(target_os = "wasi", test))]
+fn clear_pending() {
+    let pending = STATE.with(|s| {
+        let mut s = s.borrow_mut();
+        (
+            std::mem::take(&mut s.pending),
+            std::mem::take(&mut s.abandoned_locks),
+        )
+    });
+    // Callback captures may own other call futures whose drop guards use STATE.
+    drop(pending);
+}
 pub(crate) fn endpoint_name<T: 'static>() -> Result<String> {
     STATE
         .with(|s| {
@@ -224,16 +238,16 @@ pub fn test_run<F: Future>(f: F) -> F::Output {
             STATE.with(|s| {
                 let mut s = s.borrow_mut();
                 s.stopping = true;
-                s.pending.clear();
             });
+            clear_pending();
         });
         let result = f.await;
         drop(reader);
         STATE.with(|s| {
             let mut s = s.borrow_mut();
             s.active = false;
-            s.pending.clear();
         });
+        clear_pending();
         result
     })
 }
@@ -306,7 +320,16 @@ pub fn receive(value: Value) {
         }
         let _ = p.tx.send(result);
     } else {
-        eprintln!("note: response with unknown ID {}", value["id"]);
+        let lease = value["id"]
+            .as_u64()
+            .and_then(|id| STATE.with(|s| s.borrow_mut().abandoned_locks.remove(&id)));
+        if let Some(lease) = lease {
+            lease
+                .held
+                .set(value.get("error").is_none() && value["result"]["held"] == true);
+        } else {
+            eprintln!("note: response with unknown ID {}", value["id"]);
+        }
     }
 }
 
@@ -356,11 +379,11 @@ async fn serve<T: Interface>(value: T) {
                 STATE.with(|s| {
                     let mut s = s.borrow_mut();
                     s.stopping = true;
-                    s.pending.clear();
                     for (_, listener) in s.listeners.drain() {
                         listener.abort();
                     }
                 });
+                clear_pending();
                 value.stop().await;
                 break;
             }
@@ -411,8 +434,8 @@ async fn serve<T: Interface>(value: T) {
     STATE.with(|s| {
         let mut s = s.borrow_mut();
         s.active = false;
-        s.pending.clear();
     });
+    clear_pending();
     let _ = std::io::stdout().flush();
     std::process::exit(0);
 }