diff --git a/control.ml b/control.ml index 8bdbe1e..d93282c 100644 --- a/control.ml +++ b/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 diff --git a/control.mli b/control.mli index 9092697..ac89017 100644 --- a/control.mli +++ b/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 diff --git a/dune b/dune index bf95516..f21374a 100644 --- a/dune +++ b/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))) diff --git a/log.ml b/log.ml index e8dc769..d29f96d 100644 --- a/log.ml +++ b/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 diff --git a/logger.ml b/logger.ml index 7d651b6..3a27585 100644 --- a/logger.ml +++ b/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) diff --git a/test_log_rate_limit.ml b/test_log_rate_limit.ml new file mode 100644 index 0000000..29f0991 --- /dev/null +++ b/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