A possible alternative scheduling pattern for Eio that preserves effect handler semantics.

eio-plus-third-party-effect.ml

module Eio = struct
  type fiber = {
    fiber : (unit, unit) Effect.Shallow.continuation;
    mutable over : bool;
  }

  type t = { mutable fibers : fiber list }
  type _ Effect.t += Fork : fiber -> unit Effect.t | Yield : unit Effect.t

  let run f =
    let t = { fibers = [ { fiber = Effect.Shallow.fiber f; over = false } ] } in
    let rec step t =
      match t.fibers with
      | [] -> ()
      | fiber :: fibers ->
          let () =
            t.fibers <- fibers;
            Effect.Shallow.continue_with fiber.fiber ()
              {
                retc = (fun () -> fiber.over <- true);
                exnc = raise;
                effc =
                  (fun (type effect) (effect : effect Effect.t) ->
                    match effect with
                    | Fork fiber ->
                        Some
                          (fun (k : (effect, _) Effect.Shallow.continuation) ->
                            t.fibers <-
                              ({ fiber with fiber = k } :: t.fibers) @ [ fiber ])
                    | Yield ->
                        Some
                          (fun (k : (unit, unit) Effect.Shallow.continuation) ->
                            t.fibers <- t.fibers @ [ { fiber with fiber = k } ])
                    | _ -> None);
              }
          in
          step t
    in
    step t

  let yield () = Effect.perform Yield

  module Switch = struct
    type t = { mutable fibers : fiber list }

    let run f =
      let switch = { fibers = [] } in
      let res = f switch in
      let () =
        let rec join () =
          match List.find_opt (fun { over; _ } -> not over) switch.fibers with
          | None -> ()
          | Some _ -> yield ()
        in
        join ()
      in
      res
  end

  let fork ~sw f =
    let fiber = { fiber = Effect.Shallow.fiber f; over = false } in
    let () = sw.Switch.fibers <- fiber :: sw.Switch.fibers in
    Effect.perform (Fork fiber)
end

module Eio2 = struct
  type fiber = {
    fiber : (unit, unit) Effect.Shallow.continuation;
    mutable over : bool;
  }

  type t = { mutable fibers : fiber list }
  type _ Effect.t += Fork : fiber -> unit Effect.t | Yield : unit Effect.t

  let run f =
    let t = { fibers = [ { fiber = Effect.Shallow.fiber f; over = false } ] } in
    let rec step t =
      match t.fibers with
      | [] -> ()
      | fiber :: fibers ->
          let () =
            t.fibers <- fibers;
            Effect.Shallow.continue_with fiber.fiber ()
              {
                retc = (fun () -> fiber.over <- true);
                exnc = raise;
                effc =
                  (fun (type effect) (effect : effect Effect.t) ->
                    match effect with
                    | Fork fiber ->
                        Some
                          (fun (k : (effect, _) Effect.Shallow.continuation) ->
                            t.fibers <-
                              ({ fiber with fiber = k } :: t.fibers) @ [ fiber ])
                    | Yield ->
                        Some
                          (fun (k : (unit, unit) Effect.Shallow.continuation) ->
                            t.fibers <- t.fibers @ [ { fiber with fiber = k } ])
                    | _ -> None);
              }
          in
          step t
    in
    step t

  let yield () = Effect.perform Yield

  module Switch = struct
    type t = unit

    let run f = run f
  end

  let fork ~sw f =
    let fiber = { fiber = Effect.Shallow.fiber f; over = false } in
    Effect.perform (Fork fiber)
end

module Ping = struct
  type _ Effect.t += Ping : unit Effect.t

  let ping () = Effect.perform Ping

  let with_pong f =
    Effect.Deep.try_with f ()
      {
        effc =
          (fun (type effect) (effect : effect Effect.t) ->
            match effect with
            | Ping ->
                Some
                  (fun (k : (effect, _) Effect.Deep.continuation) ->
                    Format.eprintf "pong@.";
                    Effect.Deep.continue k ())
            | _ -> None);
      }
end

module Eio = Eio2

let () =
  Eio.run @@ fun () ->
  let test name =
    Format.eprintf "%s start@." name;
    Eio.yield ();
    Ping.ping ();
    Format.eprintf "%s finish@." name
  in
  let () =
    Ping.with_pong @@ fun () ->
    Eio.Switch.run @@ fun sw ->
    let () = Eio.fork ~sw (fun () -> test "forked") in
    test "main"
  in
  Format.eprintf "EOP@."