Source file throttled.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
open! Core
open! Import
open! Deferred_std
module Deferred = Deferred1

module Counting_semaphore : sig
  type t

  val wait_to_acquire_job_token : t -> unit Deferred.t
  val release_job_token : t -> unit
  val abort : t -> unit
  val create : max_concurrent_jobs:int -> t
end = struct
  type t =
    { mutable max_concurrent_jobs : int
    ; mutable waiter : unit Ivar.t option
    ; mutable aborted : bool
    }

  let wait_to_acquire_job_token ({ max_concurrent_jobs; waiter; aborted } as t) =
    match aborted with
    | true -> Deferred.never ()
    | false ->
      if max_concurrent_jobs > 0
      then (
        t.max_concurrent_jobs <- max_concurrent_jobs - 1;
        Deferred.return ())
      else (
        assert (Option.is_none waiter);
        let ivar = Ivar.create () in
        t.waiter <- Some ivar;
        Ivar.read ivar)
  ;;

  let release_job_token ({ max_concurrent_jobs; waiter; aborted = _ } as t) =
    match waiter with
    | Some ivar ->
      Ivar.fill ivar ();
      t.waiter <- None
    | None -> t.max_concurrent_jobs <- max_concurrent_jobs + 1
  ;;

  let abort t =
    t.aborted <- true;
    t.waiter <- None
  ;;

  let create ~max_concurrent_jobs =
    { max_concurrent_jobs; waiter = None; aborted = false }
  ;;
end

module T = struct
  type 'a t =
    { compute :
        Execution_context.t -> Counting_semaphore.t -> ('a Deferred.t -> unit) -> unit
    }
  [@@unboxed]

  let return x = { compute = (fun _ _ k -> k (return x)) }

  let map =
    `Custom
      (fun t ~f ->
         { compute =
             (fun exec_ctx semaphore k ->
                t.compute exec_ctx semaphore (fun d ->
                  k
                    (let%map result = d in
                     f result)))
         })
  ;;

  let apply t_f t =
    { compute =
        (fun exec_ctx semaphore k ->
           t_f.compute exec_ctx semaphore (fun df ->
             t.compute exec_ctx semaphore (fun dv ->
               k
                 (let%bind f = df in
                  let%map v = dv in
                  f v))))
    }
  ;;
end

include T
include Applicative.Make (T)

let enqueue' scheduler ctx f =
  let ivar = Ivar.create () in
  Scheduler.enqueue scheduler ctx (fun () -> upon (f ()) (Ivar.fill ivar)) ();
  Ivar.read ivar
;;

let job f =
  { compute =
      (fun exec_ctx semaphore k ->
         Deferred.upon (Counting_semaphore.wait_to_acquire_job_token semaphore) (fun () ->
           k
             (enqueue' (Scheduler.t ()) exec_ctx (fun () ->
                let%map a = f () in
                Counting_semaphore.release_job_token semaphore;
                a))))
  }
;;

let run t ~max_concurrent_jobs =
  let semaphore = Counting_semaphore.create ~max_concurrent_jobs in
  (* The name is set to the empty string in order to prevent [Monitor.send_exn]
     from appending information about this monitor to the exceptions we forward.
     This matters because we want simliar behavior to [Throttle] and not break
     existing tests. *)
  let monitor = Monitor.create ~name:"" () in
  let parent_monitor = Monitor.current () in
  Monitor.detach_and_iter_errors monitor ~f:(fun err ->
    Counting_semaphore.abort semaphore;
    Monitor.send_exn parent_monitor err);
  let exec_ctx =
    Execution_context.create_like
      ~monitor
      (Scheduler.current_execution_context (Scheduler.t ()))
  in
  let ivar = Ivar.create () in
  t.compute exec_ctx semaphore (fun r -> Deferred.upon r (Ivar.fill ivar));
  Ivar.read ivar
;;

let of_thunk thunk =
  { compute =
      (fun exec_ctx semaphore k ->
         let t = thunk () in
         t.compute exec_ctx semaphore k)
  }
;;

let ( *> ) t1 t2 =
  { compute =
      (fun exec_ctx semaphore k ->
         t1.compute exec_ctx semaphore (fun d1 ->
           t2.compute exec_ctx semaphore (fun d2 ->
             k
               (let%bind () = d1 in
                d2))))
  }
;;

let both_unit = ( *> )