|
| 1 | +(* Copyright 2025 Luca Carlon *) |
| 2 | +(* |
| 3 | + This file is part of mldonkey. |
| 4 | +
|
| 5 | + mldonkey is free software; you can redistribute it and/or modify |
| 6 | + it under the terms of the GNU General Public License as published by |
| 7 | + the Free Software Foundation; either version 2 of the License, or |
| 8 | + (at your option) any later version. |
| 9 | +
|
| 10 | + mldonkey is distributed in the hope that it will be useful, |
| 11 | + but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 12 | + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 13 | + GNU General Public License for more details. |
| 14 | +
|
| 15 | + You should have received a copy of the GNU General Public License |
| 16 | + along with mldonkey; if not, write to the Free Software |
| 17 | + Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA |
| 18 | +*) |
| 19 | + |
| 20 | +module TaskQueue = struct |
| 21 | + type 'a t = { |
| 22 | + queue : 'a Queue.t; |
| 23 | + mutex : Mutex.t; |
| 24 | + cond : Condition.t; |
| 25 | + } |
| 26 | + |
| 27 | + (* Create a new task queue *) |
| 28 | + let create () = { |
| 29 | + queue = Queue.create (); |
| 30 | + mutex = Mutex.create (); |
| 31 | + cond = Condition.create (); |
| 32 | + } |
| 33 | + |
| 34 | + (* Add a task to the queue *) |
| 35 | + let add t task = |
| 36 | + Mutex.lock t.mutex; |
| 37 | + Queue.add task t.queue; |
| 38 | + Condition.signal t.cond; |
| 39 | + Mutex.unlock t.mutex |
| 40 | + |
| 41 | + (* Retrieve a task from the queue (blocking if empty) *) |
| 42 | + let take t = |
| 43 | + Mutex.lock t.mutex; |
| 44 | + while Queue.is_empty t.queue do |
| 45 | + Condition.wait t.cond t.mutex |
| 46 | + done; |
| 47 | + let task = Queue.pop t.queue in |
| 48 | + Mutex.unlock t.mutex; |
| 49 | + task |
| 50 | +end |
| 51 | + |
| 52 | +(* Thread pool *) |
| 53 | +type t = { |
| 54 | + threads : Thread.t list; |
| 55 | + tasks : (unit -> unit) TaskQueue.t; |
| 56 | + stop_flag : bool ref; |
| 57 | +} |
| 58 | + |
| 59 | +(* Worker thread function *) |
| 60 | +let rec worker_loop tasks stop_flag = |
| 61 | + if !stop_flag then () |
| 62 | + else |
| 63 | + let task = TaskQueue.take tasks in |
| 64 | + (try task () with _ -> ()); |
| 65 | + worker_loop tasks stop_flag |
| 66 | + |
| 67 | +(* Create a thread pool with a fixed number of threads *) |
| 68 | +let create num_threads = |
| 69 | + let tasks = TaskQueue.create () in |
| 70 | + let stop_flag = ref false in |
| 71 | + let threads = List.init num_threads (fun _ -> |
| 72 | + Thread.create (fun () -> worker_loop tasks stop_flag) () |
| 73 | + ) in |
| 74 | + { threads; tasks; stop_flag } |
| 75 | + |
| 76 | +(* Add a task to the thread pool *) |
| 77 | +let add_task pool task = |
| 78 | + TaskQueue.add pool.tasks task |
| 79 | + |
| 80 | +(* Stop the thread pool and wait for all threads to finish *) |
| 81 | +let stop pool = |
| 82 | + pool.stop_flag := true; |
| 83 | + List.iter Thread.join pool.threads |
0 commit comments