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

fix pool: on shutdown, finish reading from all queues

Browse source
This commit is contained in:
Simon Cruanes 2023-10-24 09:51:40 -04:00
parent b8a31b088f
commit faeb95b49d
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
1 changed files with 45 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

69
src/pool.ml
View file

@ -108,34 +108,55 @@ let worker_thread_ (runner : t) ~on_exn ~around_task (active : bool A.t)
let num_qs = Array.length qs in let num_qs = Array.length 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
for i = 0 to num_qs - 1 do
let q = qs.((offset + i) mod num_qs) in
match Bb_queue.try_pop ~force_lock:false q with
| Some f -> raise_notrace (Got_task f)
| None -> ()
done;
None
with Got_task f -> Some f
in
(* last resort: block on my queue *)
let[@inline] pop_blocking () =
let my_q = qs.(offset mod num_qs) in
Bb_queue.pop my_q
in
let run_task task : unit =
let _ctx = before_task runner in
(* run the task now, catching errors *)
(try task ()
with e ->
let bt = Printexc.get_raw_backtrace () in
on_exn e bt);
after_task runner _ctx
in
let run_tasks_already_present () =
(* drain the queues from existing tasks *)
let continue = ref true in
while !continue do
match get_task_without_blocking () with
| None -> continue := false
| Some task -> run_task task
done
in
let main_loop () = let main_loop () =
while A.get active do while A.get active do
(* last resort: block on my queue *) run_tasks_already_present ();
let pop_blocking () =
let my_q = qs.(offset mod num_qs) in
Bb_queue.pop my_q
in
let task = (* no task available, block until one comes *)
try let task = pop_blocking () in
for i = 0 to num_qs - 1 do run_task task
let q = qs.((offset + i) mod num_qs) in done;
match Bb_queue.try_pop ~force_lock:false q with
| Some f -> raise_notrace (Got_task f)
| None -> ()
done;
pop_blocking ()
with Got_task f -> f
in
let _ctx = before_task runner in (* cleanup *)
(* run the task now, catching errors *) run_tasks_already_present ()
(try task ()
with e ->
let bt = Printexc.get_raw_backtrace () in
on_exn e bt);
after_task runner _ctx
done
in in
try try