perf: add Fut.raise_if_failed, use it in Fiber.check_if_cancelled

src/core/fut.ml

@@ -29,6 +29,11 @@ let[@inline] peek self : _ option =
   | Done x -> Some x
   | Waiting _ -> None
 
+let[@inline] raise_if_failed self : unit =
+  match A.get self.st with
+  | Done (Error ebt) -> Exn_bt.raise ebt
+  | _ -> ()
+
 let[@inline] is_done self : bool =
   match A.get self.st with
   | Done _ -> true

src/core/fut.mli

@@ -91,6 +91,10 @@ val is_failed : _ t -> bool
 (** Checks if the future is resolved with [Error _] as a result.
     @since NEXT_RELEASE *)
 
+val raise_if_failed : _ t -> unit
+(** [raise_if_failed fut] raises [e] if [fut] failed with [e].
+    @since NEXT_RELEASE *)
+
 (** {2 Combinators} *)
 
 val spawn : on:Runner.t -> (unit -> 'a) -> 'a t

src/fib/fiber.ml

@@ -291,15 +291,18 @@ let with_on_self_cancel cb (k : unit -> 'a) : 'a =
 
 module Suspend_ = Moonpool.Private.Suspend_
 
+let[@inline] check_if_cancelled_ (self : _ t) = Fut.raise_if_failed self.res
+
 let check_if_cancelled () =
   match Task_local_storage.get k_current_fiber with
   | None ->
     failwith "Fiber.check_if_cancelled: must be run from inside a fiber."
-  | Some (Any self) ->
+  | Some (Any self) -> check_if_cancelled_ self
-    (match peek self with
-    | Some (Error ebt) -> Exn_bt.raise ebt
-    | _ -> ())
 
-let[@inline] yield () : unit =
+let yield () : unit =
-  check_if_cancelled ();
+  match Task_local_storage.get k_current_fiber with
-  Suspend_.yield ()
+  | None -> failwith "Fiber.yield: must be run from inside a fiber."
+  | Some (Any self) ->
+    check_if_cancelled_ self;
+    Suspend_.yield ();
+    check_if_cancelled_ self

src/fib/fiber.mli

@@ -81,7 +81,8 @@ val wait_block : 'a t -> 'a Fut.or_error
 val check_if_cancelled : unit -> unit
 (** Check if the current fiber is cancelled, in which case this raises.
     Must be run from inside a fiber.
-    @raise Failure if not. *)
+    @raise e if the current fiber is cancelled with exception [e]
+    @raise Failure if not run from a fiber. *)
 
 val yield : unit -> unit
 (** Yield control to the scheduler from the current fiber.