Fix ref-counting error in example

The restorer table takes ownership of the reference (back-port from Eio
branch).

Author: talex5

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,18 +28,20 @@ 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 ->
+    Lwt_switch.with_switch @@ fun switch ->
+    start_server ~switch () >>= fun root_uri ->
     let vat = Capnp_rpc_unix.client_only_vat () in
     let root_sr = Capnp_rpc_unix.Vat.import vat root_uri |> or_fail in
     Sturdy_ref.with_cap_exn root_sr @@ fun root ->

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