fixup! Address PR comments

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

Author: cplaursen

ocaml/libs/rate-limit/dune

@@ -1,7 +1,7 @@
 (library
  (name xapi_rate_limit)
  (public_name xapi-rate-limit)
- (wrapped false)
+
  (libraries threads.posix mtime mtime.clock.os xapi-log xapi-stdext-threads clock)
 )
 

ocaml/libs/rate-limit/rate_limit.ml

@@ -12,15 +12,17 @@
  * GNU Lesser General Public License for more details.
  *)
 
-module D = Debug.Make (struct let name = "rate_limit" end)
+module D = Debug.Make (struct let name = __MODULE__ end)
 
 type t = {
     bucket: Token_bucket.t
   ; process_queue:
       (float * (unit -> unit)) Queue.t (* contains token cost and callback *)
   ; process_queue_lock: Mutex.t
   ; worker_thread_cond: Condition.t
-  ; should_terminate: bool ref (* signal termination to worker thread *)
+  ; should_terminate: bool ref
+        (* Signal termination to worker thread. The worker thread will
+         process all remaining items in the queue before exiting. *)
   ; worker_thread: Thread.t
 }
 
@@ -52,30 +54,27 @@ let rec worker_loop ~bucket ~process_queue ~process_queue_lock
         ~should_terminate
 
 let create ~burst_size ~fill_rate =
-  match Token_bucket.create ~burst_size ~fill_rate with
-  | Some bucket ->
-      let process_queue = Queue.create () in
-      let process_queue_lock = Mutex.create () in
-      let worker_thread_cond = Condition.create () in
-      let should_terminate = ref false in
-      let worker_thread =
-        Thread.create
-          (fun () ->
-            worker_loop ~bucket ~process_queue ~process_queue_lock
-              ~worker_thread_cond ~should_terminate
-          )
-          ()
-      in
-      {
-        bucket
-      ; process_queue
-      ; process_queue_lock
-      ; worker_thread_cond
-      ; should_terminate
-      ; worker_thread
-      }
-  | None ->
-      raise (Failure "Invalid token bucket parameters")
+  let bucket = Token_bucket.create ~burst_size ~fill_rate in
+  let process_queue = Queue.create () in
+  let process_queue_lock = Mutex.create () in
+  let worker_thread_cond = Condition.create () in
+  let should_terminate = ref false in
+  let worker_thread =
+    Thread.create
+      (fun () ->
+        worker_loop ~bucket ~process_queue ~process_queue_lock
+          ~worker_thread_cond ~should_terminate
+      )
+      ()
+  in
+  {
+    bucket
+  ; process_queue
+  ; process_queue_lock
+  ; worker_thread_cond
+  ; should_terminate
+  ; worker_thread
+  }
 
 let delete data =
   with_lock data.process_queue_lock (fun () ->
@@ -84,11 +83,22 @@ let delete data =
   ) ;
   Thread.join data.worker_thread
 
+let check_not_terminated should_terminate =
+  if !should_terminate then
+    invalid_arg "Rate_limit: submit called on a deleted rate limiter"
+
 (* The callback should return quickly - if it is a longer task it is
    responsible for creating a thread to do the task *)
 let submit_async
-    ({bucket; process_queue; process_queue_lock; worker_thread_cond; _} as _data)
-    ~callback amount =
+    {
+      bucket
+    ; process_queue
+    ; process_queue_lock
+    ; worker_thread_cond
+    ; should_terminate
+    ; _
+    } ~callback amount =
+  check_not_terminated should_terminate ;
   let run_immediately =
     with_lock process_queue_lock (fun () ->
         let immediate =
@@ -104,10 +114,11 @@ let submit_async
   if run_immediately then
     callback ()
   else
-    D.debug "rate limiting sync call"
+    D.debug "%s: rate limiting call" __FUNCTION__
 
 (* Block and execute on the same thread *)
 let submit_sync bucket_data ~callback amount =
+  check_not_terminated bucket_data.should_terminate ;
   let channel_opt =
     with_lock bucket_data.process_queue_lock (fun () ->
         if

ocaml/libs/rate-limit/rate_limit.mli

@@ -20,14 +20,17 @@ type t
 
 val create : burst_size:float -> fill_rate:float -> t
 (** [create ~burst_size ~fill_rate] creates a new rate limiter with the given
-    token bucket parameters. Raises [Failure] if the parameters are invalid
+    token bucket parameters.
+    @raises Invalid_argument if the parameters are invalid
     (e.g. non-positive fill rate).
     @param burst_size Maximum number of tokens in the bucket
     @param fill_rate Number of tokens added per second *)
 
 val delete : t -> unit
-(** [delete t] signals the worker thread to terminate and waits for it to
-    finish processing any remaining queued callbacks. *)
+(** [delete t] signals the worker thread to terminate. The worker thread
+    processes any remaining queued callbacks, then exits. Blocks the caller
+    until the worker thread has finished. Subsequent calls to [submit_async]
+    or [submit_sync] will raise [Invalid_argument]. *)
 
 val submit_async : t -> callback:(unit -> unit) -> float -> unit
 (** [submit_async t ~callback amount] submits a callback under rate limiting.

ocaml/libs/rate-limit/token_bucket.ml

@@ -18,10 +18,10 @@ type t = {burst_size: float; fill_rate: float; state: state Atomic.t}
 
 let create_with_timestamp timestamp ~burst_size ~fill_rate =
   if fill_rate <= 0. then
-    None
+    invalid_arg "Token_bucket.create: fill_rate must be positive"
   else
     let state = Atomic.make {tokens= burst_size; last_refill= timestamp} in
-    Some {burst_size; fill_rate; state}
+    {burst_size; fill_rate; state}
 
 let create = create_with_timestamp (Mtime_clock.elapsed ())
 

ocaml/libs/rate-limit/token_bucket.mli

@@ -18,6 +18,9 @@
     up to [burst_size]. Consumers may take tokens (if available), or query when
     enough tokens will become available.
 
+    Token counts are represented as floats rather than integers to allow
+    fractional token costs and fine-grained refill rates.
+
     Token buckets implement rate limiting by allowing operations to proceed
     only when sufficient tokens are available - otherwise, the operations can
     be delayed until enough tokens are available.
@@ -34,9 +37,9 @@
 
 type t
 
-val create : burst_size:float -> fill_rate:float -> t option
+val create : burst_size:float -> fill_rate:float -> t
 (** Create token bucket with given parameters.
-    Returns None if the fill rate is 0 or negative.
+    @raises Invalid_argument if the fill rate is 0 or negative.
     @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
     *)
@@ -63,15 +66,18 @@ val get_delay_until_available : t -> float -> float
 *)
 
 val delay_then_consume : t -> float -> unit
+(** [delay_then_consume tb amount] sleeps the calling thread until [amount]
+    tokens are available, then consumes them. Thread-safe but does not
+    guarantee fairness between competing callers. *)
 
 (**/**)
 
 (* Fuctions accepting a timestamp are meant for testing only *)
 
 val create_with_timestamp :
-  Mtime.span -> burst_size:float -> fill_rate:float -> t option
-(** Create token bucket with given parameters and supplied inital timestamp
-    Returns None if the fill_rate is 0 or negative.
+  Mtime.span -> burst_size:float -> fill_rate:float -> t
+(** Create token bucket with given parameters and supplied inital timestamp.
+    @raises Invalid_argument if the fill_rate is 0 or negative.
     @param timestamp Initial timestamp
     @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