Leaky effect handling

BRZO

;

leak.ml

(* Example from the manual https://ocaml.org/manual/5.5/effects.html
   Modified with a fork that continues with the forker, not the forkee.
   This ~~implementation leaks~~
   No it doesn't you are just missing a yield at each iteration! See
   https://discuss.ocaml.org/t/leaking-effect-handling/18351

   ocamlopt leak.ml
   ./a.out  # look at your top
*)

open Effect
open Effect.Deep

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

let fork f = perform (Fork f)
let yield () = perform Yield
let xchg v = perform (Xchg v)

(* A concurrent round-robin scheduler *)
let run (main : unit -> unit) : unit =
  let exchanger : (int * (int, unit) continuation) option ref =
    ref None (* waiting exchanger *)
  in
  let run_q = Queue.create () in (* scheduler queue *)
  let enqueue_task task = Queue.push task run_q in
  let enqueue k v = enqueue_task (fun () -> continue k v) in
  let dequeue () =
    if Queue.is_empty run_q then () (* done *)
    else begin
      let task = Queue.pop run_q in
      task ()
    end
  in
  let rec spawn (f : unit -> unit) : unit =
    match f () with
    | () -> dequeue ()
    | exception e ->
        print_endline (Printexc.to_string e);
        dequeue ()
    | effect Yield, k -> enqueue k (); dequeue ()
    | effect (Fork f), k -> enqueue_task (fun () -> spawn f); continue k ()
    | effect (Xchg n), k ->
        begin match !exchanger with
        | Some (n', k') -> exchanger := None; enqueue k' n; continue k n'
        | None -> exchanger := Some (n, k); dequeue ()
        end
  in
  spawn main

let do_trace = false
let count = 100_000_000

let trace fmt =
  Printf.ksprintf (fun m -> if do_trace then print_endline m else ()) fmt

let () = run (fun () ->
    for i = 1 to count do
      fork (fun () ->
          trace "[t1] Sending 0";
          let v = xchg 0 in
          trace "[t1] received %d" v);
      fork (fun () ->
          trace "[t2] Sending 1";
          let v = xchg 1 in
          trace "[t2] received %d" v)
      yield ();
    done)