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

sdk/src/calls.rs

@@ -1,5 +1,19 @@
 //! Outgoing core calls and notification callbacks.
 use super::*;
+/// Removes a dropped call's local routing state, without cancelling host work.
+struct PendingCall(u64);
+impl Drop for PendingCall {
+    fn drop(&mut self) {
+        let pending = STATE.with(|s| s.borrow_mut().pending.remove(&self.0));
+        if let Some(mut pending) = pending {
+            // A lock may still be granted: release it when its late response arrives.
+            if let Some(lease) = pending.unlock.take() {
+                STATE.with(|s| s.borrow_mut().abandoned_locks.insert(self.0, lease));
+            }
+            // Drop user-owned callback captures outside the STATE borrow.
+        }
+    }
+}
 /// An outgoing request builder.
 #[must_use = "await the call to send it"]
 pub struct Call {
@@ -26,12 +40,16 @@ impl Call {
         self
     }
     /// Receive progress notifications before the result arrives.
+    /// Dropping the call future removes its local pending entry and callback;
+    /// later notifications and responses are ignored. It does not cancel host work.
     ///
     /// For `core.watch`, poll the call future alongside work in the handler, then
     /// drop it when the handler returns. Do not put it in `background()` just to
     /// keep the watch alive: background tasks keep the process alive, while the
     /// watch only ends when the process ends. Dropping the future does not cancel
-    /// the host watch, but it lets an otherwise idle process exit normally.
+    /// the host watch, but it lets an otherwise idle process exit normally. A
+    /// long-lived service should keep one watch per purpose for its lifetime,
+    /// rather than repeatedly opening and dropping watches on the host.
     ///
     /// Scan only after `armed`, and rescan on `changed`. Events may coalesce, but
     /// are not rate-limited: watch the narrowest useful path to avoid rescanning
@@ -71,6 +89,7 @@ impl IntoFuture for Call {
                 );
                 Ok(id)
             })?;
+            let _pending = PendingCall(id);
             let (method, params) = if self.detach {
                 (
                     "core.call".to_string(),
@@ -80,7 +99,8 @@ impl IntoFuture for Call {
                 (self.method, self.params)
             };
             if let Err(e) = send(json!({"jsonrpc":"2.0","id":id,"method":method,"params":params})) {
-                STATE.with(|s| s.borrow_mut().pending.remove(&id));
+                let pending = STATE.with(|s| s.borrow_mut().pending.remove(&id));
+                drop(pending);
                 return Err(e);
             }
             rx.await
@@ -88,3 +108,154 @@ impl IntoFuture for Call {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{STATE, State, block_on, call, clear_pending, lock, receive};
+    use crate::Result;
+    use futures::{FutureExt, future::LocalBoxFuture};
+    use serde_json::{Value, json};
+    use std::{cell::RefCell, future::IntoFuture, rc::Rc, task::Context};
+
+    fn activate() {
+        STATE.with(|s| {
+            *s.borrow_mut() = State {
+                active: true,
+                ..State::default()
+            };
+        });
+    }
+    fn start(
+        mut call: LocalBoxFuture<'static, Result<Value>>,
+    ) -> (u64, LocalBoxFuture<'static, Result<Value>>) {
+        assert!(
+            call.as_mut()
+                .poll(&mut Context::from_waker(std::task::Waker::noop()))
+                .is_pending()
+        );
+        (STATE.with(|s| s.borrow().next), call)
+    }
+    fn notification(id: u64, name: &str) {
+        receive(json!({"method":"core.notify","params":{"call":id,"name":name,"params":{}}}));
+    }
+    fn response(id: u64, result: Value) {
+        receive(json!({"id":id,"result":result}));
+    }
+
+    #[test]
+    fn dropped_call_removes_pending_and_callback() {
+        activate();
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let captured = events.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
+                .into_future(),
+        );
+        notification(id, "armed");
+        assert_eq!(&*events.borrow(), &["armed"]);
+        assert_eq!(Rc::strong_count(&events), 2);
+        drop(future);
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert!(s.borrow().abandoned_locks.is_empty());
+        });
+        assert_eq!(Rc::strong_count(&events), 1);
+        notification(id, "changed");
+        response(id, Value::Null);
+        receive(json!({"id":id,"error":{"code":-32000,"message":"late error"}}));
+        assert_eq!(&*events.borrow(), &["armed"]);
+    }
+
+    #[test]
+    fn kept_watch_delivers_until_result() {
+        activate();
+        let events = Rc::new(RefCell::new(Vec::new()));
+        let captured = events.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |name, _| captured.borrow_mut().push(name.to_owned()))
+                .into_future(),
+        );
+        notification(id, "armed");
+        notification(id, "changed");
+        notification(id, "changed");
+        response(id, json!({"done":true}));
+        assert_eq!(block_on(future).unwrap(), json!({"done":true}));
+        notification(id, "late");
+        assert_eq!(&*events.borrow(), &["armed", "changed", "changed"]);
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+        assert_eq!(Rc::strong_count(&events), 1);
+    }
+
+    #[test]
+    fn callback_can_drop_its_own_call_without_being_restored() {
+        activate();
+        let holder = Rc::new(RefCell::new(None));
+        let captured = holder.clone();
+        let (id, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |_, _| {
+                    // Panics on a second invocation; the first must remove the entry.
+                    drop(captured.borrow_mut().take().expect("live call"));
+                })
+                .into_future(),
+        );
+        *holder.borrow_mut() = Some(future);
+        notification(id, "armed");
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+        notification(id, "changed");
+        response(id, Value::Null);
+        assert_eq!(Rc::strong_count(&holder), 1);
+    }
+
+    #[test]
+    fn eof_cleanup_can_drop_callback_owned_futures() {
+        activate();
+        let holder = Rc::new(RefCell::new(None));
+        let weak = Rc::downgrade(&holder);
+        let captured = holder.clone();
+        let (_, future) = start(
+            call("core.watch", json!({"path":"/mailbox"}))
+                .on_notify(move |_, _| {
+                    let _ = &captured;
+                })
+                .into_future(),
+        );
+        *holder.borrow_mut() = Some(future);
+        drop(holder);
+        STATE.with(|s| s.borrow_mut().stopping = true);
+        clear_pending();
+        assert!(weak.upgrade().is_none());
+        STATE.with(|s| assert!(s.borrow().pending.is_empty()));
+    }
+
+    #[test]
+    fn unpolled_call_never_registers() {
+        activate();
+        drop(call("core.watch", json!({"path":"/mailbox"})).into_future());
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert_eq!(s.borrow().next, 0);
+        });
+    }
+
+    #[test]
+    fn dropped_lock_keeps_only_late_grant_cleanup() {
+        activate();
+        let (id, future) = start(
+            lock("/late")
+                .into_future()
+                .map(|r| r.map(|_| Value::Null))
+                .boxed_local(),
+        );
+        drop(future);
+        STATE.with(|s| {
+            assert!(s.borrow().pending.is_empty());
+            assert_eq!(s.borrow().abandoned_locks.len(), 1);
+        });
+        notification(id, "ignored");
+        response(id, json!({"held":true}));
+        STATE.with(|s| assert!(s.borrow().abandoned_locks.is_empty()));
+    }
+}