Number of alive fibers need not be atomic in the FIFO scheduler

Author: polytypic

lib/picos_mux.fifo/picos_mux_fifo.ml

@@ -2,8 +2,6 @@ open Picos
 
 let[@inline never] quota_non_positive () = invalid_arg "quota must be positive"
 
-(* As a minor optimization, we avoid allocating closures, which take slightly
-   more memory than values of this type. *)
 type ready =
   | Spawn of Fiber.t * (Fiber.t -> unit)
   | Continue of Fiber.t * (unit, unit) Effect.Deep.continuation
@@ -17,7 +15,6 @@ module Mpscq = Picos_aux_mpscq
 type t = {
   ready : ready Mpscq.t;
   needs_wakeup : bool Atomic.t;
-  num_alive_fibers : int Atomic.t;
   mutex : Mutex.t;
   condition : Condition.t;
   resume :
@@ -33,6 +30,7 @@ type t = {
   quota : int;
   mutable fiber : Fiber.Maybe.t;
   mutable remaining_quota : int;
+  mutable num_alive_fibers : int;
 }
 
 let rec next t =
@@ -55,15 +53,11 @@ let rec next t =
       Fiber.resume fiber k
   | exception Mpscq.Empty ->
       t.fiber <- Fiber.Maybe.nothing;
-      if Atomic.get t.num_alive_fibers <> 0 then begin
+      if t.num_alive_fibers <> 0 then begin
         if Atomic.get t.needs_wakeup then begin
           Mutex.lock t.mutex;
           match
-            if Atomic.get t.needs_wakeup then
-              (* We assume that there is no poll point after the above
-                 [Mutex.lock] and before the below [Condition.wait] is ready to
-                 be woken up by a [Condition.broadcast]. *)
-              Condition.wait t.condition t.mutex
+            if Atomic.get t.needs_wakeup then Condition.wait t.condition t.mutex
           with
           | () -> Mutex.unlock t.mutex
           | exception exn ->
@@ -85,15 +79,13 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
   Select.check_configured ();
   let ready = Mpscq.create ~padded:true ()
   and needs_wakeup = Atomic.make false |> Multicore_magic.copy_as_padded
-  and num_alive_fibers = Atomic.make 1 |> Multicore_magic.copy_as_padded
   and mutex = Mutex.create ()
   and condition = Condition.create () in
   let rec t =
     {
       ready;
       fiber = Fiber.Maybe.nothing;
       needs_wakeup;
-      num_alive_fibers;
       mutex;
       condition;
       resume;
@@ -104,11 +96,9 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
       handler;
       quota;
       remaining_quota = quota;
+      num_alive_fibers = 1;
     }
   and current =
-    (* The current handler must never propagate cancelation, but it would be
-       possible to continue some other fiber and resume the current fiber
-       later. *)
     Some
       (fun k ->
         let fiber = Fiber.Maybe.to_fiber t.fiber in
@@ -140,24 +130,17 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
   and[@alert "-handler"] effc :
       type a. a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option =
     function
-    | Fiber.Current ->
-        (* We handle [Current] first as it is perhaps the most latency
-           sensitive effect. *)
-        t.current
+    | Fiber.Current -> t.current
     | Fiber.Spawn r ->
-        (* We check cancelation status once and then either perform the
-           whole operation or discontinue the fiber. *)
         let fiber = Fiber.Maybe.to_fiber t.fiber in
         if Fiber.is_canceled fiber then t.discontinue
         else begin
-          Atomic.incr t.num_alive_fibers;
+          t.num_alive_fibers <- t.num_alive_fibers + 1;
           Mpscq.push t.ready (Spawn (r.fiber, r.main));
           t.return
         end
     | Fiber.Yield -> t.yield
     | Computation.Cancel_after r -> begin
-        (* We check cancelation status once and then either perform the
-           whole operation or discontinue the fiber. *)
         let fiber = Fiber.Maybe.to_fiber t.fiber in
         if Fiber.is_canceled fiber then t.discontinue
         else
@@ -186,33 +169,20 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
               end)
     | _ -> None
   and retc () =
-    Atomic.decr t.num_alive_fibers;
+    t.num_alive_fibers <- t.num_alive_fibers - 1;
     next t
   and resume trigger fiber k =
     let resume = Resume (fiber, k) in
-    if Fiber.unsuspend fiber trigger then
-      (* The fiber has not been canceled, so we queue the fiber normally. *)
-      Mpscq.push t.ready resume
-    else
-      (* The fiber has been canceled, so we give priority to it in this
-         scheduler. *)
-      Mpscq.push_head t.ready resume;
-    (* As the trigger might have been signaled from another domain or systhread
-       outside of the scheduler, we check whether the scheduler needs to be
-       woken up and take care of it if necessary. *)
+    if Fiber.unsuspend fiber trigger then Mpscq.push t.ready resume
+    else Mpscq.push_head t.ready resume;
     if
       Atomic.get t.needs_wakeup
       && Atomic.compare_and_set t.needs_wakeup true false
     then begin
       begin
         match Mutex.lock t.mutex with
         | () -> Mutex.unlock t.mutex
-        | exception Sys_error _ ->
-            (* This should mean that [resume] was called from a signal handler
-               running on the scheduler thread.  If the assumption about not
-               having poll points holds, the [Condition.broadcast] should now be
-               able to wake up the [Condition.wait] in the scheduler. *)
-            ()
+        | exception Sys_error _ -> ()
       end;
       Condition.broadcast t.condition
     end