Rate limiting: token buckets with zero or negative fill rate fail

Zero or negative rate limits can cause issues in the behaviour of rate
limiting. In particular, zero fill rate leads to a division by zero in time
calculations. Rather than account for this, we forbid the creation of token
buckets with a bad fill rate by returning None.

Signed-off-by: Christian Pardillo Laursen <christian.pardillolaursen@citrix.com>

Author: cplaursen

ocaml/libs/rate-limit/bucket_table.ml

@@ -17,8 +17,13 @@ type t = (string, Token_bucket.t) Hashtbl.t
 let create () = Hashtbl.create 16
 
 let add_bucket table ~user_agent ~burst_size ~fill_rate =
-  let bucket = Token_bucket.create ~burst_size ~fill_rate in
-  Hashtbl.add table user_agent bucket
+  let bucket_option = Token_bucket.create ~burst_size ~fill_rate in
+  match bucket_option with
+  | Some bucket ->
+      Hashtbl.add table user_agent bucket ;
+      true
+  | None ->
+      false
 
 let delete_bucket table ~user_agent = Hashtbl.remove table user_agent
 

ocaml/libs/rate-limit/bucket_table.mli

@@ -12,17 +12,29 @@
  * GNU Lesser General Public License for more details.
  *)
 
+(** Hash table mapping client identifiers to their token buckets for rate limiting. *)
 type t = (string, Token_bucket.t) Hashtbl.t
 
 val create : unit -> t
+(** [create ()] creates a new empty bucket table. *)
 
 val add_bucket :
-  t -> user_agent:string -> burst_size:float -> fill_rate:float -> unit
+  t -> user_agent:string -> burst_size:float -> fill_rate:float -> bool
+(** [add_bucket table ~user_agent ~burst_size ~fill_rate] adds a token bucket
+    for the given user agent. Returns [false] if a bucket already exists, or if
+    the bucket configuration is invalid, e.g. negative fill rate. *)
 
 val peek : t -> user_agent:string -> float option
+(** [peek table ~user_agent] returns the current token count for the user agent,
+    or [None] if no bucket exists. *)
 
 val delete_bucket : t -> user_agent:string -> unit
+(** [delete_bucket table ~user_agent] removes the bucket for the user agent. *)
 
 val try_consume : t -> user_agent:string -> float -> bool
+(** [try_consume table ~user_agent amount] attempts to consume tokens.
+    Returns [true] on success, [false] if insufficient tokens. *)
 
 val consume_and_block : t -> user_agent:string -> float -> unit
+(** [consume_and_block table ~user_agent amount] consumes tokens, blocking
+    until sufficient tokens are available. *)

ocaml/libs/rate-limit/test/test_token_bucket.ml

@@ -2,11 +2,22 @@ open Thread
 open Rate_limit
 open QCheck
 
+let test_bad_fill_rate () =
+  let tb_zero = Token_bucket.create ~burst_size:1.0 ~fill_rate:0.0 in
+  Alcotest.(check bool)
+    "Creating a token bucket with 0 fill rate should fail" true (tb_zero = None) ;
+  let tb_negative = Token_bucket.create ~burst_size:1.0 ~fill_rate:~-.1.0 in
+  Alcotest.(check bool)
+    "Creating a token bucket with negative fill rate should fail" true
+    (tb_negative = None)
+
 let test_consume_removes_correct_amount () =
   let initial_time = Mtime.Span.of_uint64_ns 0L in
   let tb =
-    Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
-      ~fill_rate:2.0
+    Option.get
+      (Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
+         ~fill_rate:2.0
+      )
   in
 
   Alcotest.(check (float 0.0))
@@ -25,8 +36,10 @@ let test_consume_removes_correct_amount () =
 let test_consume_more_than_available () =
   let initial_time = Mtime.Span.of_uint64_ns 0L in
   let tb =
-    Token_bucket.create_with_timestamp initial_time ~burst_size:5.0
-      ~fill_rate:1.0
+    Option.get
+      (Token_bucket.create_with_timestamp initial_time ~burst_size:5.0
+         ~fill_rate:1.0
+      )
   in
 
   let _ = Token_bucket.consume_with_timestamp (fun () -> initial_time) tb 4.0 in
@@ -43,8 +56,10 @@ let test_consume_more_than_available () =
 let test_consume_refills_before_removing () =
   let initial_time = Mtime.Span.of_uint64_ns 0L in
   let tb =
-    Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
-      ~fill_rate:2.0
+    Option.get
+      (Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
+         ~fill_rate:2.0
+      )
   in
 
   let first_consume =
@@ -67,8 +82,10 @@ let test_consume_refills_before_removing () =
 let test_peek_respects_burst_size () =
   let initial_time = Mtime.Span.of_uint64_ns 0L in
   let tb =
-    Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
-      ~fill_rate:5.0
+    Option.get
+      (Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
+         ~fill_rate:5.0
+      )
   in
 
   let _ = Token_bucket.consume_with_timestamp (fun () -> initial_time) tb 8.0 in
@@ -80,8 +97,10 @@ let test_peek_respects_burst_size () =
 
 let test_concurrent_access () =
   let tb =
-    Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:15.0
-      ~fill_rate:0.0
+    Option.get
+      (Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:15.0
+         ~fill_rate:0.01
+      )
   in
   let threads =
     Array.init 10 (fun _ ->
@@ -101,14 +120,14 @@ let test_concurrent_access () =
     5.0
 
 let test_sleep () =
-  let tb = Token_bucket.create ~burst_size:20.0 ~fill_rate:5.0 in
+  let tb = Option.get (Token_bucket.create ~burst_size:20.0 ~fill_rate:5.0) in
   let _ = Token_bucket.consume tb 10.0 in
   Thread.delay 1.0 ;
   Alcotest.(check (float 0.2))
     "Sleep 1 should refill token bucket by fill_rate" 15.0 (Token_bucket.peek tb)
 
 let test_system_time_versions () =
-  let tb = Token_bucket.create ~burst_size:10.0 ~fill_rate:2.0 in
+  let tb = Option.get (Token_bucket.create ~burst_size:10.0 ~fill_rate:2.0) in
 
   let initial_peek = Token_bucket.peek tb in
   Alcotest.(check (float 0.01))
@@ -122,7 +141,7 @@ let test_system_time_versions () =
     "After consume, should have 7 tokens" 7.0 after_consume_peek
 
 let test_concurrent_system_time () =
-  let tb = Token_bucket.create ~burst_size:100.0 ~fill_rate:10.0 in
+  let tb = Option.get (Token_bucket.create ~burst_size:100.0 ~fill_rate:10.0) in
   let num_threads = 20 in
   let consume_per_thread = 3 in
 
@@ -149,8 +168,10 @@ let test_concurrent_system_time () =
 
 let test_consume_more_than_available_concurrent () =
   let tb =
-    Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:5.0
-      ~fill_rate:0.0
+    Option.get
+      (Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:5.0
+         ~fill_rate:0.1
+      )
   in
   let num_threads = 10 in
   let consume_per_thread = 1 in
@@ -180,15 +201,17 @@ let test_consume_more_than_available_concurrent () =
 
   Alcotest.(check int)
     "Only 5 consumptions should succeed" 5 !successful_consumes ;
-  Alcotest.(check (float 0.0))
+  Alcotest.(check (float 0.1))
     "Bucket should be empty after consumptions" 0.0
     (Token_bucket.peek_with_timestamp Mtime.Span.zero tb)
 
 let test_delay_until_available () =
   let initial_time = Mtime.Span.of_uint64_ns 0L in
   let tb =
-    Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
-      ~fill_rate:2.0
+    Option.get
+      (Token_bucket.create_with_timestamp initial_time ~burst_size:10.0
+         ~fill_rate:2.0
+      )
   in
 
   let _ =
@@ -201,57 +224,40 @@ let test_delay_until_available () =
   Alcotest.(check (float 0.01))
     "Delay for 4 tokens at 2 tokens/sec should be 2 seconds" 2.0 delay ;
 
-  let tb_fresh = Token_bucket.create ~burst_size:10.0 ~fill_rate:2.0 in
+  let tb_fresh =
+    Option.get (Token_bucket.create ~burst_size:10.0 ~fill_rate:2.0)
+  in
   let _ = Token_bucket.consume tb_fresh 10.0 in
   let delay_system = Token_bucket.delay_until_available tb_fresh 4.0 in
 
   Alcotest.(check (float 0.1))
     "System time delay should be approximately 2 seconds" 2.0 delay_system
 
 let test_edge_cases () =
-  let tb_zero_rate =
-    Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:5.0
-      ~fill_rate:0.0
-  in
-  let _ =
-    Token_bucket.consume_with_timestamp
-      (fun () -> Mtime.Span.zero)
-      tb_zero_rate 2.0
-  in
-  let later_time = Mtime.Span.of_uint64_ns 1_000_000_000_000L in
-  let available = Token_bucket.peek_with_timestamp later_time tb_zero_rate in
-  Alcotest.(check (float 0.0))
-    "Zero fill rate should not add tokens" 3.0 available ;
-
-  let tb_zero_amount =
-    Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:5.0
-      ~fill_rate:1.0
+  let tb =
+    Option.get
+      (Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:5.0
+         ~fill_rate:1.0
+      )
   in
   let success =
-    Token_bucket.consume_with_timestamp
-      (fun () -> Mtime.Span.zero)
-      tb_zero_amount 0.0
+    Token_bucket.consume_with_timestamp (fun () -> Mtime.Span.zero) tb 0.0
   in
   Alcotest.(check bool) "Consuming zero tokens should succeed" true success ;
 
   let tb_small =
-    Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:1.0
-      ~fill_rate:0.1
+    Option.get
+      (Token_bucket.create_with_timestamp Mtime.Span.zero ~burst_size:1.0
+         ~fill_rate:0.1
+      )
   in
   let success_small =
     Token_bucket.consume_with_timestamp
       (fun () -> Mtime.Span.zero)
       tb_small 0.001
   in
   Alcotest.(check bool)
-    "Consuming very small amount should succeed" true success_small ;
-
-  let tb_zero = Token_bucket.create ~burst_size:0.0 ~fill_rate:0.0 in
-  let success_zero = Token_bucket.consume tb_zero 0.0 in
-  let success_small = Token_bucket.consume tb_zero 0.001 in
-  Alcotest.(check bool) "Consuming zero tokens should succeed" true success_zero ;
-  Alcotest.(check bool)
-    "Consuming very small amount should fail" false success_small
+    "Consuming very small amount should succeed" true success_small
 
 let test_consume_quickcheck =
   let open QCheck.Gen in
@@ -289,7 +295,8 @@ let test_consume_quickcheck =
   let property (burst_size, fill_rate, operations) =
     let initial_time = Mtime.Span.of_uint64_ns 0L in
     let tb =
-      Token_bucket.create_with_timestamp initial_time ~burst_size ~fill_rate
+      Option.get
+        (Token_bucket.create_with_timestamp initial_time ~burst_size ~fill_rate)
     in
 
     let rec check_operations op_num time_ns last_refill_ns current_tokens ops =
@@ -345,7 +352,9 @@ let test_consume_quickcheck =
   in
 
   let gen_all =
-    map3 (fun burst fill ops -> (burst, fill, ops)) pfloat pfloat gen_operations
+    map3
+      (fun burst fill ops -> (burst, fill, ops))
+      pfloat (float_range 1e-9 1e9) gen_operations
   in
 
   let arb_all =
@@ -371,7 +380,11 @@ let test_consume_quickcheck =
 
 let test =
   [
-    ( "Consume removes correct amount"
+    ( "A bucket with zero or negative fill rate cannot be created"
+    , `Quick
+    , test_bad_fill_rate
+    )
+  ; ( "Consume removes correct amount"
     , `Quick
     , test_consume_removes_correct_amount
     )

ocaml/libs/rate-limit/token_bucket.ml

@@ -23,13 +23,17 @@ type t = {
 }
 
 let create_with_timestamp timestamp ~burst_size ~fill_rate =
-  {
-    burst_size
-  ; fill_rate
-  ; tokens= burst_size
-  ; last_refill= timestamp
-  ; mutex= Mutex.create ()
-  }
+  if fill_rate <= 0. then
+    None
+  else
+    Some
+      {
+        burst_size
+      ; fill_rate
+      ; tokens= burst_size
+      ; last_refill= timestamp
+      ; mutex= Mutex.create ()
+      }
 
 let create = create_with_timestamp (Mtime_clock.elapsed ())
 

ocaml/libs/rate-limit/token_bucket.mli

@@ -34,7 +34,7 @@
 
 type t
 
-val create : burst_size:float -> fill_rate:float -> t
+val create : burst_size:float -> fill_rate:float -> t option
 (** Create token bucket with given parameters.
     @param burst_size Maximum number of tokens that can fit in the bucket
     @param fill_rate Number of tokens added to the bucket per second
@@ -66,7 +66,7 @@ val delay_until_available : t -> float -> float
 (* Fuctions accepting a timestamp are meant for testing only *)
 
 val create_with_timestamp :
-  Mtime.span -> burst_size:float -> fill_rate:float -> t
+  Mtime.span -> burst_size:float -> fill_rate:float -> t option
 (** Create token bucket with given parameters and supplied inital timestamp
     @param timestamp Initial timestamp
     @param burst_size Maximum number of tokens that can fit in the bucket

ocaml/xapi/xapi_rate_limit.ml

@@ -27,8 +27,15 @@ let register_xapi_globs () =
               "Adding user agent %s to bucket table with burst size %f and \
                fill rate %f"
               user_agent burst_size fill_rate ;
-            Rate_limit.Bucket_table.add_bucket bucket_table ~user_agent
-              ~burst_size ~fill_rate
+            if
+              not
+                (Rate_limit.Bucket_table.add_bucket bucket_table ~user_agent
+                   ~burst_size ~fill_rate
+                )
+            then
+              D.error
+                "Bucket creation failed for user agent %s: invalid fill rate %f"
+                user_agent fill_rate
         | _ ->
             D.debug "Skipping invalid numeric values in: %s\n" s
       )