Source file rate_limiter.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
(** Route-based rate limiting for HTTP clients. *)

let src = Logs.Src.create "polymarket.rate_limiter" ~doc:"Rate limiter"

module Log = (val Logs.src_log src : Logs.LOG)

(* Re-export core types *)
type behavior = Types.behavior = Delay | Error

type route_pattern = Types.route_pattern = {
  host : string option;
  method_ : string option;
  path_prefix : string option;
}

type limit_config = Types.limit_config = {
  requests : int;
  window_seconds : float;
}

type route_config = Types.route_config = {
  pattern : route_pattern;
  limits : limit_config list;
  behavior : behavior;
}

type error = Types.error =
  | Rate_limited of { retry_after : float; route_key : string }

(* Rate limiter state *)
type t = {
  mutable routes : route_config list;
  routes_mutex : Eio.Mutex.t;
  state : State.t;
  sleep : float -> unit;
}

let create ~routes ~clock ?max_idle_time () =
  let state = State.create ~clock ?max_idle_time () in
  let sleep duration = Eio.Time.sleep clock duration in
  let routes_mutex = Eio.Mutex.create () in
  { routes; routes_mutex; state; sleep }

let update_routes t routes =
  Eio.Mutex.use_rw ~protect:true t.routes_mutex (fun () -> t.routes <- routes)

(* Wait for rate limit slot with Delay behavior, retrying until successful *)
let rec wait_for_slot t ~route_key ~limits =
  match State.check_limits t.state ~route_key ~limits with
  | Ok () -> ()
  | Error retry ->
      Log.debug (fun m ->
          m "Rate limited (delay): %s, waiting %.2fs" route_key retry);
      t.sleep retry;
      wait_for_slot t ~route_key ~limits

(* Check rate limits, returns None if allowed or Some error *)
let check t ~method_ ~uri =
  let matching =
    Eio.Mutex.use_ro t.routes_mutex (fun () ->
        Matcher.find_matching_routes ~method_ ~uri t.routes)
  in
  let rec check_routes routes max_error =
    match routes with
    | [] -> max_error
    | (route : route_config) :: rest ->
        let route_key = Matcher.make_route_key ~method_ ~uri route.pattern in
        let result =
          State.check_limits t.state ~route_key ~limits:route.limits
        in
        let new_max_error =
          match (result, route.behavior, max_error) with
          | Ok (), _, max_err -> max_err
          | Error retry, Error, None ->
              Log.debug (fun m ->
                  m "Rate limited (error): %s, retry after %.2fs" route_key
                    retry);
              Some (Rate_limited { retry_after = retry; route_key })
          | Error retry, Error, Some (Rate_limited prev) ->
              Some
                (Rate_limited
                   {
                     retry_after = Float.max retry prev.retry_after;
                     route_key = prev.route_key;
                   })
          | Error _, Delay, _ ->
              wait_for_slot t ~route_key ~limits:route.limits;
              max_error
        in
        check_routes rest new_max_error
  in
  check_routes matching None

exception Rate_limit_exceeded of error

let before_request t ~method_ ~uri =
  match check t ~method_ ~uri with
  | None -> ()
  | Some err -> raise (Rate_limit_exceeded err)

let before_request_result t ~method_ ~uri =
  match check t ~method_ ~uri with None -> Ok () | Some err -> Error err

(* State management *)
let cleanup t = State.cleanup t.state
let state_count t = State.state_count t.state
let reset t = State.reset t.state

(* Sub-modules *)
module Types = Types
module Builder = Builder
module Gcra = Gcra
module State = State
module Matcher = Matcher