feature(fiber): add Cancel.make_child for linked cancellation

src/fiber/src/cancel.ml

@@ -13,15 +13,41 @@ type handlers =
       ; mutable prev : handlers
       }
 
-module State = struct
-  type t =
-    | Cancelled
-    | Not_cancelled of { mutable handlers : handlers }
-end
+type state =
+  | Cancelled
+  | Not_cancelled of
+      { mutable handlers : handlers
+      ; mutable children : t list
+      ; parent : t option
+      }
+
+and t = { mutable state : state }
 
-type t = { mutable state : State.t }
+let create () =
+  { state = Not_cancelled { handlers = End_of_handlers; children = []; parent = None } }
+;;
 
-let create () = { state = Not_cancelled { handlers = End_of_handlers } }
+let make_child parent =
+  match parent.state with
+  | Cancelled -> { state = Cancelled }
+  | Not_cancelled p ->
+    let child =
+      { state =
+          Not_cancelled
+            { handlers = End_of_handlers; children = []; parent = Some parent }
+      }
+    in
+    p.children <- child :: p.children;
+    child
+;;
+
+let remove_me_from_parent t ~parent =
+  match parent with
+  | None -> ()
+  | Some { state = Cancelled; _ } -> ()
+  | Some { state = Not_cancelled p; _ } ->
+    p.children <- List.filter (fun c -> not (c == t)) p.children
+;;
 
 let rec invoke_handlers = function
   | Handler { ivar; next; prev = _ } ->
@@ -30,13 +56,20 @@ let rec invoke_handlers = function
   | End_of_handlers -> return ()
 ;;
 
-let fire t =
-  of_thunk (fun () ->
+let fire =
+  let rec gather acc t =
     match t.state with
-    | Cancelled -> return ()
-    | Not_cancelled { handlers } ->
+    | Cancelled -> acc
+    | Not_cancelled { handlers; children; parent = _ } ->
       t.state <- Cancelled;
-      invoke_handlers handlers)
+      List.fold_left gather (handlers :: acc) children
+  in
+  fun t ->
+    of_thunk (fun () ->
+      (match t.state with
+       | Cancelled -> ()
+       | Not_cancelled { parent; _ } -> remove_me_from_parent t ~parent);
+      gather [] t |> List.rev |> parallel_iter ~f:invoke_handlers)
 ;;
 
 let rec fills_of_handlers acc = function
@@ -45,12 +78,13 @@ let rec fills_of_handlers acc = function
   | End_of_handlers -> List.rev acc
 ;;
 
-let fire' t =
+let rec fire' t =
   match t.state with
   | Cancelled -> []
-  | Not_cancelled { handlers } ->
+  | Not_cancelled { handlers; children; parent } ->
     t.state <- Cancelled;
-    fills_of_handlers [] handlers
+    remove_me_from_parent t ~parent;
+    fills_of_handlers [] handlers @ List.concat_map fire' children
 ;;
 
 let fired t =
@@ -63,7 +97,7 @@ let with_handler t f ~on_cancel =
   match t.state with
   | Cancelled ->
     let+ x, y = fork_and_join f on_cancel in
-    x, Cancelled y
+    x, (Cancelled y : _ outcome)
   | Not_cancelled h ->
     let ivar = Ivar.create () in
     let node = Handler { ivar; next = h.handlers; prev = End_of_handlers } in
@@ -97,6 +131,6 @@ let with_handler t f ~on_cancel =
          >>= function
          | Cancelled () ->
            let+ x = on_cancel () in
-           Cancelled x
-         | Not_cancelled -> return Not_cancelled)
+           (Cancelled x : _ outcome)
+         | Not_cancelled -> return (Not_cancelled : _ outcome))
 ;;

src/fiber/src/fiber.mli

@@ -522,6 +522,15 @@ module Cancel : sig
 
   val create : unit -> t
 
+  (** [make_child t] returns a new cancellation that is fired whenever [t] is
+      fired (but not vice versa).
+
+      Note that the order in which the handlers of a parent and its children run
+      is left unspecified. Note also that a child is not reclaimed until either
+      it or its parent is fired (or the parent is collected), so repeatedly
+      attaching children to a long-lived parent may leak memory. *)
+  val make_child : t -> t
+
   (** Activate a cancellation request.
 
       [fire] is idempotent, so calling [fire t] more than once has no effect. *)

src/fiber/test/cancel_tests.ml

@@ -71,3 +71,100 @@ let%expect_test "cancel_test4" =
      | Not_cancelled -> "FAIL");
   [%expect {| PASS |}]
 ;;
+
+let%expect_test "cancelling parent triggers child cancellation" =
+  let run () =
+    let parent = Fiber.Cancel.create () in
+    let child = Fiber.Cancel.make_child parent in
+    let parent_cancelled = Fiber.Ivar.create () in
+    let child_running = Fiber.Ivar.create () in
+    let child_cancelled = Fiber.Ivar.create () in
+    let+ (), res =
+      Fiber.Cancel.with_handler
+        parent
+        ~on_cancel:(fun () ->
+          print_endline "parent cancel";
+          Fiber.Ivar.fill parent_cancelled ())
+        (fun () ->
+          print_endline "with parent";
+          Fiber.fork_and_join_unit
+            (fun () ->
+              let+ (), res =
+                Fiber.Cancel.with_handler
+                  child
+                  ~on_cancel:(Fiber.Ivar.fill child_cancelled)
+                  (fun () ->
+                    let* () = Fiber.Ivar.fill child_running () in
+                    Fiber.Ivar.read child_cancelled)
+              in
+              match res with
+              | Not_cancelled -> assert false
+              | Cancelled () -> print_endline "child cancelled")
+            (fun () ->
+              let* () = Fiber.Ivar.read child_running in
+              let* () = Fiber.Cancel.fire parent in
+              Fiber.Ivar.read parent_cancelled))
+    in
+    match res with
+    | Not_cancelled -> assert false
+    | Cancelled () -> print_endline "parent cancelled"
+  in
+  Scheduler.run (Fiber.of_thunk run);
+  [%expect
+    {|
+    with parent
+    parent cancel
+    child cancelled
+    parent cancelled
+    |}]
+;;
+
+let%expect_test "child of cancelled" =
+  let run () =
+    let parent = Fiber.Cancel.create () in
+    let+ () = Fiber.Cancel.fire parent in
+    let child = Fiber.Cancel.make_child parent in
+    printf "child should be cancelled: %b" (Fiber.Cancel.fired child)
+  in
+  Scheduler.run (Fiber.of_thunk run);
+  [%expect {| child should be cancelled: true |}]
+;;
+
+let%expect_test "cancelling child does not cancel parent" =
+  let run () =
+    let parent = Fiber.Cancel.create () in
+    let child = Fiber.Cancel.make_child parent in
+    let+ (), res =
+      Fiber.Cancel.with_handler
+        parent
+        ~on_cancel:(fun () ->
+          print_endline "cancelling parent";
+          Fiber.return ())
+        (fun () ->
+          let* (), res =
+            Fiber.Cancel.with_handler
+              child
+              ~on_cancel:(fun () ->
+                print_endline "cancelling child";
+                Fiber.return ())
+              (fun () -> Fiber.Cancel.fire child)
+          in
+          match res with
+          | Not_cancelled -> assert false
+          | Cancelled () ->
+            print_endline "child cancelled";
+            Fiber.Cancel.fire parent)
+    in
+    match res with
+    | Not_cancelled -> assert false
+    | Cancelled () -> print_endline "parent cancelled"
+  in
+  Scheduler.run (Fiber.of_thunk run);
+  [%expect
+    {|
+    cancelling child
+    child cancelled
+    cancelling parent
+    parent cancelled
+    |}]
+;;