Using `[@poll error]` attribute to implement systhread safe data structures

article.md

For a long time OCaml has supported lightweight threads exposed via the Thread module. These threads are often called “systhreads”, but I will simply call them “threads” in this post.

The OCaml Stdlib also provides many mutable data structures such as Hashtbl, Queue, and Stack. As the documentation alerts, none of these are thread-safe.

In this post I will very briefly describe an approach to implementing lock-free thread-safe data structures.

In OCaml 4 and in OCaml 5, within a single Domain, only a single thread may run at a time. In other words, threads do not run in parallel except when they run in different domains in OCaml 5. The OCaml runtime schedules threads and semi pre-emptively switches (within a domain) between threads (created within the domain) during “safe points”. In other words, thread switches cannot happen at arbitrary points — they may only happen at safe points. Memory allocations are safe points. Additional safe points (where no actual memory allocation happens) are inserted into various constructs such as loops.

This means that within a block of code where there are no safe points it is possible to make multiple read and write accesses atomically with respect to threads (within the domain).

How does one ensure that a block of code does not include a safe point?

The OCaml compiler provides an annotation [@poll error] that one can use on a function to ensure that the function does not include a safe point.

IOW, using [@poll error] one can essentially create functions that are executed atomically with respect to threads (within a domain).

With a particular application in mind, I have created a lock-free thread-safe (integer keyed) hash table, thread-table.

As mentioned in the README, the implementation has “zero synchronization overhead on lookups”. Indeed, if you look at the find operation implementation

let find t k' =
  let h = Mix.int k' in
  let buckets = t.buckets in
  let n = Array.length buckets in
  let i = h land (n - 1) in
  find k' (Array.unsafe_get buckets i)

it includes no synchronization. In this case not even a [@poll error] attribute is needed.

For other operations the thread-table implementation uses functions annotated with the [@poll error] attribute (to make atomic updates) and familiar lock-free programming patterns such as retry loops and cooperation to avoid starvation. As an example, see the add operation implementation:

let[@poll error] add_atomically t buckets n i before after =
  t.rehash = 0 && buckets == t.buckets
  && before == Array.unsafe_get buckets i
  && begin
       Array.unsafe_set buckets i after;
       let length = t.length + 1 in
       t.length <- length;
       if n < length && n < max_buckets_div_2 then t.rehash <- n * 2;
       true
     end

let rec add t k' v' =
  let h = Mix.int k' in
  maybe_rehash t;
  let buckets = t.buckets in
  let n = Array.length buckets in
  let i = h land (n - 1) in
  let before = Array.unsafe_get buckets i in
  let after = Cons (k', v', before) in
  if not (add_atomically t buckets n i before after) then add t k' v'

Compared to e.g. using a Stdlib Mutex to protect a data structure against concurrent accesses by threads, this sort of lock-free implementation can give better performance (especially for read-only operations) and also allows use of the operations in contexts, such as signal handlers, where locks are not appropriate.

Note that this technique is not sufficient for parallelism-safe implementation of data structures.