Repository: capnproto/capnproto-rust (language: rust)

Issue:
Using `CapabilityServerSet` prevents Server objects from being dropped at the appropriate time
When I use a `CapabilityServerSet` to create clients for Server objects, those Server objects are no longer dropped when the remote connection drops them. These Server objects are kept alive until the `CapabilityServerSet` itself is dropped. 

The behavior I want is for the reference to the Server object to be removed from the `CapabilityServerSet` when it is dropped by the remote connection, so that the Server object may be dropped immediately.

Is it possible to achieve this?


Discussion / hints:
Good point. Looks like C++ implementation acts as you describe: https://github.com/capnproto/capnproto/blob/851a9e23ea4fb2f1dd0a2d34924407ee612d1da2/c%2B%2B/src/capnp/capability.h#L670-L671

Hm... maybe we could fix that by using `std::rc::Weak` instead of `std::rc::Rc` in the definition of `CapabilityServerSet`: 
https://github.com/capnproto/capnproto-rust/blob/27721b7be92b7d7654a56bae98a6a2ee6b3acf4b/capnp-rpc/src/lib.rs#L325
> Hm... maybe we could fix that by using `std::rc::Weak` instead of `std::rc::Rc` in the definition of `CapabilityServerSet`:

I tested this and I think it does what I want. However note that it would change the return types of the accessor methods in `CapabilityServerSet`.

I will create a PR shortly with the changes I tested.
Ah, I see -- the problem is that we currently return a `&Rc`. That will need to change into an `Rc` (and return `None` if the `Weak` cannot be upgraded).

Maybe for compatibility purposes, `CapabilityServerSet` could have
```
   caps: std::collections::HashMap<usize, Weak<RefCell<C::Dispatch>>>,
   owned_caps: std::collections::HashMap<usize, Rc<RefCell<C::Dispatch>>>,
```

and `owned_caps` only gets populated when the existing `get_local_server()` or `get_local_server_of_resolved()` get called. Then we could add new methods `upgrade_local_server()` and `upgrade_local_server_of_resolved()` that have the new behavior. That way we wouldn't need to make a breaking change.
Or maybe we make just a separate struct named something like `WeakCapabilityServerSet` with the new semantics. Then on the next version bump we can remove the existing one.
Another downside to this solution is that the invalid entries in the hashmap will never get deleted after the underlying objects get dropped, which could be a source of increasing memory usage for long-running connections with a lot of object creations and deletions.

I doubt removing failed lookups on access would mitigate this since the object almost certainly will still exist when the get_local_* methods are called.

I don't think this would be a concern for my intended usage though.
I was thinking we would add a `gc(&mut self)` method that would walk through the existing objects and drop the weak pointers that no longer point to anything.
> Or maybe we make just a separate struct named something like `WeakCapabilityServerSet` with the new semantics. Then on the next version bump we can remove the existing one.

I think I would prefer to keep the same name as the C++ library concept, unless you think this would be a more idiomatic name for rust.
On the version bump we'll rename it back to `CapabilityServerSet` (and maybe keep around a deprecated type alias to `WeakCapabilityServerSet`).
Ah, I see. That sounds good to me then.
I can take a stab at the changes if you'd like.
Yes, please go ahead if you're interested.

Fix patch:
diff --git a/capnp-rpc/src/lib.rs b/capnp-rpc/src/lib.rs
index bede45b6d..a9162e745 100644
--- a/capnp-rpc/src/lib.rs
+++ b/capnp-rpc/src/lib.rs
@@ -65,7 +65,7 @@ use futures::channel::oneshot;
 use futures::{Future, FutureExt, TryFutureExt};
 use std::cell::RefCell;
 use std::pin::Pin;
-use std::rc::Rc;
+use std::rc::{Rc, Weak};
 use std::task::{Context, Poll};
 
 pub use crate::rpc::Disconnector;
@@ -383,6 +383,80 @@ where
     }
 }
 
+/// Allows a server to recognize its own capabilities when passed back to it, and obtain the
+/// underlying Server objects associated with them. Holds only weak references to Server objects
+/// allowing Server objects to be dropped when dropped by the remote client. Call the `gc` method
+/// to reclaim memory used for Server objects that have been dropped.
+pub struct WeakCapabilityServerSet<S, C>
+where
+    C: capnp::capability::FromServer<S>,
+{
+    caps: std::collections::HashMap<usize, Weak<RefCell<C::Dispatch>>>,
+}
+
+impl<S, C> Default for WeakCapabilityServerSet<S, C>
+where
+    C: capnp::capability::FromServer<S>,
+{
+    fn default() -> Self {
+        Self {
+            caps: std::default::Default::default(),
+        }
+    }
+}
+
+impl<S, C> WeakCapabilityServerSet<S, C>
+where
+    C: capnp::capability::FromServer<S>,
+{
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Adds a new capability to the set and returns a client backed by it.
+    pub fn new_client(&mut self, s: S) -> C {
+        let dispatch = <C as capnp::capability::FromServer<S>>::from_server(s);
+        let wrapped = Rc::new(RefCell::new(dispatch));
+        let ptr = wrapped.as_ptr() as usize;
+        self.caps.insert(ptr, Rc::downgrade(&wrapped));
+        capnp::capability::FromClientHook::new(Box::new(local::Client::from_rc(wrapped)))
+    }
+
+    /// Looks up a capability and returns its underlying server object, if found.
+    /// Fully resolves the capability before looking it up.
+    pub async fn get_local_server(&self, client: &C) -> Option<Rc<RefCell<C::Dispatch>>>
+    where
+        C: capnp::capability::FromClientHook,
+    {
+        let resolved: C = capnp::capability::get_resolved_cap(
+            capnp::capability::FromClientHook::new(client.as_client_hook().add_ref()),
+        )
+        .await;
+        let hook = resolved.into_client_hook();
+        let ptr = hook.get_ptr();
+        self.caps.get(&ptr).and_then(|c| c.upgrade())
+    }
+
+    /// Looks up a capability and returns its underlying server object, if found.
+    /// Does *not* attempt to resolve the capability first, so you will usually want
+    /// to call `get_resolved_cap()` before calling this. The advantage of this method
+    /// over `get_local_server()` is that this one is synchronous and borrows `self`
+    /// over a shorter span (which can be very important if `self` is inside a `RefCell`).
+    pub fn get_local_server_of_resolved(&self, client: &C) -> Option<Rc<RefCell<C::Dispatch>>>
+    where
+        C: capnp::capability::FromClientHook,
+    {
+        let hook = client.as_client_hook();
+        let ptr = hook.get_ptr();
+        self.caps.get(&ptr).and_then(|c| c.upgrade())
+    }
+
+    /// Reclaim memory used for Server objects that no longer exist.
+    pub fn gc(&mut self) {
+        self.caps.retain(|_, c| c.strong_count() > 0);
+    }
+}
+
 /// Converts a promise for a client into a client that queues up any calls that arrive
 /// before the promise resolves.
 // TODO: figure out a better way to allow construction of promise clients.


Test patch:
diff --git a/capnp-rpc/test/test.rs b/capnp-rpc/test/test.rs
index 42dc2c233..fd812f54b 100644
--- a/capnp-rpc/test/test.rs
+++ b/capnp-rpc/test/test.rs
@@ -916,11 +916,11 @@ fn capability_list() {
 fn capability_server_set() {
     use crate::impls;
     use crate::test_capnp::test_interface;
-    use capnp_rpc::CapabilityServerSet;
-    let mut set1: CapabilityServerSet<impls::TestInterface, test_interface::Client> =
-        CapabilityServerSet::new();
-    let mut set2: CapabilityServerSet<impls::TestInterface, test_interface::Client> =
-        CapabilityServerSet::new();
+    use capnp_rpc::WeakCapabilityServerSet;
+    let mut set1: WeakCapabilityServerSet<impls::TestInterface, test_interface::Client> =
+        WeakCapabilityServerSet::new();
+    let mut set2: WeakCapabilityServerSet<impls::TestInterface, test_interface::Client> =
+        WeakCapabilityServerSet::new();
 
     let client_standalone = capnp_rpc::new_client(impls::TestInterface::new());
 
@@ -932,6 +932,10 @@ fn capability_server_set() {
     let own_server2_counter = own_server2.get_call_count();
     let client2 = set2.new_client(own_server2);
 
+    // gc() doesn't remove valid entries
+    set1.gc();
+    set2.gc();
+
     // Getting the local server using the correct set works.
     let own_server1_again = futures::executor::block_on(set1.get_local_server(&client1)).unwrap();
     assert_eq!(