unikernel/duniverse/lwt/src/unix/lwt_preemptive.ml

@@ -0,0 +1,263 @@
+(* This file is part of Lwt, released under the MIT license. See LICENSE.md for
+   details, or visit https://github.com/ocsigen/lwt/blob/master/LICENSE.md. *)
+
+
+
+(* [Lwt_sequence] is deprecated – we don't want users outside Lwt using it.
+   However, it is still used internally by Lwt. So, briefly disable warning 3
+   ("deprecated"), and create a local, non-deprecated alias for
+   [Lwt_sequence] that can be referred to by the rest of the code in this
+   module without triggering any more warnings. *)
+module Lwt_sequence = Lwt_sequence
+
+open Lwt.Infix
+
+(* +-----------------------------------------------------------------+
+   | Parameters                                                      |
+   +-----------------------------------------------------------------+ *)
+
+(* Minimum number of preemptive threads: *)
+let min_threads : int ref = ref 0
+
+(* Maximum number of preemptive threads: *)
+let max_threads : int ref = ref 0
+
+(* Size of the waiting queue: *)
+let max_thread_queued = ref 1000
+
+let get_max_number_of_threads_queued _ =
+  !max_thread_queued
+
+let set_max_number_of_threads_queued n =
+  if n < 0 then invalid_arg "Lwt_preemptive.set_max_number_of_threads_queued";
+  max_thread_queued := n
+
+(* The total number of preemptive threads currently running: *)
+let threads_count = ref 0
+
+(* +-----------------------------------------------------------------+
+   | Preemptive threads management                                   |
+   +-----------------------------------------------------------------+ *)
+
+module CELL :
+sig
+  type 'a t
+
+  val make : unit -> 'a t
+  val get : 'a t -> 'a
+  val set : 'a t -> 'a -> unit
+end =
+struct
+  type 'a t = {
+    m  : Mutex.t;
+    cv : Condition.t;
+    mutable cell : 'a option;
+  }
+
+  let make () = { m = Mutex.create (); cv = Condition.create (); cell = None }
+
+  let get t =
+    let rec await_value t =
+      match t.cell with
+      | None ->
+        Condition.wait t.cv t.m;
+        await_value t
+      | Some v ->
+        t.cell <- None;
+        Mutex.unlock t.m;
+        v
+    in
+    Mutex.lock t.m;
+    await_value t
+
+  let set t v =
+    Mutex.lock t.m;
+    t.cell <- Some v;
+    Mutex.unlock t.m;
+    Condition.signal t.cv
+end
+
+type thread = {
+  task_cell: (int * (unit -> unit)) CELL.t;
+  (* Channel used to communicate notification id and tasks to the
+     worker thread. *)
+
+  mutable thread : Thread.t;
+  (* The worker thread. *)
+
+  mutable reuse : bool;
+  (* Whether the thread must be re-added to the pool when the work is
+     done. *)
+}
+
+(* Pool of worker threads: *)
+let workers : thread Queue.t = Queue.create ()
+
+(* Queue of clients waiting for a worker to be available: *)
+let waiters : thread Lwt.u Lwt_sequence.t = Lwt_sequence.create ()
+
+(* Code executed by a worker: *)
+let rec worker_loop worker =
+  let id, task = CELL.get worker.task_cell in
+  task ();
+  (* If there is too much threads, exit. This can happen if the user
+     decreased the maximum: *)
+  if !threads_count > !max_threads then worker.reuse <- false;
+  (* Tell the main thread that work is done: *)
+  Lwt_unix.send_notification id;
+  if worker.reuse then worker_loop worker
+
+(* create a new worker: *)
+let make_worker () =
+  incr threads_count;
+  let worker = {
+    task_cell = CELL.make ();
+    thread = Thread.self ();
+    reuse = true;
+  } in
+  worker.thread <- Thread.create worker_loop worker;
+  worker
+
+(* Add a worker to the pool: *)
+let add_worker worker =
+  match Lwt_sequence.take_opt_l waiters with
+  | None ->
+    Queue.add worker workers
+  | Some w ->
+    Lwt.wakeup w worker
+
+(* Wait for worker to be available, then return it: *)
+let get_worker () =
+  if not (Queue.is_empty workers) then
+    Lwt.return (Queue.take workers)
+  else if !threads_count < !max_threads then
+    Lwt.return (make_worker ())
+  else
+    (Lwt.add_task_r [@ocaml.warning "-3"]) waiters
+
+(* +-----------------------------------------------------------------+
+   | Initialisation, and dynamic parameters reset                    |
+   +-----------------------------------------------------------------+ *)
+
+let get_bounds () = (!min_threads, !max_threads)
+
+let set_bounds (min, max) =
+  if min < 0 || max < min then invalid_arg "Lwt_preemptive.set_bounds";
+  let diff = min - !threads_count in
+  min_threads := min;
+  max_threads := max;
+  (* Launch new workers: *)
+  for _i = 1 to diff do
+    add_worker (make_worker ())
+  done
+
+let initialized = ref false
+
+let init min max _errlog =
+  initialized := true;
+  set_bounds (min, max)
+
+let simple_init () =
+  if not !initialized then begin
+    initialized := true;
+    set_bounds (0, 4)
+  end
+
+let nbthreads () = !threads_count
+let nbthreadsqueued () = Lwt_sequence.fold_l (fun _ x -> x + 1) waiters 0
+let nbthreadsbusy () = !threads_count - Queue.length workers
+
+(* +-----------------------------------------------------------------+
+   | Detaching                                                       |
+   +-----------------------------------------------------------------+ *)
+
+let init_result = Result.Error (Failure "Lwt_preemptive.detach")
+
+let detach f args =
+  simple_init ();
+  let result = ref init_result in
+  (* The task for the worker thread: *)
+  let task () =
+    try
+      result := Result.Ok (f args)
+    with exn when Lwt.Exception_filter.run exn ->
+      result := Result.Error exn
+  in
+  get_worker () >>= fun worker ->
+  let waiter, wakener = Lwt.wait () in
+  let id =
+    Lwt_unix.make_notification ~once:true
+      (fun () -> Lwt.wakeup_result wakener !result)
+  in
+  Lwt.finalize
+    (fun () ->
+       (* Send the id and the task to the worker: *)
+       CELL.set worker.task_cell (id, task);
+       waiter)
+    (fun () ->
+       if worker.reuse then
+         (* Put back the worker to the pool: *)
+         add_worker worker
+       else begin
+         decr threads_count;
+         (* Or wait for the thread to terminates, to free its associated
+            resources: *)
+         Thread.join worker.thread
+       end;
+       Lwt.return_unit)
+
+(* +-----------------------------------------------------------------+
+   | Running Lwt threads in the main thread                          |
+   +-----------------------------------------------------------------+ *)
+
+(* Queue of [unit -> unit Lwt.t] functions. *)
+let jobs = Queue.create ()
+
+(* Mutex to protect access to [jobs]. *)
+let jobs_mutex = Mutex.create ()
+
+let job_notification =
+  Lwt_unix.make_notification
+    (fun () ->
+       (* Take the first job. The queue is never empty at this
+          point. *)
+       Mutex.lock jobs_mutex;
+       let thunk = Queue.take jobs in
+       Mutex.unlock jobs_mutex;
+       ignore (thunk ()))
+
+let run_in_main_dont_wait f =
+  (* Add the job to the queue. *)
+  Mutex.lock jobs_mutex;
+  Queue.add f jobs;
+  Mutex.unlock jobs_mutex;
+  (* Notify the main thread. *)
+  Lwt_unix.send_notification job_notification
+
+(* There is a potential performance issue from creating a cell every time this
+   function is called. See:
+   https://github.com/ocsigen/lwt/issues/218
+   https://github.com/ocsigen/lwt/pull/219
+   https://github.com/ocaml/ocaml/issues/7158 *)
+let run_in_main f =
+  let cell = CELL.make () in
+  (* Create the job. *)
+  let job () =
+    (* Execute [f] and wait for its result. *)
+    Lwt.try_bind f
+      (fun ret -> Lwt.return (Result.Ok ret))
+      (fun exn -> Lwt.return (Result.Error exn)) >>= fun result ->
+    (* Send the result. *)
+    CELL.set cell result;
+    Lwt.return_unit
+  in
+  run_in_main_dont_wait job;
+  (* Wait for the result. *)
+  match CELL.get cell with
+  | Result.Ok ret -> ret
+  | Result.Error exn -> raise exn
+
+(* This version shadows the one above, adding an exception handler *)
+let run_in_main_dont_wait f handler =
+  let f () = Lwt.catch f (fun exc -> handler exc; Lwt.return_unit) in
+  run_in_main_dont_wait f