Fix some ref-counting errors in examples

- In `sturdy-refs-3`, `for_alice` wasn't explicitly released. If it gets
  GC'd while the connection is still up, it logs a warning. This was
  causing CI to fail in some cases.

- The restorer table takes ownership of the reference (back-port
  from Eio branch). Ensure vats are shutdown cleanly to make errors
  easier to find (Eio does this automatically).

README.md

@@ -939,9 +939,10 @@ the admin can request the sturdy ref like this:
 <!-- $MDX include,file=examples/sturdy-refs-3/main.ml,part=save -->
 ```ocaml
     (* The admin creates a logger for Alice and saves it: *)
-    let for_alice = Logger.sub root "alice" in
-    Persistence.save_exn for_alice >>= fun uri ->
-    Capnp_rpc_unix.Cap_file.save_uri uri "alice.cap" |> or_fail;
+    Capability.with_ref (Logger.sub root "alice") (fun for_alice ->
+        Persistence.save_exn for_alice >|= fun uri ->
+        Capnp_rpc_unix.Cap_file.save_uri uri "alice.cap" |> or_fail
+      ) >>= fun () ->
     (* Alice uses it: *)
     run_client "alice.cap"
 ```

examples/sturdy-refs-3/logger.ml

@@ -23,11 +23,13 @@ let rec local ~services sr label =
       let id = Capnp_rpc_net.Restorer.Id.generate () in
       let sr = Capnp_rpc_net.Restorer.Table.sturdy_ref services id in
       let sub = local ~services sr (Printf.sprintf "%s/%s" label sub_label) in
-      Capnp_rpc_net.Restorer.Table.add services id sub;
       let response, results = Service.Response.create Results.init_pointer in
       Results.logger_set results (Some sub);
-      Capability.dec_ref sub;
+      Capnp_rpc_net.Restorer.Table.add services id sub; (* Takes ownership of [sub] *)
       Service.return response
+
+    method! pp f =
+      Fmt.pf f "Logger(%s)" label
   end
 
 module Logger = Api.Client.Logger

examples/sturdy-refs-3/main.ml

@@ -14,10 +14,11 @@ let or_fail = function
   | Ok x -> x
   | Error (`Msg m) -> failwith m
 
-let start_server () =
+let start_server ~switch () =
   let config = Capnp_rpc_unix.Vat_config.create ~secret_key listen_address in
   let make_sturdy = Capnp_rpc_unix.Vat_config.sturdy_uri config in
   let services = Restorer.Table.create make_sturdy in
+  Lwt_switch.add_hook (Some switch) (fun () -> Restorer.Table.clear services; Lwt.return_unit);
   let restore = Restorer.of_table services in
   (* $MDX part-begin=root *)
   let root_id = Capnp_rpc_unix.Vat_config.derived_id config "root" in
@@ -27,27 +28,30 @@ let start_server () =
   in
   (* $MDX part-end *)
   Restorer.Table.add services root_id root;
-  Capnp_rpc_unix.serve config ~restore >|= fun _vat ->
+  Capnp_rpc_unix.serve ~switch config ~restore >|= fun _vat ->
   Capnp_rpc_unix.Vat_config.sturdy_uri config root_id
 
 let run_client cap_file =
-  let vat = Capnp_rpc_unix.client_only_vat () in
+  Lwt_switch.with_switch @@ fun switch ->
+  let vat = Capnp_rpc_unix.client_only_vat ~switch () in
   let sr = Capnp_rpc_unix.Cap_file.load vat cap_file |> or_fail in
   Sturdy_ref.with_cap_exn sr @@ fun for_alice ->
   Logger.log for_alice "Message from Alice"
 
 let () =
   Lwt_main.run begin
-    start_server () >>= fun root_uri ->
-    let vat = Capnp_rpc_unix.client_only_vat () in
+    Lwt_switch.with_switch @@ fun switch ->
+    start_server ~switch () >>= fun root_uri ->
+    let vat = Capnp_rpc_unix.client_only_vat ~switch () in
     let root_sr = Capnp_rpc_unix.Vat.import vat root_uri |> or_fail in
     Sturdy_ref.with_cap_exn root_sr @@ fun root ->
     Logger.log root "Message from Admin" >>= fun () ->
     (* $MDX part-begin=save *)
     (* The admin creates a logger for Alice and saves it: *)
-    let for_alice = Logger.sub root "alice" in
-    Persistence.save_exn for_alice >>= fun uri ->
-    Capnp_rpc_unix.Cap_file.save_uri uri "alice.cap" |> or_fail;
+    Capability.with_ref (Logger.sub root "alice") (fun for_alice ->
+        Persistence.save_exn for_alice >|= fun uri ->
+        Capnp_rpc_unix.Cap_file.save_uri uri "alice.cap" |> or_fail
+      ) >>= fun () ->
     (* Alice uses it: *)
     run_client "alice.cap"
     (* $MDX part-end *)

examples/sturdy-refs-4/logger.ml

@@ -31,6 +31,9 @@ let local ~persist_new sr label =
         Capability.dec_ref logger;
         Ok response
 (* $MDX part-end *)
+
+    method! pp f =
+      Fmt.pf f "Logger(%s)" label
   end
 
 module Logger = Api.Client.Logger

examples/sturdy-refs-4/main.ml

@@ -57,7 +57,7 @@ let sub cap_file label =
     let vat = Capnp_rpc_unix.client_only_vat () in
     let sr = Capnp_rpc_unix.Cap_file.load vat cap_file |> or_fail in
     Sturdy_ref.with_cap_exn sr @@ fun logger ->
-    let sub = Logger.sub logger label in
+    Capability.with_ref (Logger.sub logger label) @@ fun sub ->
     Persistence.save_exn sub >>= fun uri ->
     Capnp_rpc_unix.Cap_file.save_uri uri sub_file |> or_fail;
     Printf.printf "Wrote %S\n%!" sub_file;