simon/moonpool
1
0
Fork 0
mirror of https://github.com/c-cube/moonpool.git synced 2025-12-06 03:05:30 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix pool: rework scheduler to use one condition

Browse source
This commit is contained in:
Simon Cruanes 2023-10-24 13:47:53 -04:00
parent 69faea0bcb
commit d15bfb07f2
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
1 changed files with 45 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
src/pool.ml
View file

@ -84,26 +84,29 @@ let run_direct_ (self : state) (task : task) : unit =
let n_qs = Array.length self.qs in let n_qs = Array.length self.qs in
let offset = A.fetch_and_add self.cur_q 1 in let offset = A.fetch_and_add self.cur_q 1 in
(* blocking push, last resort *) (* push that forces lock acquisition, last resort *)
let[@inline] push_wait f = let[@inline] push_wait f =
let q_idx = offset mod Array.length self.qs in let q_idx = offset mod Array.length self.qs in
let q = self.qs.(q_idx) in let q = self.qs.(q_idx) in
TS_queue.push q f TS_queue.push q f
in in
let old_num_tasks = A.fetch_and_add self.num_tasks 1 in (try
(* try each queue with a round-robin initial offset *)
for _retry = 1 to 10 do
for i = 0 to n_qs - 1 do
let q_idx = (i + offset) mod Array.length self.qs in
let q = self.qs.(q_idx) in
try if TS_queue.try_push q task then raise_notrace Exit
(* try each queue with a round-robin initial offset *) done
for _retry = 1 to 10 do done;
for i = 0 to n_qs - 1 do push_wait task
let q_idx = (i + offset) mod Array.length self.qs in with Exit -> ());
let q = self.qs.(q_idx) in
if TS_queue.try_push q task then raise_notrace Exit (* successfully pushed, now see if we need to wakeup workers *)
done let old_num_tasks = A.fetch_and_add self.num_tasks 1 in
done; if old_num_tasks < size_ self then awake_workers_ self
push_wait task
with Exit -> if old_num_tasks < size_ self then awake_workers_ self
let rec run_async_ (self : state) (task : task) : unit = let rec run_async_ (self : state) (task : task) : unit =
let task' () = let task' () =
@ -157,13 +160,18 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task
let num_qs = Array.length self.qs in let num_qs = Array.length self.qs in
let (AT_pair (before_task, after_task)) = around_task in let (AT_pair (before_task, after_task)) = around_task in
let get_task_without_blocking () : _ option = (* try to get a task that is already in one of the queues.
@param force_lock if true, we force acquisition of the queue's mutex,
which is slower but always succeeds to get a task if there's one. *)
let get_task_already_in_queues ~force_lock () : _ option =
try try
for i = 0 to num_qs - 1 do for _retry = 1 to 3 do
let q = self.qs.((offset + i) mod num_qs) in for i = 0 to num_qs - 1 do
match TS_queue.try_pop ~force_lock:false q with let q = self.qs.((offset + i) mod num_qs) in
| Some f -> raise_notrace (Got_task f) match TS_queue.try_pop ~force_lock q with
| None -> () | Some f -> raise_notrace (Got_task f)
| None -> ()
done
done; done;
None None
with Got_task f -> with Got_task f ->
@ -171,17 +179,21 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task
Some f Some f
in in
(* last resort: block on condition or raise Closed *) (* slow path: force locking when trying to get tasks,
and wait on [self.cond] if no task is currently available. *)
let pop_blocking () : task = let pop_blocking () : task =
Mutex.lock self.mutex;
try try
while A.get self.active do while A.get self.active do
match get_task_without_blocking () with match get_task_already_in_queues ~force_lock:true () with
| Some t -> | Some t -> raise_notrace (Got_task t)
Mutex.unlock self.mutex; | None ->
raise_notrace (Got_task t) Mutex.lock self.mutex;
| None -> Condition.wait self.cond self.mutex (* NOTE: be careful about race conditions: we must only
block if the [shutdown] that sets [active] to [false]
has not broadcast over this condition first. Otherwise
we might miss the signal and wait here forever. *)
if A.get self.active then Condition.wait self.cond self.mutex;
Mutex.unlock self.mutex
done; done;
raise Closed raise Closed
with Got_task t -> t with Got_task t -> t
@ -197,11 +209,12 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task
after_task runner _ctx after_task runner _ctx
in in
let run_tasks_already_present () = (* drain the queues from existing tasks. If [force_lock=false]
(* drain the queues from existing tasks *) then it is best effort. *)
let run_tasks_already_present ~force_lock () =
let continue = ref true in let continue = ref true in
while !continue do while !continue do
match get_task_without_blocking () with match get_task_already_in_queues ~force_lock () with
| None -> continue := false | None -> continue := false
| Some task -> run_task task | Some task -> run_task task
done done
@ -209,7 +222,7 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task
let main_loop () = let main_loop () =
while A.get self.active do while A.get self.active do
run_tasks_already_present (); run_tasks_already_present ~force_lock:false ();
(* no task available, block until one comes *) (* no task available, block until one comes *)
match pop_blocking () with match pop_blocking () with
@ -218,7 +231,7 @@ let worker_thread_ (self : state) (runner : t) ~on_exn ~around_task
done; done;
(* cleanup *) (* cleanup *)
run_tasks_already_present () run_tasks_already_present ~force_lock:true ()
in in
try try