package domain-local-await

README.md

API reference

domain-local-await — Scheduler independent blocking

A low level mechanism intended for writing higher level libraries that need to block in a scheduler friendly manner.

A library that needs to suspend and later resume the current thread of execution may simply call prepare_for_await to obtain a pair of await and release operations for the purpose.

To provide an efficient and scheduler friendly implementation of the mechanism, schedulers may install an implementation by wrapping the scheduler main loop with a call to using. The implementation is then stored in a domain, and optionally thread, local variable. The overhead that this imposes on a scheduler should be insignificant.

An application can then choose to use schedulers that provide the necessary implementation. An implementation that works with plain domains and threads is provided as a default.

The end result is effective interoperability between schedulers and concurrent programming libraries.

Contents

• Example: Concurrency-safe lazy

• Example: Scheduler-friendly Mutex

• Example: Awaitable atomic locations

• Example: Transparently asynchronous IO

• References

Example: Concurrency-safe lazy

At the time of writing this, the documentation of the Stdlib Lazy module includes the following note:

Let's build a draft of a concurrency-safe version of lazy using atomics and domain-local-await!

First we need to require the library:

# #require "domain-local-await"

Here is a pair of types to represent the internal state of a lazy computation:

type 'a state =
  | Fun of (unit -> 'a)
  | Run of (unit -> unit) list
  | Val of 'a
  | Exn of exn

type 'a lazy_t = 'a state Atomic.t

A lazy computation starts as a thunk:

# let from_fun th = Atomic.make (Fun th)
val from_fun : (unit -> 'a) -> 'a state Atomic.t = <fun>

Or can be directly constructed with the given value:

# let from_val v = Atomic.make (Val v)
val from_val : 'a -> 'a state Atomic.t = <fun>

The interesting bits are in the force implementation:

# let rec force t =
    match Atomic.get t with
    | Val v -> v
    | Exn e -> raise e
    | Fun th as before ->
      if Atomic.compare_and_set t before (Run []) then
        let result =
          match th () with
          | v -> Val v
          | exception e -> Exn e
        in
        match Atomic.exchange t result with
        | (Val _ | Exn _ | Fun _) ->
          failwith "impossible"
        | Run waiters ->
          List.iter ((|>) ()) waiters;
          force t
      else
        force t
    | Run waiters as before ->
      let dla = Domain_local_await.prepare_for_await () in
      let after = Run (dla.release :: waiters) in
      if Atomic.compare_and_set t before after then
        match dla.await () with
        | () ->
          force t
        | exception cancelation_exn ->
          let rec cleanup () =
            match Atomic.get t with
            | (Val _ | Exn _ | Fun _) ->
              ()
            | Run waiters as before ->
              let after = Run (List.filter ((!=) dla.release) waiters) in
              if not (Atomic.compare_and_set t before after) then
                cleanup ()
          in
          cleanup ();
          raise cancelation_exn
      else
        force t
val force : 'a state Atomic.t -> 'a = <fun>

First force examines the state of the lazy computation. In case the result is already known, the value is returned or the exception is raised. Otherwise either the computation is started or the current thread of execution is suspended using domain-local-await. Once the thunk returns, the lazy is updated with the new state, any awaiters are released, and then all the force attempts will retry to examine the result. Notice also that the above force implementation is careful to perform a cleanup in case the await call raises an exception, which indicates cancellation.

Let's then try it by creating a lazy computation and forcing it from two different domains:

# let hello =
    from_fun (fun () ->
    Unix.sleepf 0.25;
    "Hello!")
val hello : string state Atomic.t = <abstr>

# let other = Domain.spawn (fun () -> force hello)
val other : string Domain.t = <abstr>

# force hello
- : string = "Hello!"

# Domain.join other
- : string = "Hello!"

Hello, indeed!

Note that the above implementation of lazy is intentionally kept relatively simple. It could be optimized slightly to reduce allocations and proper propagation of exception backtraces should be implemented. It could also be useful to have a scheduler independent mechanism to get a unique id corresponding to the current fiber, systhread, or domain and store that in the lazy state to be able to give an error in case of recursive forcing.

Example: Scheduler-friendly Mutex

At the time of writing this, the Stdlib Mutex implementation does not take into account the possibility of having an effects based scheduler and simply blocks the current domain (or (sys)thread) without giving a potential scheduler the opportunity to schedule another fiber on the domain.

Let's build a draft of a scheduler-friendly mutex using atomics and domain-local-await.

Here is a pair of types to represent a mutex:

type state =
  | Unlocked
  | Locked of (unit -> unit) list

type mutex = state Atomic.t

Essentially, a mutex is either unlocked or locked with a list of awaiters.

To construct a mutex we simply allocate a new atomic:

# let mutex () = Atomic.make Unlocked
val mutex : unit -> state Atomic.t = <fun>

The unlock operation just marks the mutex as unlocked and then wakes up all the awaiters:

# let rec unlock t =
    match Atomic.exchange t Unlocked with
    | Unlocked -> invalid_arg "mutex: already unlocked"
    | Locked awaiters -> List.iter ((|>) ()) awaiters
val unlock : state Atomic.t -> unit = <fun>

The lock operation is more complex:

# let rec lock t =
    match Atomic.get t with
    | Unlocked ->
      if not (Atomic.compare_and_set t Unlocked (Locked [])) then
        lock t
    | Locked awaiters as before ->
      let dla = Domain_local_await.prepare_for_await () in
      let after = Locked (dla.release :: awaiters) in
      if Atomic.compare_and_set t before after then
        match dla.await () with
        | () -> lock t
        | exception cancellation_exn ->
          let rec cleanup () =
            match Atomic.get t with
            | Unlocked -> ()
            | Locked awaiters as before ->
              if List.for_all ((==) dla.release) awaiters then
                let after =
                  Locked (List.filter ((!=) dla.release) awaiters)
                in
                if not (Atomic.compare_and_set t before after) then
                  cleanup ()
          in
          cleanup ();
          raise cancellation_exn
      else
        lock t
val lock : state Atomic.t -> unit = <fun>

In case the mutex is already locked, domain-local-await is used to await until the mutex is unlocked and the corresponding release is called. In case await raises, unlock makes sure to remove the release operation from the mutex to avoid a potential space leak.

Let's then use the mutex in a simple example of increment a counter from multiple domains:

# let mutex = mutex ()
val mutex : state Atomic.t = <abstr>

# let counter = ref 0
val counter : int ref = {contents = 0}

# let domains = List.init 3 @@ fun _ ->
    Domain.spawn @@ fun () ->
    for _ = 1 to 10000 do
      lock mutex;
      incr counter;
      unlock mutex;
    done
val domains : unit Domain.t list = [<abstr>; <abstr>; <abstr>]

# List.iter Domain.join domains
- : unit = ()

# !counter
- : int = 30000

Note that, like with the previous lazy implementation, the above mutex implementation is intentionally kept relatively simple and can be improved in various ways. It would make sense to use a backoff in case of contention. The representation could also be optimized to reduce memory usage. The above mutex implementation is also unfair.

Example: Awaitable atomic locations

Let's implement a simple awaitable atomic location abstraction. An awaitable location contains both the current value of the location and a list of awaiters, which are just unit -> unit functions:

type 'a awaitable_atomic = ('a * (unit -> unit) list) Atomic.t

The constructor of awaitable locations just pairs the initial value with an empty list of awaiters:

# let awaitable_atomic v : _ awaitable_atomic = Atomic.make (v, [])
val awaitable_atomic : 'a -> 'a awaitable_atomic = <fun>

Operations that modify awaitable locations, like fetch_and_add, need to call the awaiters to wake them up after a successful modification:

# let rec fetch_and_add x n =
    let (i, awaiters) as was = Atomic.get x in
      if Atomic.compare_and_set x was (i+n, []) then begin
          List.iter ((|>) ()) awaiters;
          i
        end
      else
        fetch_and_add x n
val fetch_and_add : (int * (unit -> unit) list) Atomic.t -> int -> int =
  <fun>

We can also have read-only operations, like get_as, that can be used to await for an awaitable location to have a specific value:

# let rec get_as fn x =
    let (v, awaiters) as was = Atomic.get x in
    match fn v with
    | Some w -> w
    | None ->
      let dla = Domain_local_await.prepare_for_await () in
      if Atomic.compare_and_set x was (v, dla.release :: awaiters) then
        match dla.await () with
        | () -> get_as fn x
        | exception cancelation_exn ->
          let rec cleanup () =
            let (w, awaiters) as was = Atomic.get x in
            if v == w then
              let awaiters = List.filter ((!=) dla.release) awaiters in
              if not (Atomic.compare_and_set x was (w, awaiters))
              then cleanup ()
          in
          cleanup ();
          raise cancelation_exn
      else
        get_as fn x
val get_as : ('a -> 'b option) -> ('a * (unit -> unit) list) Atomic.t -> 'b =
  <fun>

Notice that we carefully cleaned up in case the await was canceled.

We could, of course, also have operations that potentially awaits for the location to have an acceptable value before attempting modification. Let's leave that as an exercise.

To test awaitable locations, let's first create a location:

# let x = awaitable_atomic 0
val x : int awaitable_atomic = <abstr>

And let's then create a thread that awaits until the value of the location has changed and then modifies the value of the location:

# let a_thread =
    ()
    |> Thread.create @@ fun () ->
       get_as (fun x -> if x = 0 then None else Some ()) x;
       fetch_and_add x 21 |> ignore
val a_thread : Thread.t = <abstr>

The other thread is now awaiting for the initial modification:

# assert (0 = fetch_and_add x 21)
- : unit = ()

And we can await for the thread to perform its modification:

# get_as (fun x -> if x <> 21 then Some x else None) x;
- : int = 42

Let's then finish by joining with the other thread:

# Thread.join a_thread
- : unit = ()

Example: Transparently asynchronous IO

As a final example, let's sketch out an implementation of something a bit more involved — transparently asynchronous IO. The idea is that we implement operations such as read and write on Unix file descriptors in such a way that they block in a scheduler friendly manner allowing other fibers to run while waiting for the IO.

But first, we want to perform certain operations atomically. For that purpose we extend the Atomic module with a couple of helpers:

module Atomic = struct
  include Stdlib.Atomic

  let rec update t fn =
    let before = Atomic.get t in
    let after = fn before in
    if Atomic.compare_and_set t before after then
      before
    else
      update t fn

  let modify t fn = update t fn |> ignore
end

Below is the asynchronous IO module. It exposes read, write, and accept operations on Unix file descriptors. The operations block in a scheduler friendly manner. The implementation automatically manages a systhread per domain that runs a select loop, which takes care of awaiting for IO operations to be immediately executable. The operations on file descriptors communicate with the select loop thread.

module Async_io : sig
  open Unix
  val read : file_descr -> bytes -> int -> int -> int
  val write : file_descr -> bytes -> int -> int -> int
  val accept : ?cloexec:bool -> file_descr -> file_descr * sockaddr
end = struct
  module Awaiter = struct
    type t = { file_descr : Unix.file_descr; release : unit -> unit }

    let file_descr_of t = t.file_descr

    let rec signal aws file_descr =
      match aws with
      | [] -> ()
      | aw :: aws ->
          if aw.file_descr == file_descr then
            aw.release ()
          else signal aws file_descr

    let signal_or_wakeup wakeup aws file_descr =
      if file_descr == wakeup then begin
        let n = Unix.read file_descr (Bytes.create 1) 0 1 in
        assert (n = 1)
      end
      else signal aws file_descr

    let reject file_descr =
      List.filter (fun aw -> aw.file_descr != file_descr)
  end

  type state = {
    mutable state : [ `Init | `Locked | `Alive | `Dead ];
    mutable pipe_out : Unix.file_descr;
    reading : Awaiter.t list Atomic.t;
    writing : Awaiter.t list Atomic.t;
  }

  let key =
    Domain.DLS.new_key @@ fun () -> {
      state = `Init;
      pipe_out =
        (* Unfortunately we cannot safely allocate a pipe here,
           so we use stdin as a dummy value. *)
        Unix.stdin;
      reading = Atomic.make [];
      writing = Atomic.make [];
    }

  let[@poll error] try_lock s =
    s.state == `Init && begin
      s.state <- `Locked;
      true
    end

  let needs_init s =
    s.state != `Alive

  let[@poll error] unlock s pipe_out =
    s.pipe_out <- pipe_out;
    s.state <- `Alive

  let wakeup s =
    let n = Unix.write s.pipe_out (Bytes.create 1) 0 1 in
    assert (n = 1)

  let rec init s =
    (* DLS initialization may be run multiple times, so we
       perform more involved initialization here. *)
    if try_lock s then begin
      (* The pipe is used to wake up the select after changing
         the lists of reading and writing file descriptors. *)
      let pipe_inn, pipe_out = Unix.pipe ~cloexec:true () in
      unlock s pipe_out;
      let t =
        ()
        |> Thread.create @@ fun () ->
           (* This is the IO select loop that performs select and
              then wakes up fibers blocked on IO. *)
           while s.state != `Dead do
             let rs, ws, _ =
               Unix.select
                 (pipe_inn
                  :: List.map Awaiter.file_descr_of (Atomic.get s.reading))
                 (List.map Awaiter.file_descr_of (Atomic.get s.writing))
                 []
                 (-1.0)
             in
             List.iter
               (Awaiter.signal_or_wakeup pipe_inn (Atomic.get s.reading))
               rs;
             List.iter (Awaiter.signal (Atomic.get s.writing)) ws;
             Atomic.modify s.reading (List.fold_right Awaiter.reject rs);
             Atomic.modify s.writing (List.fold_right Awaiter.reject ws);
         done;
         Unix.close pipe_inn;
         Unix.close pipe_out
      in
      Domain.at_exit @@ fun () ->
        s.state <- `Dead;
        wakeup s;
        Thread.join t
    end
    else if needs_init s then begin
      Thread.yield ();
      init s;
    end

  let get () =
    let s = Domain.DLS.get key in
    if needs_init s then
      init s;
    s

  let await s r file_descr =
    let Domain_local_await.{ await; release } =
      Domain_local_await.prepare_for_await ()
    in
    let awaiter = Awaiter.{ file_descr; release } in
    Atomic.modify r (List.cons awaiter);
    wakeup s;
    try await ()
    with cancellation_exn ->
      Atomic.modify r (List.filter ((!=) awaiter));
      raise cancellation_exn

  let read file_descr bytes pos len =
    let s = get () in
    await s s.reading file_descr;
    Unix.read file_descr bytes pos len

  let write file_descr bytes pos len =
    let s = get () in
    await s s.writing file_descr;
    Unix.write file_descr bytes pos len

  let accept ?cloexec file_descr =
    let s = get () in
    await s s.reading file_descr;
    Unix.accept ?cloexec file_descr
end

To demonstrate that we can perform IO operations without blocking the thread we implement a very minimalistic effects based toy scheduler. We could also use any existing scheduler that provides support for domain-local-await (see).

module Toy_scheduler : sig
  val fiber : (unit -> unit) -> unit
  val run : (unit -> unit) -> unit
end = struct
  type _ Effect.t +=
    | Suspend : (('a, unit) Effect.Deep.continuation -> unit) -> 'a Effect.t

  let ready = Atomic.make []
  let num_alive_fibers = ref 0

  let fiber thunk =
    incr num_alive_fibers;
    let thunk () =
      thunk ();
      decr num_alive_fibers
    in
    Atomic.modify ready (List.cons thunk)

  let run program =
    let needs_wakeup = Atomic.make false in
    let pipe_inn, pipe_out = Unix.pipe ~cloexec:true () in
    let rec scheduler () =
      match Atomic.update ready (function [] -> [] | _::xs -> xs) with
      | work::_ ->
        let effc (type a) : a Effect.t -> _ = function
          | Suspend ef -> Some ef
          | _ -> None in
        Effect.Deep.try_with work () { effc };
        scheduler ()
      | [] ->
        if !num_alive_fibers <> 0 then begin
          if Atomic.get needs_wakeup then
            (* There are blocked fibers, so we wait for them to
               become unblocked. *)
            let _ = Unix.select [pipe_inn] [] [] (-1.0) in
            let n = Unix.read pipe_inn (Bytes.create 1) 0 1 in
            assert (n = 1)
          else
            (* There are blocked fibers, so we need to wait for
               them to become ready.  But we need to check the
               ready list once more before we do so. *)
            Atomic.set needs_wakeup true;
          scheduler ()
        end
    in
    let prepare_for_await _ =
      let state = Atomic.make `Init in
      let release () =
        if Atomic.get state != `Released then
          match Atomic.exchange state `Released with
          | `Awaiting k ->
            let thunk = Effect.Deep.continue k in
            Atomic.modify ready (List.cons thunk);
            if Atomic.get needs_wakeup &&
               Atomic.compare_and_set needs_wakeup true false then
              (* The scheduler is potentially waiting on select,
                 so we need to perform a wakeup. *)
              let n = Unix.write pipe_out (Bytes.create 1) 0 1 in
              assert (n = 1)
          | _ -> () in
      let await () =
        if Atomic.get state != `Released then
          Effect.perform @@ Suspend (fun k ->
            if not (Atomic.compare_and_set state `Init (`Awaiting k)) then
              Effect.Deep.continue k ())
      in
      Domain_local_await.{ release; await } in
    Domain_local_await.using
      ~prepare_for_await
      ~while_running:(fun () ->
        incr num_alive_fibers;
        let program () =
          program ();
          decr num_alive_fibers
        in
        Atomic.modify ready (List.cons program);
        scheduler ())
end

The toy scheduler and the async IO implementation do not depend on each other and, more generally, know nothing about each other. They simply interoperate through the use of domain-local-await!

Finally here is an example program that runs a client and a server fiber that communicate through sockets:

# Toy_scheduler.run @@ fun () ->

  let n = 100 in
  let port =
    Random.self_init ();
    Random.int 1000 + 3000
  in
  let server_addr = Unix.ADDR_INET (Unix.inet_addr_loopback, port) in

  let () =
    Toy_scheduler.fiber @@ fun () ->
    Printf.printf "  Client running\n%!";
    let socket = Unix.socket ~cloexec:true PF_INET SOCK_STREAM 0 in
    Fun.protect ~finally:(fun () -> Unix.close socket) @@ fun () ->
    Unix.connect socket server_addr;
    Printf.printf "  Client connected\n%!";
    let bytes = Bytes.create n in
    let n = Async_io.write socket bytes 0 (Bytes.length bytes) in
    Printf.printf "  Client wrote %d\n%!" n;
    let n = Async_io.read socket bytes 0 (Bytes.length bytes) in
    Printf.printf "  Client read %d\n%!" n
  in

  let () =
    Toy_scheduler.fiber @@ fun () ->
    Printf.printf "  Server running\n%!";
    let client, _client_addr =
      let socket = Unix.socket ~cloexec:true PF_INET SOCK_STREAM 0 in
      Fun.protect ~finally:(fun () -> Unix.close socket) @@ fun () ->
      Unix.set_nonblock socket;
      Unix.bind socket server_addr;
      Unix.listen socket 1;
      Printf.printf "  Server listening\n%!";
      Async_io.accept ~cloexec:true socket
    in
    Fun.protect ~finally:(fun () -> Unix.close client) @@ fun () ->
    Unix.set_nonblock client;
    let bytes = Bytes.create n in
    let n = Async_io.read client bytes 0 (Bytes.length bytes) in
    Printf.printf "  Server read %d\n%!" n;
    let n = Async_io.write client bytes 0 (n / 2) in
    Printf.printf "  Server wrote %d\n%!" n
  in

  Printf.printf "Client server test\n%!"
Client server test
  Server running
  Server listening
  Client running
  Client connected
  Client wrote 100
  Server read 100
  Server wrote 50
  Client read 50
- : unit = ()

This proof-of-concept shows that using just domain-local-await and a systhread we can implement scheduler agnostic transparently asynchronous IO. There is a lot of room for optimizations and other kinds of improvements.

References

DLA is used to implement blocking operations by the following libraries:

• Kcas

DLA support is provided by the following schedulers:

• Eio ^{(>= 0.10)}

• Domainslib ^{(>= 0.5.1)}

• Moonpool ^{(>= 0.3)}