Make multififo threads periodically steal to balance load

Author: polytypic

lib/picos_aux.mpmcq/picos_aux_mpmcq.ml

@@ -186,18 +186,25 @@ let rec push_head t value backoff =
         end
     end
 
-let rec length t =
-  let head = Atomic.get t.head in
-  let tail = Atomic.fenceless_get t.tail in
-  if head != Atomic.get t.head then length t
-  else
-    let head_at =
-      match head with H (Cons r) -> r.counter | H (Head r) -> r.counter
-    in
-    let tail_at =
-      match tail with T (Snoc r) -> r.counter | T (Tail r) -> r.counter
-    in
-    tail_at - head_at + 1
+let[@inline] length t =
+  let t_head = t.head in
+  let t_tail = t.tail in
+  let head = ref (Atomic.get t_head) in
+  let tail = ref (Atomic.fenceless_get t_tail) in
+  while
+    head := Atomic.get t_head;
+    tail := Atomic.fenceless_get t_tail;
+    !head != Atomic.get t_head
+  do
+    ()
+  done;
+  let head_at =
+    match !head with H (Cons r) -> r.counter | H (Head r) -> r.counter
+  in
+  let tail_at =
+    match !tail with T (Snoc r) -> r.counter | T (Tail r) -> r.counter
+  in
+  tail_at - head_at + 1
 
 let[@inline] is_empty t = length t == 0
 let[@inline] pop_exn t = pop t Backoff.default (Atomic.fenceless_get t.head)

lib/picos_mux.multififo/picos_mux_multififo.ml

@@ -45,6 +45,7 @@ and _ tdt =
       mutable num_stopped : int;
       mutable fiber : Fiber.Maybe.t;
       mutable remaining_quota : int;
+      mutable countdown_to_steal : int;
     }
       -> [> `Per_thread ] tdt
 
@@ -116,7 +117,27 @@ let exec ready (Per_thread p : per_thread) t =
   | Resume (fiber, k) -> Fiber.resume fiber k
 
 let rec next (Per_thread p as pt : per_thread) =
-  match Mpmcq.pop_exn p.ready with
+  let ready =
+    let c = p.countdown_to_steal in
+    if 0 < c then begin
+      p.countdown_to_steal <- c - 1;
+      p.ready
+    end
+    else begin
+      let t = p.context in
+      if 1 < t.threads_num then begin
+        let (Per_thread v) = get_thread t (Random.int t.threads_num) in
+        p.countdown_to_steal <- Mpmcq.length p.ready + 1;
+        if p.countdown_to_steal <= Mpmcq.length v.ready then v.ready
+        else p.ready
+      end
+      else begin
+        p.countdown_to_steal <- 1_000;
+        p.ready
+      end
+    end
+  in
+  match Mpmcq.pop_exn ready with
   | ready ->
       let t = p.context in
       relaxed_wakeup t ~known_not_empty:false p.ready;
@@ -195,6 +216,7 @@ let per_thread context =
         num_stopped = 0;
         fiber = Fiber.Maybe.nothing;
         remaining_quota = 0;
+        countdown_to_steal = 1;
       }
   in
   p.resume <-

lib/picos_mux.multififo/picos_mux_multififo.mli

@@ -3,9 +3,11 @@
 
     This scheduler uses a queue per thread to implement a mostly FIFO scheduler.
     If a thread runs out of fibers to run, it will try to take a fiber from the
-    queues of other threads, which means that fibers can move from one thread to
-    another. This scheduler also gives priority to fibers woken up due to being
-    canceled.
+    queues of other threads. Furthermore, threads periodically consider taking
+    fibers from other threads to balance the number of fibers per thread. All of
+    this means that this scheduler should act relatively fairly and work well
+    for concurrent workloads and workloads where fairness matters. This
+    scheduler also gives priority to fibers woken up due to being canceled.
 
     🐌 Due to mostly FIFO scheduling this scheduler performs poorly on highly
     parallel workloads.