Really leaky effect handling with 5.5

BRZO

;

leak.ml

(*---------------------------------------------------------------------------
   Copyright (c) 2026 The affect programmers. All rights reserved.
   SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

(* ocamlopt leak.ml && ./a.out

   On OCaml 5.5 on both MacOS and Linux the [live_stack_words]
   constantly increases when used with a total number of
   domains > 2.

   ./a.out -P 2 # no leak
   ./a.out -P 3 # starts leaking

   The main domain runs a for loop which on each iteration forks two fibers
   that just yield once and return. These fibers are picked up by busy waiting
   worker domains for execution. The main domain loop then reaps them by busy
   waiting and starts over. Every 10_000 iterations we Gc.compact and print
   stats that show the leak. Busy waiting is used to simplify the test case
   but the leak also shows up if thread sleeping with conditions is used.
*)

let rec atomic_update f a =
  let old = Atomic.get a in
  let new' = f old in
  if old == new' || Atomic.compare_and_set a old new' then () else
  (Domain.cpu_relax (); atomic_update f a)

let rec atomic_fold_update f a =
  let old = Atomic.get a in
  let new', v = f old in
  if old == new' || Atomic.compare_and_set a old new' then v else
  (Domain.cpu_relax (); atomic_fold_update f a)

module Synchronized_queue = struct
  type 'a t = 'a list Atomic.t
  let make () = Atomic.make []
  let is_empty q = List.is_empty (Atomic.get q)
  let add q v = atomic_update (fun q -> v :: q) q
  let take q =
    let take = function [] -> [], None | v :: vs -> List.rev vs, Some v in
    atomic_fold_update (fun q -> take (List.rev q)) q
end

type _ Effect.t +=
| Yield : unit Effect.t
| Fork : (unit -> unit) -> unit Effect.t
| Reap : unit Effect.t

let fork f = Effect.perform (Fork f)
let yield () = Effect.perform Yield
let reap () = Effect.perform Reap

let stop = Atomic.make false
let ready = Synchronized_queue.make ()
let reapable = Atomic.make 0
let incr_reapable () = Atomic.incr reapable

let busy_wait () = for i = 1 to 1000 do Domain.cpu_relax () done

let rec busy_run_next_ready () =
  assert (not (Domain.is_main_domain ()));
  match Synchronized_queue.take ready with
  | Some continue -> continue ()
  | None -> busy_wait (); if Atomic.get stop then () else busy_run_next_ready ()

let rec do_busy_reap () =
  assert (Domain.is_main_domain ());
  let count = Atomic.get reapable in
  if count = 0 then (busy_wait (); do_busy_reap ()) else
  if Atomic.compare_and_set reapable count (count - 1) then () else
  (Domain.cpu_relax (); do_busy_reap ())

(* Detect if we leak continuations ourselves manually *)

let to_continue = Atomic.make 0
let incr_to_continue () = Atomic.incr to_continue
let deep_continue k v = Atomic.decr to_continue; Effect.Deep.continue k v

let rec run f = match f () with
| () ->
    if Domain.is_main_domain ()
    then Atomic.set stop true
    else busy_run_next_ready ()
| effect Yield, k ->
    incr_to_continue ();
    Synchronized_queue.add ready (fun () -> deep_continue k ());
    busy_run_next_ready ()
| effect (Fork f), k ->
    incr_to_continue ();
    let run_f () = run (fun () -> f (); incr_reapable ()) in
    Synchronized_queue.add ready run_f;
    deep_continue k ()
| effect Reap, k ->
    incr_to_continue ();
    do_busy_reap ();
    deep_continue k ()

let log_status i ~repeat =
  Gc.compact ();
  let stat = Gc.stat () in
  Printf.printf
    "compact %d/%d live_stack_words: %d diff:%d to_continue: %d\n%!"
    i repeat
    stat.Gc.live_stacks_words (stat.Gc.live_stacks_words - stat.stack_size)
    (Atomic.get to_continue)

let leak ~domain_count ~repeat () =
  let spawn i = Domain.spawn (fun i -> run (fun () -> ())) in
  let workers = Iarray.init (domain_count - 1) spawn in
  begin run @@ fun () ->
    for i = 1 to repeat do
      let f = fun () -> yield () in
      fork f; fork f;
      reap (); reap ();
      assert (Atomic.get to_continue = 0);
      if i mod 10_000 = 0 then log_status i ~repeat;
    done;
  end;
  Iarray.iter Domain.join workers

let main () =
  let domain_count = ref None in
  let repeat = ref 50_000_000 in
  let parse_domain_count n =
    if n < 2
    then raise (Arg.Bad "Parallel count must be >= 2")
    else domain_count := Some n
  in
  let speclist =
    ["-R", Arg.Set_int repeat, "Loop count";
     "-P", Arg.Int parse_domain_count,
     "Set the total number of parallel domains (>= 2)" ]
  in
  let usage = "leak [-P COUNT] [-R COUNT]" in
  Arg.parse speclist (fun s -> raise (Arg.Bad "no pos arg supported")) usage;
  let domain_count = match !domain_count with
  | None -> Domain.recommended_domain_count ()
  | Some domain_count -> domain_count
  in
  leak ~domain_count ~repeat:!repeat ();
  0

let () = if !Sys.interactive then () else exit (main ())