Restorer.Table.create now takes a switch

The table is now cleared when the switch finishes. Avoids having to
remember to set this up manually.

Also, no longer yield when returning a cached ref. Seems unnecessary,
and could cause a leak if cancelled.

README.md

@@ -846,7 +846,7 @@ let make_service ~config ~services name =
 let start_server ~sw net =
   let config = Capnp_rpc_unix.Vat_config.create ~secret_key ~net listen_address in
   let make_sturdy = Capnp_rpc_unix.Vat_config.sturdy_uri config in
-  let services = Restorer.Table.create make_sturdy in
+  let services = Restorer.Table.create ~sw make_sturdy in
   let restore = Restorer.of_table services in
   let services = List.map (make_service ~config ~services) ["alice"; "bob"] in
   let vat = Capnp_rpc_unix.serve ~sw ~restore config in

capnp-rpc-net/capnp_rpc_net.mli

@@ -104,15 +104,17 @@ module Restorer : sig
     type t
     (** A restorer that keeps a hashtable mapping IDs to capabilities in memory. *)
 
-    val create : (Id.t -> Uri.t) -> t
-    (** [create make_sturdy] is a new in-memory-only table.
+    val create : sw:Eio.Switch.t -> (Id.t -> Uri.t) -> t
+    (** [create ~sw make_sturdy] is a new in-memory-only table.
         [make_sturdy id] converts an ID to a full URI, by adding the
-        hosting vat's address and fingerprint. *)
+        hosting vat's address and fingerprint.
+        @param sw {!clear} is called when [sw] is turned off. *)
 
     val of_loader : sw:Eio.Switch.t -> (module LOADER with type t = 'loader) -> 'loader -> t
     (** [of_loader ~sw (module Loader) l] is a new caching table that uses
         [Loader.load l sr (Loader.hash id)] to restore services that aren't in the cache.
-        The load function runs in a new fiber in [sw]. *)
+        @param sw The load function runs in a new fiber in [sw].
+                  {!clear} is called when [sw] is turned off. *)
 
     val add : t -> Id.t -> 'a Capability.t -> unit
     (** [add t id cap] adds a mapping to [t].

capnp-rpc-net/restorer.ml

@@ -78,6 +78,7 @@ module Table = struct
     | Manual of Core_types.cap          (* We hold a ref on the cap *)
 
   type t = {
+    sw : Switch.t;
     hash : Digestif.hash';
     cache : (digest, entry) Hashtbl.t;
     load : Id.t -> digest -> resolution Promise.or_exn;
@@ -86,11 +87,21 @@ module Table = struct
 
   (* [cache] contains promises or capabilities with positive ref-counts. *)
 
-  let create make_sturdy =
+  let release = function
+    | Manual cap -> Core_types.dec_ref cap;
+    | Cached _ -> ()
+
+  let clear t =
+    Hashtbl.iter (fun _ v -> release v) t.cache;
+    Hashtbl.clear t.cache
+
+  let create ~sw make_sturdy =
     let hash = `SHA256 in
     let cache = Hashtbl.create 53 in
     let load _ _ = Promise.create_resolved (Ok unknown_service_id) in
-    { hash; cache; load; make_sturdy }
+    let t = { sw; hash; cache; load; make_sturdy } in
+    Switch.on_release sw (fun () -> clear t);
+    t
 
   let hash t id =
     Id.digest t.hash id
@@ -102,13 +113,9 @@ module Table = struct
       Core_types.inc_ref cap;
       Ok cap
     | Cached res ->
-      begin match Promise.await_exn res with
-        | Error _ as e -> e
-        | Ok cap ->
-          Core_types.inc_ref cap;
-          Fiber.yield ();
-          Ok cap
-      end
+      let res = Promise.await_exn res in
+      Result.iter Core_types.inc_ref res;
+      res
     | exception Not_found ->
       let cap = t.load id digest in
       Hashtbl.add t.cache digest (Cached cap);
@@ -138,13 +145,15 @@ module Table = struct
           end in
           L.load loader (Cast.sturdy_of_raw sr) digest
         )
-    and t = { hash; cache; load; make_sturdy = L.make_sturdy loader } in
+    and t = { sw; hash; cache; load; make_sturdy = L.make_sturdy loader } in
+    Switch.on_release sw (fun () -> clear t);
     t
 
   let add t id cap =
     let cap = Cast.cap_to_raw cap in
     let id = hash t id in
     assert (not (Hashtbl.mem t.cache id));
+    Switch.check t.sw;
     Hashtbl.add t.cache id (Manual cap)
 
   let sturdy_ref t id =
@@ -153,21 +162,13 @@ module Table = struct
       method to_uri_with_secrets = t.make_sturdy id
     end
 
-  let release = function
-    | Manual cap -> Core_types.dec_ref cap;
-    | Cached _ -> ()
-
   let remove t id =
     let id = hash t id in
     match Hashtbl.find t.cache id with
     | exception Not_found -> failwith "Service ID not in restorer table"
     | value ->
       release value;
       Hashtbl.remove t.cache id
-
-  let clear t =
-    Hashtbl.iter (fun _ v -> release v) t.cache;
-    Hashtbl.clear t.cache
 end
 
 let of_table = Table.resolve

examples/sturdy-refs-2/main.ml

@@ -17,7 +17,7 @@ let or_fail = function
 let start_server ~sw net =
   let config = Capnp_rpc_unix.Vat_config.create ~secret_key ~net listen_address in
   let make_sturdy = Capnp_rpc_unix.Vat_config.sturdy_uri config in
-  let services = Restorer.Table.create make_sturdy in
+  let services = Restorer.Table.create ~sw make_sturdy in
   let restore = Restorer.of_table services in
   let root_id = Capnp_rpc_unix.Vat_config.derived_id config "root" in
   let root = Logger.local "root" in

examples/sturdy-refs-3/main.ml

@@ -17,8 +17,7 @@ let or_fail = function
 let start_server ~sw net =
   let config = Capnp_rpc_unix.Vat_config.create ~secret_key ~net listen_address in
   let make_sturdy = Capnp_rpc_unix.Vat_config.sturdy_uri config in
-  let services = Restorer.Table.create make_sturdy in
-  Switch.on_release sw (fun () -> Restorer.Table.clear services);
+  let services = Restorer.Table.create ~sw make_sturdy in
   let restore = Restorer.of_table services in
   (* $MDX part-begin=root *)
   let root_id = Capnp_rpc_unix.Vat_config.derived_id config "root" in

examples/sturdy-refs/main.ml

@@ -24,7 +24,7 @@ let make_service ~config ~services name =
 let start_server ~sw net =
   let config = Capnp_rpc_unix.Vat_config.create ~secret_key ~net listen_address in
   let make_sturdy = Capnp_rpc_unix.Vat_config.sturdy_uri config in
-  let services = Restorer.Table.create make_sturdy in
+  let services = Restorer.Table.create ~sw make_sturdy in
   let restore = Restorer.of_table services in
   let services = List.map (make_service ~config ~services) ["alice"; "bob"] in
   let vat = Capnp_rpc_unix.serve ~sw ~restore config in

test/test.ml

@@ -403,7 +403,7 @@ let except_ty = Alcotest.testable Capnp_rpc.Exception.pp_ty (=)
 let test_table_restorer ~net:_ =
   Switch.run @@ fun sw ->
   let make_sturdy id = Uri.make ~path:(Restorer.Id.to_string id) () in
-  let table = Restorer.Table.create make_sturdy in
+  let table = Restorer.Table.create ~sw make_sturdy in
   let echo_id = Restorer.Id.public "echo" in
   let registry_id = Restorer.Id.public "registry" in
   let broken_id = Restorer.Id.public "broken" in