rate limiter in logger, opt-in

control.ml

@@ -25,3 +25,61 @@ let with_output_bin name k = with_open_out_bin name (fun ch -> bracket (IO.outpu
 let with_output_txt name k = with_open_out_txt name (fun ch -> bracket (IO.output_channel ch) IO.flush k)
 
 let with_opendir dir = bracket (Unix.opendir dir) Unix.closedir
+
+(* token bucket
+   https://en.wikipedia.org/wiki/Token_bucket *)
+module Rate_limit = struct
+  type t =
+    | None
+    | RL of {
+      mutable tokens: float;
+      mutable count_silenced: int;
+      mutable last_update: float;
+      capacity: float;
+      rate: float; (** new tokens/sec *)
+    }
+
+  let none = None
+
+  let create ?burst_capacity ~allowed_per_sec () : t =
+    if classify_float allowed_per_sec <> FP_normal || allowed_per_sec <= 0. then
+      invalid_arg "Rate_limit.create: allowed_per_sec must be finite and positive";
+    let capacity = match burst_capacity with
+      | Some n ->
+        if n < 1 then invalid_arg "Rate_limit.create: burst capacity must be >= 1";
+        float n
+      | None ->
+        (* default: burst of 5sec worth of tokens *)
+        max 1. @@ min max_float @@ allowed_per_sec *. 5.
+    in
+    RL {
+      tokens=capacity; last_update=Time.now(); count_silenced=0; capacity;
+      rate=allowed_per_sec;
+    }
+
+  let take_rate_limited_count = function
+    | None -> 0
+    | RL rl ->
+        let n = rl.count_silenced in
+        rl.count_silenced <- 0;
+        n
+
+  let attempt = function
+    | None -> true
+    | RL rl ->
+      let now = Time.now() in
+
+      if now > rl.last_update then (
+        rl.tokens <- min rl.capacity
+          (rl.tokens +. rl.rate *. (now -. rl.last_update));
+        rl.last_update <- now;
+      );
+
+      if rl.tokens >= 1. then (
+        rl.tokens <- rl.tokens -. 1.;
+        true
+      ) else (
+        rl.count_silenced <- 1 + rl.count_silenced;
+        false
+      )
+end

control.mli

@@ -43,3 +43,22 @@ val with_output_txt : string -> (unit IO.output -> 'a) -> 'a
 (** Misc. *)
 
 val with_opendir : string -> (Unix.dir_handle -> 'b) -> 'b
+
+
+module Rate_limit : sig
+  type t
+  val none : t
+  val create : ?burst_capacity:int -> allowed_per_sec:float -> unit -> t
+  (** Create a token-bucket limiter with the given sustained rate and capacity
+      for ten seconds of traffic (at least one token). The bucket starts full.
+      @param burst_capacity limits the size of a burst when token bucket is full
+      @param allowed_per_sec number of tokens refilled per second, ie asymptotic
+        max throughtput
+      @raise Invalid_argument if [allowed_per_sec] is not finite and positive. *)
+
+  val take_rate_limited_count: t -> int
+  (** How many attempts have been rate limited since last time this was called? *)
+
+  val attempt : t -> bool
+  (** Attempt to perform one action. Return [true] if allowed by rate limiter. *)
+end

dune

@@ -36,7 +36,8 @@
            memory_jemalloc
            test
            test_gzip
-           test_httpev)
+           test_httpev
+           test_log_rate_limit)
   (preprocess
     (per_module
       ((pps lwt_ppx)
@@ -76,6 +77,11 @@
   (libraries devkit extlib)
   (modules test_gzip))
 
+(test
+  (name test_log_rate_limit)
+  (libraries devkit unix)
+  (modules test_log_rate_limit))
+
 (rule
   (alias runtest)
   (action (run ./test.exe)))

log.ml

@@ -184,26 +184,30 @@ let read_env_config = State.read_env_config
 
   param [structured_pairs] key/value pairs to use for structured log formats only. Plain logging will discard.
 *)
-type 'a pr = ?exn:exn -> ?lines:bool -> ?backtrace:bool -> ?saved_backtrace:string list -> ?ts:Time.t -> ?structured_pairs:Logger.Pairs.t -> ?pairs:Logger.Pairs.t -> ('a, unit, string, unit) format4 -> 'a
+type 'a pr = ?rate_limit:Control.Rate_limit.t -> ?exn:exn -> ?lines:bool -> ?backtrace:bool -> ?saved_backtrace:string list -> ?ts:Time.t -> ?structured_pairs:Logger.Pairs.t -> ?pairs:Logger.Pairs.t -> ('a, unit, string, unit) format4 -> 'a
 
-class logger facil =
-  let make_s (output_line:Logger.facil -> Time.t -> Logger.Pairs.t -> string -> unit) =
+class logger ?(logger=State.logger) facil =
+  let make_s (logger: Logger.t) (level:Logger.level) =
   let output = function
   | true ->
       fun facil ts pairs s ->
         if String.contains s '\n' then
-          List.iter (output_line facil ts pairs) @@ String.nsplit s "\n"
+          List.iter (logger.put level facil ts pairs) @@ String.nsplit s "\n"
         else
-          output_line facil ts pairs s
-  | false -> output_line
+          logger.put level facil ts pairs s
+  | false -> logger.put level
   in
   let print_bt lines exn bt ts pairs s =
     output lines facil ts pairs (s ^ " : exn " ^ Exn.str exn ^ (if bt = [] then " (no backtrace)" else ""));
-    List.iter (fun line -> output_line facil ts pairs ("    " ^ line)) bt
+    List.iter (fun line -> logger.put level facil ts pairs ("    " ^ line)) bt
   in
-  fun ?exn ?(lines=true) ?(backtrace=false) ?saved_backtrace ?(ts=Unix.gettimeofday()) ?(structured_pairs=[]) ?(pairs=[]) s ->
+  fun ?(rate_limit=Control.Rate_limit.none) ?exn ?(lines=true) ?(backtrace=false) ?saved_backtrace ?(ts=Unix.gettimeofday()) ?(structured_pairs=[]) ?(pairs=[]) s ->
+    if logger.allowed facil level && Control.Rate_limit.attempt rate_limit then
     let pairs = if State.is_structured_format () then List.rev_append structured_pairs pairs else pairs in
     try
+      if Logger.allowed facil `Warn then
+        let rate_limited = Control.Rate_limit.take_rate_limited_count rate_limit in
+        if rate_limited > 0 then logger.put `Warn facil ts [] (sprintf "(%d messages have been rate limited)" rate_limited);
       match exn with
       | None -> output lines facil ts pairs s
       | Some exn ->
@@ -214,17 +218,17 @@ class logger facil =
       | true -> print_bt lines exn (Exn.get_backtrace ()) ts pairs s
       | false -> output lines facil ts pairs (s ^ " : exn " ^ Exn.str exn)
     with exn ->
-      output_line facil ts pairs (sprintf "LOG FAILED : %S with message %S" (Exn.str exn) s)
+      logger.put level facil ts pairs (sprintf "LOG FAILED : %S with message %S" (Exn.str exn) s)
 in
-let make : _ -> _ pr = fun output ?exn ?lines ?backtrace ?saved_backtrace ?ts ?structured_pairs ?pairs fmt ->
-    ksprintf (fun s -> output ?exn ?lines ?backtrace ?saved_backtrace ?ts ?structured_pairs ?pairs s) fmt
+let make : _ -> _ pr = fun output ?rate_limit ?exn ?lines ?backtrace ?saved_backtrace ?ts ?structured_pairs ?pairs fmt ->
+    ksprintf (fun s -> output ?rate_limit ?exn ?lines ?backtrace ?saved_backtrace ?ts ?structured_pairs ?pairs s) fmt
 in
-let debug_s = make_s (State.logger.put `Debug) in
-let warn_s = make_s (State.logger.put `Warn) in
-let info_s = make_s (State.logger.put `Info) in
-let error_s = make_s (State.logger.put `Error) in
-let critical_s = make_s (State.logger.put `Critical) in
-let put_s level = make_s (State.logger.put level) in
+let debug_s = make_s logger `Debug in
+let warn_s = make_s logger `Warn in
+let info_s = make_s logger `Info in
+let error_s = make_s logger `Error in
+let critical_s = make_s logger `Critical in
+let put_s level = make_s logger level in
 object
 method debug_s = debug_s
 method warn_s = warn_s

logger.ml

@@ -47,10 +47,12 @@ type target = {
 
 (** A logger *)
 type t = {
-  put : level -> facil -> Time.t -> Pairs.t -> string -> unit
-} [@@unboxed]
+  put : level -> facil -> Time.t -> Pairs.t -> string -> unit;
+  allowed : facil -> level -> bool;
+}
 
 let put_simple (t:target) : t = {
+  allowed;
   put = fun level facil ts pairs str ->
     if allowed facil level then
       t.output level facil (t.format level facil ts pairs str)

test_log_rate_limit.ml

@@ -0,0 +1,63 @@
+open Devkit
+
+let fail expected actual =
+  Printf.eprintf "expected:\n%s\nactual:\n%s\n" expected actual;
+  exit 1
+
+let expect_invalid_rate rate =
+  match Control.Rate_limit.create ~allowed_per_sec:rate () with
+  | exception Invalid_argument _ -> ()
+  | _ -> fail "Invalid_argument" "rate limiter created"
+
+let expect_invalid_capacity burst_capacity =
+  match Control.Rate_limit.create ~burst_capacity ~allowed_per_sec:1. () with
+  | exception Invalid_argument _ -> ()
+  | _ -> fail "Invalid_argument" "rate limiter created"
+
+let logging_lines count =
+  let rec loop i acc =
+    if i < 0 then acc else loop (i - 1) (Printf.sprintf "logging %d" i :: acc)
+  in
+  loop (count - 1) []
+
+let () =
+  List.iter expect_invalid_rate [0.; -1.; infinity; nan];
+  List.iter expect_invalid_capacity [0; -1];
+
+  (* A very low rate must still have capacity for its initial token. *)
+  let slow = Control.Rate_limit.create ~allowed_per_sec:0.01 () in
+  if not (Control.Rate_limit.attempt slow) then fail "allowed" "rate limited";
+  if Control.Rate_limit.attempt slow then fail "rate limited" "allowed";
+  if Control.Rate_limit.take_rate_limited_count slow <> 1 then
+    fail "one rate-limited attempt" "unexpected count";
+  if Control.Rate_limit.take_rate_limited_count slow <> 0 then
+    fail "reset rate-limited count" "non-zero count";
+
+  let output = Buffer.create 256 in
+  let target = { Logger.
+    format = (fun _level _facility _timestamp _pairs message -> message);
+    output = (fun _level _facility message ->
+      Buffer.add_string output message;
+      Buffer.add_char output '\n');
+  } in
+  let logger = Logger.put_simple target in
+  let log = new Log.logger ~logger (Log.facility "rate-limit-test") in
+  let rate_limit = Control.Rate_limit.create ~burst_capacity:7 ~allowed_per_sec:2. () in
+  let emit count =
+    for i = 0 to count - 1 do
+      log#info ~rate_limit "logging %d" i
+    done
+  in
+  emit 10_000;
+  Unix.sleep 2;
+  (* Emit only the number guaranteed to have been refilled. This keeps a
+     delayed test process from changing the expected output. *)
+  emit 4;
+  let expected =
+    String.concat "\n"
+      (logging_lines 7 @
+       ["(9993 messages have been rate limited)"] @
+       logging_lines 4 @ [""])
+  in
+  let actual = Buffer.contents output in
+  if actual <> expected then fail expected actual