diff --git a/src/fiber/src/cancel.ml b/src/fiber/src/cancel.ml index 5537544c354..f81f0f6106c 100644 --- a/src/fiber/src/cancel.ml +++ b/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)) ;; diff --git a/src/fiber/src/fiber.mli b/src/fiber/src/fiber.mli index c69688f30b3..29554261f17 100644 --- a/src/fiber/src/fiber.mli +++ b/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. *) diff --git a/src/fiber/test/cancel_tests.ml b/src/fiber/test/cancel_tests.ml index 83e38eb15d0..75910c3b8f5 100644 --- a/src/fiber/test/cancel_tests.ml +++ b/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 + |}] +;;