package async_kernel

  1. Overview
  2. Docs
package async_kernel
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file limiter_async.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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
open Core_kernel
open Async_kernel
open Limiter.Infinite_or_finite

module Outcome = struct
  type 'a t =
    | Ok of 'a
    | Aborted
    | Raised of exn
  [@@deriving sexp_of]
end

module Job = struct
  type t =
    | Immediate : Monitor.t * ('a -> unit) * 'a -> t
    | Deferred  : ('a -> 'b Deferred.t) * 'a * 'b Outcome.t Ivar.t -> t
end

module Expert = struct
  type t =
    { continue_on_error : bool
    (* [is_dead] is true if [t] was killed due to a job raising an exception or [kill t]
       being called. *)
    ; mutable is_dead       : bool
    (* Ivar that is filled the next time return_to_hopper is called. *)
    ; mutable hopper_filled : unit Ivar.t option
    ; limiter               : Limiter.t
    ; throttle_queue        : ((int * Job.t) Queue.t [@sexp.opaque]) }
  [@@deriving sexp_of]

  let to_jane_limiter t = t.limiter

  let cycle_start () = Async_kernel_scheduler.cycle_start_ns ()

  let create_exn
        ~hopper_to_bucket_rate_per_sec
        ~bucket_limit
        ~in_flight_limit
        ~initial_bucket_level
        ~initial_hopper_level
        ~continue_on_error
    =
    let limiter =
      Limiter.Expert.create_exn
        ~now:(cycle_start ())
        ~hopper_to_bucket_rate_per_sec
        ~bucket_limit
        ~in_flight_limit
        ~initial_bucket_level
        ~initial_hopper_level
    in
    let throttle_queue = Queue.create () in
    { continue_on_error
    ; is_dead       = false
    ; hopper_filled = None
    ; limiter
    ; throttle_queue
    }
  ;;

  let is_dead t = t.is_dead

  let kill_job = function
    | Job.Deferred  (_, _, i)       -> Ivar.fill_if_empty i Aborted
    | Job.Immediate (monitor, _, _) ->
      Monitor.send_exn monitor ~backtrace:`Get (Failure "Limiter killed")
  ;;

  let kill t =
    if not t.is_dead then begin
      t.is_dead <- true;
      Queue.iter t.throttle_queue ~f:(fun (_, job) -> kill_job job)
    end;
  ;;

  let saw_error t = if not t.continue_on_error then (kill t)

  let wait_for_hopper_fill t =
    match t.hopper_filled with
    | Some i -> Ivar.read i
    | None   ->
      let i = Ivar.create () in
      t.hopper_filled <- Some i;
      Ivar.read i
  ;;

  let return_to_hopper t ~now amount =
    begin match t.hopper_filled with
    | None   -> ()
    | Some i ->
      Ivar.fill i ();
      t.hopper_filled <- None;
    end;
    Limiter.Expert.return_to_hopper t.limiter ~now amount
  ;;

  let run_job_now t job ~return_after : unit =
    if t.is_dead
    then (kill_job job)
    else begin
      match job with
      | Job.Immediate (monitor, f, v)    ->
        begin try
          f v
        with
        | e -> Monitor.send_exn monitor ~backtrace:`Get e
        end;
        return_to_hopper t ~now:(cycle_start ()) return_after
      | Job.Deferred  (f, v, i) ->
        Monitor.try_with (fun () ->
          f v)
        >>> fun res ->
        return_to_hopper t ~now:(cycle_start ()) return_after;
        match res with
        | Error e ->
          Ivar.fill_if_empty i (Raised e);
          saw_error t
        | Ok v    ->
          Ivar.fill_if_empty i (Ok v)
    end
  ;;

  (* given a job, immediately creates and runs a job that fails with the given (as a
     format string) message *)
  let fail_job t job k =
    ksprintf (fun s ->
      let f () = failwith s in
      let job =
        match job with
        | Job.Immediate (monitor, _, _) -> Job.Immediate (monitor, f, ())
        | Job.Deferred (_, _, i)        -> Job.Deferred (f, (), i)
      in
      run_job_now t job ~return_after:0)
      k
  ;;

  let rec run_throttled_jobs_until_empty t =
    if Queue.length t.throttle_queue = 0
    then ()
    else begin
      let amount, job = Queue.peek_exn t.throttle_queue in
      let now         = cycle_start () in
      match Limiter.Expert.try_take t.limiter ~now amount with
      | Asked_for_more_than_bucket_limit ->
        fail_job t job !"job asked for more tokens (%i) than possible (%i)"
          amount (Limiter.bucket_limit t.limiter);
        run_throttled_jobs_until_empty t
      | Taken ->
        (* Safe, because we checked the length above.  And, we're guaranteed that
           dequeue_exn gets out the same job that peek_exn does.  *)
        ignore (Queue.dequeue_exn t.throttle_queue : (int * Job.t));
        run_job_now t job ~return_after:amount;
        run_throttled_jobs_until_empty t
      | Unable ->
        begin match Limiter.Expert.tokens_may_be_available_when t.limiter ~now amount with
        | Never_because_greater_than_bucket_limit ->
          fail_job t job !"job asked for more tokens (%i) than possible (%i)"
            amount (Limiter.bucket_limit t.limiter);
          run_throttled_jobs_until_empty t
        | When_return_to_hopper_is_called   ->
          wait_for_hopper_fill t
          >>> fun () ->
          run_throttled_jobs_until_empty t
        | At expected_fill_time ->
          let min_fill_time =
            Time_ns.add (cycle_start ()) (Async_kernel_scheduler.event_precision_ns ())
          in
          Clock_ns.at (Time_ns.max expected_fill_time min_fill_time)
          >>> fun () ->
          run_throttled_jobs_until_empty t
        end
    end
  ;;

  let enqueue_job_and_maybe_start_queue_runner t amount job ~allow_immediate_run =
    let bucket_limit = Limiter.bucket_limit t.limiter in
    if bucket_limit < amount
    then (fail_job t job !"requested job size (%i) exceeds the possible size (%i)"
            amount bucket_limit);
    if t.is_dead
    then (kill_job job)
    else begin
      if Queue.length t.throttle_queue > 0
      then (Queue.enqueue t.throttle_queue (amount, job))
      else begin
        let now = cycle_start () in
        match Limiter.Expert.try_take t.limiter ~now amount with
        | Asked_for_more_than_bucket_limit ->
          fail_job t job !"requested job size (%i) exceeds the possible size (%i)"
            amount bucket_limit;
        | Taken  ->
          (* These semantics are copied from the current Throttle, and it was
             important enough there to add a specific unit test.  If you have

             do_f ();
             enqueue thing_to_do_later;
             do_g ();

             it is surprising if any portion of the closure thing_to_do_later happens, so
             we always schedule the work for later on the Async queue.

             This isn't as efficient as it could be for immediate jobs and can be avoided
             with [run_or_enqueue].
          *)
          if allow_immediate_run
          then (run_job_now t job ~return_after:amount)
          else (Async_kernel_scheduler.enqueue_job
                  Execution_context.main
                  (fun t ->
                     run_job_now t job ~return_after:amount) t)
        | Unable ->
          Queue.enqueue t.throttle_queue (amount, job);
          run_throttled_jobs_until_empty t
      end
    end
  ;;

  let enqueue_exn t ?(allow_immediate_run=false) amount f v =
    enqueue_job_and_maybe_start_queue_runner t amount ~allow_immediate_run
      (Immediate (Monitor.current (), f, v))
  ;;

  let enqueue' t amount f v =
    Deferred.create (fun i ->
      try
        enqueue_job_and_maybe_start_queue_runner t amount (Deferred (f, v, i))
          ~allow_immediate_run:false
      with e -> Ivar.fill i (Raised e))
  ;;

  let cost_of_jobs_waiting_to_start t =
    Queue.fold t.throttle_queue ~init:0
      ~f:(fun sum (cost, _) -> cost + sum)
  ;;

end

open Expert
type t = Expert.t [@@deriving sexp_of]
type limiter = t [@@deriving sexp_of]

module Common = struct
  let to_limiter (t:t) = t
  let kill = Expert.kill
  let is_dead = Expert.is_dead
end

module type Common = sig
  type _ t

  (** kills [t], which aborts all enqueued jobs that haven't started and all jobs enqueued
      in the future.  If [t] has already been killed, then calling [kill t] has no effect.
      Note that kill does not effect currently running jobs in any way. *)
  val kill : _ t -> unit

  (** [is_dead t] returns [true] if [t] was killed, either by [kill] or by an unhandled
      exception in a job. *)
  val is_dead : _ t -> bool

  val to_limiter : _ t -> limiter
end

module Token_bucket = struct
  type t = limiter [@@deriving sexp_of]
  type _ u = t

  let create_exn
        ~burst_size:bucket_limit
        ~sustained_rate_per_sec:fill_rate
        ~continue_on_error
        ?in_flight_limit
        ?(initial_burst_size = 0)
        ()
    =
    let in_flight_limit =
      match in_flight_limit with
      | None       -> Infinite
      | Some limit -> Finite limit
    in
    Expert.create_exn
      ~bucket_limit
      ~in_flight_limit
      ~hopper_to_bucket_rate_per_sec:(Finite fill_rate)
      ~initial_bucket_level:initial_burst_size
      ~initial_hopper_level:Infinite
      ~continue_on_error
  ;;

  let enqueue_exn        = Expert.enqueue_exn
  let enqueue'           = Expert.enqueue'

  include Common
end

module Throttle = struct
  type t = limiter
  [@@deriving sexp_of]
  type _ u = t

  let create_exn
        ~concurrent_jobs_target
        ~continue_on_error
        ?burst_size
        ?sustained_rate_per_sec
        ()
    =
    if concurrent_jobs_target < 1
    then (failwithf !"concurrent_jobs_target < 1 (%i) doesn't make sense"
            concurrent_jobs_target ());
    let concurrent_jobs_target = concurrent_jobs_target in
    let hopper_to_bucket_rate_per_sec =
      match sustained_rate_per_sec with
      | None      -> Infinite
      | Some rate -> Finite rate
    in
    let bucket_limit =
      match burst_size with
      | None            -> concurrent_jobs_target
      | Some burst_size -> burst_size
    in
    let initial_bucket_level = bucket_limit in
    Expert.create_exn
      ~bucket_limit
      ~in_flight_limit:(Finite concurrent_jobs_target)
      ~hopper_to_bucket_rate_per_sec
      ~initial_bucket_level
      ~initial_hopper_level:(Finite 0)
      ~continue_on_error
  ;;

  let enqueue_exn t ?allow_immediate_run f v =
    Expert.enqueue_exn t ?allow_immediate_run 1 f v
  ;;

  let enqueue' t f v =
    Expert.enqueue' t 1 f v

  let jlimiter = Expert.to_jane_limiter

  let concurrent_jobs_target t =
    jlimiter t
    |> Limiter.bucket_limit
  ;;

  let num_jobs_waiting_to_start t = Queue.length t.throttle_queue

  let num_jobs_running t =
    Limiter.in_flight (jlimiter t) ~now:(Async_kernel_scheduler.cycle_start_ns ())
  ;;

  include Common
end

module Sequencer = struct
  include Throttle

  let create ?(continue_on_error=false) ?burst_size ?sustained_rate_per_sec () =
    create_exn ~concurrent_jobs_target:1
      ~continue_on_error ?burst_size ?sustained_rate_per_sec ()
  ;;

  include Common
end

module Resource_throttle = struct
  type 'a t =
    { throttle  : Throttle.t
    ; resources : 'a Queue.t }
  [@@deriving sexp_of]

  let create_exn ~resources ~continue_on_error ?burst_size ?sustained_rate_per_sec () =
    let resources = Queue.of_list resources in
    let max_concurrent_jobs = Queue.length resources in
    let throttle =
      Throttle.create_exn
        ~concurrent_jobs_target:max_concurrent_jobs
        ~continue_on_error
        ?burst_size
        ?sustained_rate_per_sec
        ()
    in
    { throttle; resources }
  ;;

  let enqueue_gen t ?allow_immediate_run f enqueue =
    let f () =
      let v = Queue.dequeue_exn t.resources in
      protect ~f:(fun () -> f v)
        ~finally:(fun () -> Queue.enqueue t.resources v)
    in
    enqueue t.throttle ?allow_immediate_run f ()
  ;;

  let enqueue_exn t ?allow_immediate_run f =
    enqueue_gen t ?allow_immediate_run f Throttle.enqueue_exn
  ;;

  let enqueue' t f =
    let f () =
      let v = Queue.dequeue_exn t.resources in
      Monitor.protect (fun () -> f v)
        ~finally:(fun () -> Queue.enqueue t.resources v; Deferred.unit)
    in
    Throttle.enqueue' t.throttle f ()
  ;;

  let max_concurrent_jobs t = Throttle.concurrent_jobs_target t.throttle

  let to_limiter t = t.throttle
  let kill       t = kill    t.throttle
  let is_dead    t = is_dead t.throttle
end
OCaml

Innovation. Community. Security.

GitHub Discord Twitter Peertube RSS
About
Changelog Releases Industrial Users Academic Users Why OCaml
Ecosystem
Packages Community Events OCaml Planet Jobs
Resources
Install OCaml Get Started Platform Tools Language Manual Standard Library API Books Exercises Papers OCaml Playground Logo
Policies
Carbon Footprint Governance Privacy Code of Conduct