Improve cancellation handling

capnp-rpc-lwt/service.ml

@@ -46,6 +46,10 @@ let local (s:#generic) =
         Payload.content_get p |> Schema.ReaderOps.deref_opt_struct_pointer |> Schema.ReaderOps.cast_struct in
       match m contents release_params with
       | r -> results#resolve r
+      | exception (Eio.Cancel.Cancelled _ as ex) ->
+        release_params ();
+        Core_types.resolve_payload results (Error `Cancelled);
+        raise ex
       | exception ex ->
         release_params ();
         Log.warn (fun f -> f "Uncaught exception handling %a: %a" pp_method (interface_id, method_id) Fmt.exn ex);

capnp-rpc-net/tls_wrapper.ml

@@ -17,7 +17,7 @@ let connect_as_server ~sw flow secret_key =
     let tls_config = Secret_key.tls_server_config key in
     match Tls_eio.server_of_flow tls_config flow with
     | exception (Failure msg) -> error "TLS connection failed: %s" msg
-    | exception ex -> error "TLS connection failed: %a" Fmt.exn ex
+    | exception ex -> Eio.Fiber.check (); error "TLS connection failed: %a" Fmt.exn ex
     | flow ->
       match Tls_eio.epoch flow with
       | Error () -> failwith "Unknown error getting TLS epoch data"
@@ -36,5 +36,5 @@ let connect_as_client ~sw flow secret_key auth =
     Log.info (fun f -> f "Doing TLS client-side handshake...");
     match Tls_eio.client_of_flow tls_config flow with
     | exception (Failure msg) -> error "TLS connection failed: %s" msg
-    | exception ex -> error "TLS connection failed: %a" Fmt.exn ex
+    | exception ex -> Eio.Fiber.check (); error "TLS connection failed: %a" Fmt.exn ex
     | flow -> Ok (Endpoint.of_flow ~sw ~peer_id:auth flow)

unix/network.ml

@@ -98,6 +98,7 @@ let connect net ~sw ~secret_key (addr, auth) =
     end;
     Tls_wrapper.connect_as_client ~sw socket secret_key auth
   | exception ex ->
+    Fiber.check ();
     error "@[<v2>Network connection for %a failed:@,%a@]" Location.pp addr Fmt.exn ex
 
 let accept_connection ~sw ~secret_key flow =