From 436e6773158b14e42cdcd2277f4a6dd06191c73d Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes@ahrefs.com>
Date: Tue, 25 Aug 2026 08:46:44 -0400
Subject: [PATCH 1/5] add optional rate limiter to Log

https://ahrefs.slack.com/archives/C03254B74/p1787278689533289
---
 log.ml    | 85 ++++++++++++++++++++++++++++++++++++++++++++-----------
 logger.ml |  6 ++--
 2 files changed, 72 insertions(+), 19 deletions(-)

diff --git a/log.ml b/log.ml
index e8dc769..98b6cc5 100644
--- a/log.ml
+++ b/log.ml
@@ -38,6 +38,54 @@ open Printf
 open ExtLib
 open Prelude
 
+module Rate_limit : sig
+  type t
+  val none : t
+  val create : max:int -> period:Time.t -> unit -> t
+  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 = struct
+  type t =
+    | None
+    | RL of {
+      backlog: Time.t Queue.t; (** Deadlines where tokens available again *)
+      mutable count_silenced: int;
+      max: int;
+      period: Time.t;
+    }
+
+  let none = None
+  let create ~max ~period () : t =
+    if max < 1 || period < Time.msec 1 then invalid_arg "Log.Rate_limit: max>=1, period>=1ms";
+    RL { backlog=Queue.create(); count_silenced=0; max; period }
+  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 ->
+      (* inspired from ahrefskit *)
+      let now = Time.now() in
+      if Queue.length rl.backlog < rl.max then (
+        Queue.push now rl.backlog;
+        true
+      ) else match Queue.peek rl.backlog with
+        | ts when ts < now -. rl.period ->
+          ignore (Queue.pop rl.backlog : Time.t);
+          Queue.push now rl.backlog;
+          true
+        | _ ->
+          rl.count_silenced <- 1 + rl.count_silenced;
+          false
+end
+
 (** Global logger state *)
 module State = struct
   let all = Hashtbl.create 10
@@ -184,26 +232,29 @@ 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: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=Rate_limit.none) ?exn ?(lines=true) ?(backtrace=false) ?saved_backtrace ?(ts=Unix.gettimeofday()) ?(structured_pairs=[]) ?(pairs=[]) s ->
+    if logger.allowed facil level && Rate_limit.attempt rate_limit then
     let pairs = if State.is_structured_format () then List.rev_append structured_pairs pairs else pairs in
     try
+      let rate_limited = Rate_limit.take_rate_limited_count rate_limit in
+      if rate_limited > 0 then logger.put level 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 +265,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)

From 9f3e096ac48d2b5f5a77f00abc3411754037eab4 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes@ahrefs.com>
Date: Tue, 25 Aug 2026 08:46:50 -0400
Subject: [PATCH 2/5] basic test for rate limiter

---
 dune                   |  8 +++++++-
 test_log_rate_limit.ml | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 test_log_rate_limit.ml

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/test_log_rate_limit.ml b/test_log_rate_limit.ml
new file mode 100644
index 0000000..1c5b262
--- /dev/null
+++ b/test_log_rate_limit.ml
@@ -0,0 +1,37 @@
+open Devkit
+
+let fail expected actual =
+  Printf.eprintf "expected:\n%s\nactual:\n%s\n" expected actual;
+  exit 1
+
+let () =
+  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 = Log.Rate_limit.create ~max:2 ~period:1. () in
+  let emit () =
+    for i = 0 to 9_999 do
+      log#info ~rate_limit "logging %d" i
+    done
+  in
+  emit ();
+  Unix.sleep 2;
+  emit ();
+  let expected =
+    String.concat "\n" [
+      "logging 0";
+      "logging 1";
+      "(9998 messages have been rate limited)";
+      "logging 0";
+      "logging 1";
+      "";
+    ]
+  in
+  let actual = Buffer.contents output in
+  if actual <> expected then fail expected actual

From a27529af8f8737df189d8edb0d3990e6ddad3aa6 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes@ahrefs.com>
Date: Tue, 25 Aug 2026 08:54:20 -0400
Subject: [PATCH 3/5] notifications about rate limiting are at Warn level

---
 log.ml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/log.ml b/log.ml
index 98b6cc5..00e6446 100644
--- a/log.ml
+++ b/log.ml
@@ -253,8 +253,9 @@ class logger ?(logger=State.logger) facil =
     if logger.allowed facil level && Rate_limit.attempt rate_limit then
     let pairs = if State.is_structured_format () then List.rev_append structured_pairs pairs else pairs in
     try
-      let rate_limited = Rate_limit.take_rate_limited_count rate_limit in
-      if rate_limited > 0 then logger.put level facil ts [] (sprintf "(%d messages have been rate limited)" rate_limited);
+      if Logger.allowed facil `Warn then
+        let rate_limited = 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 ->

From d6991e02438647633f431e5c9ef123f81eb69d9f Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes@ahrefs.com>
Date: Tue, 25 Aug 2026 11:16:52 -0400
Subject: [PATCH 4/5] post-review refactor

---
 log.ml | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/log.ml b/log.ml
index 00e6446..5cc63cb 100644
--- a/log.ml
+++ b/log.ml
@@ -71,19 +71,19 @@ end = struct
   let attempt = function
     | None -> true
     | RL rl ->
-      (* inspired from ahrefskit *)
       let now = Time.now() in
       if Queue.length rl.backlog < rl.max then (
-        Queue.push now rl.backlog;
+        Queue.push (now +. rl.period) rl.backlog;
         true
-      ) else match Queue.peek rl.backlog with
-        | ts when ts < now -. rl.period ->
-          ignore (Queue.pop rl.backlog : Time.t);
-          Queue.push now rl.backlog;
-          true
-        | _ ->
-          rl.count_silenced <- 1 + rl.count_silenced;
-          false
+      ) else if Queue.peek rl.backlog < now then (
+        (* first deadline is past, this attempt succeeds *)
+        ignore (Queue.pop rl.backlog : Time.t);
+        Queue.push (now +. rl.period) rl.backlog;
+        true
+      ) else (
+        rl.count_silenced <- 1 + rl.count_silenced;
+        false
+      )
 end
 
 (** Global logger state *)

From 0e8f7cb3ca0e484c91211cb562ca95151080135a Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes@ahrefs.com>
Date: Tue, 25 Aug 2026 22:00:04 -0400
Subject: [PATCH 5/5] move log rate limiting to Control, implement token bucket

---
 control.ml             | 58 ++++++++++++++++++++++++++++++++++++++++++
 control.mli            | 19 ++++++++++++++
 log.ml                 | 56 +++-------------------------------------
 test_log_rate_limit.ml | 52 +++++++++++++++++++++++++++----------
 4 files changed, 120 insertions(+), 65 deletions(-)

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/log.ml b/log.ml
index 5cc63cb..d29f96d 100644
--- a/log.ml
+++ b/log.ml
@@ -38,54 +38,6 @@ open Printf
 open ExtLib
 open Prelude
 
-module Rate_limit : sig
-  type t
-  val none : t
-  val create : max:int -> period:Time.t -> unit -> t
-  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 = struct
-  type t =
-    | None
-    | RL of {
-      backlog: Time.t Queue.t; (** Deadlines where tokens available again *)
-      mutable count_silenced: int;
-      max: int;
-      period: Time.t;
-    }
-
-  let none = None
-  let create ~max ~period () : t =
-    if max < 1 || period < Time.msec 1 then invalid_arg "Log.Rate_limit: max>=1, period>=1ms";
-    RL { backlog=Queue.create(); count_silenced=0; max; period }
-  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 Queue.length rl.backlog < rl.max then (
-        Queue.push (now +. rl.period) rl.backlog;
-        true
-      ) else if Queue.peek rl.backlog < now then (
-        (* first deadline is past, this attempt succeeds *)
-        ignore (Queue.pop rl.backlog : Time.t);
-        Queue.push (now +. rl.period) rl.backlog;
-        true
-      ) else (
-        rl.count_silenced <- 1 + rl.count_silenced;
-        false
-      )
-end
-
 (** Global logger state *)
 module State = struct
   let all = Hashtbl.create 10
@@ -232,7 +184,7 @@ 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 = ?rate_limit: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
+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 ?(logger=State.logger) facil =
   let make_s (logger: Logger.t) (level:Logger.level) =
@@ -249,12 +201,12 @@ class logger ?(logger=State.logger) facil =
     output lines facil ts pairs (s ^ " : exn " ^ Exn.str exn ^ (if bt = [] then " (no backtrace)" else ""));
     List.iter (fun line -> logger.put level facil ts pairs ("    " ^ line)) bt
   in
-  fun ?(rate_limit=Rate_limit.none) ?exn ?(lines=true) ?(backtrace=false) ?saved_backtrace ?(ts=Unix.gettimeofday()) ?(structured_pairs=[]) ?(pairs=[]) s ->
-    if logger.allowed facil level && Rate_limit.attempt rate_limit then
+  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 = Rate_limit.take_rate_limited_count rate_limit in
+        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
diff --git a/test_log_rate_limit.ml b/test_log_rate_limit.ml
index 1c5b262..29f0991 100644
--- a/test_log_rate_limit.ml
+++ b/test_log_rate_limit.ml
@@ -4,7 +4,35 @@ 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);
@@ -14,24 +42,22 @@ let () =
   } in
   let logger = Logger.put_simple target in
   let log = new Log.logger ~logger (Log.facility "rate-limit-test") in
-  let rate_limit = Log.Rate_limit.create ~max:2 ~period:1. () in
-  let emit () =
-    for i = 0 to 9_999 do
+  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 ();
+  emit 10_000;
   Unix.sleep 2;
-  emit ();
+  (* 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 0";
-      "logging 1";
-      "(9998 messages have been rate limited)";
-      "logging 0";
-      "logging 1";
-      "";
-    ]
+    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