From c19b8dc16f80cb982b254de0b9862fc8703803f6 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 18 Jan 2024 22:05:23 -0500
Subject: [PATCH] add histograms to prometheus

---
 dune-project                             |  2 +
 src/dune                                 |  2 +-
 src/prometheus/dune                      | 15 +++--
 src/prometheus/time_.default.ml          |  3 +
 src/prometheus/time_.mli                 |  1 +
 src/prometheus/time_.mtime.ml            |  3 +
 src/prometheus/tiny_httpd_prometheus.ml  | 84 +++++++++++++++++++++++-
 src/prometheus/tiny_httpd_prometheus.mli | 15 +++++
 tests/prometheus/t_prom.expected         | 20 ++++++
 tests/prometheus/t_prom.ml               | 17 +++++
 tiny_httpd.opam                          |  3 +
 11 files changed, 157 insertions(+), 8 deletions(-)
 create mode 100644 src/prometheus/time_.default.ml
 create mode 100644 src/prometheus/time_.mli
 create mode 100644 src/prometheus/time_.mtime.ml

diff --git a/dune-project b/dune-project
index 415fe81c..95927d24 100644
--- a/dune-project
+++ b/dune-project
@@ -13,6 +13,8 @@
   (name tiny_httpd)
   (synopsis "Minimal HTTP server using threads")
   (tags (http thread server tiny_httpd http_of_dir simplehttpserver))
+  (depopts
+    (mtime (>= 2.0)))
   (depends
     seq
     base-threads
diff --git a/src/dune b/src/dune
index 89ba2327..680ebc43 100644
--- a/src/dune
+++ b/src/dune
@@ -7,7 +7,7 @@
 (library
   (name tiny_httpd)
   (public_name tiny_httpd)
-  (libraries threads seq)
+  (libraries threads seq unix)
   (wrapped false))
 
 (rule
diff --git a/src/prometheus/dune b/src/prometheus/dune
index 9cf17df0..5da724e5 100644
--- a/src/prometheus/dune
+++ b/src/prometheus/dune
@@ -1,7 +1,12 @@
 
+
 (library
-  (name tiny_httpd_prometheus)
-  (public_name tiny_httpd.prometheus)
-  (synopsis "Metrics using prometheus")
-  (private_modules common_)
-  (libraries tiny_httpd))
+ (name tiny_httpd_prometheus)
+ (public_name tiny_httpd.prometheus)
+ (synopsis "Metrics using prometheus")
+ (private_modules common_ time_)
+ (libraries
+  tiny_httpd unix
+  (select time_.ml from
+   (mtime mtime.clock.os -> time_.mtime.ml)
+   (-> time_.default.ml))))
diff --git a/src/prometheus/time_.default.ml b/src/prometheus/time_.default.ml
new file mode 100644
index 00000000..86dd302c
--- /dev/null
+++ b/src/prometheus/time_.default.ml
@@ -0,0 +1,3 @@
+let[@inline] now_us () =
+  let t = Unix.gettimeofday () in
+  t *. 1e6 |> ceil
diff --git a/src/prometheus/time_.mli b/src/prometheus/time_.mli
new file mode 100644
index 00000000..d6824fba
--- /dev/null
+++ b/src/prometheus/time_.mli
@@ -0,0 +1 @@
+val now_us : unit -> float
diff --git a/src/prometheus/time_.mtime.ml b/src/prometheus/time_.mtime.ml
new file mode 100644
index 00000000..65e2ec73
--- /dev/null
+++ b/src/prometheus/time_.mtime.ml
@@ -0,0 +1,3 @@
+let[@inline] now_us () =
+  let t = Mtime_clock.now_ns () in
+  Int64.(div t 1000L |> to_float)
diff --git a/src/prometheus/tiny_httpd_prometheus.ml b/src/prometheus/tiny_httpd_prometheus.ml
index 249d14f4..bbede737 100644
--- a/src/prometheus/tiny_httpd_prometheus.ml
+++ b/src/prometheus/tiny_httpd_prometheus.ml
@@ -9,7 +9,20 @@ let bpf = Printf.bprintf
 type tags = (string * string) list
 type counter = { name: string; tags: tags; descr: string option; c: int A.t }
 type gauge = { name: string; tags: tags; descr: string option; g: int A.t }
-type registry = { mutable counters: counter list; mutable gauges: gauge list }
+
+type histogram = {
+  name: string;
+  tags: tags;
+  descr: string option;
+  sum: float A.t;
+  buckets: (float * int A.t) array;
+}
+
+type registry = {
+  mutable counters: counter list;
+  mutable gauges: gauge list;
+  mutable hists: histogram list;
+}
 
 let validate_descr_ what s =
   if String.contains s '\n' then
@@ -73,14 +86,71 @@ module Gauge = struct
   let[@inline] decr_by self n = ignore (A.fetch_and_add self.g (-n) : int)
 end
 
+module Histogram = struct
+  type t = histogram
+
+  let create reg ?(tags = []) ?descr ~buckets name : t =
+    opt_iter_ (validate_descr_ "histogram") descr;
+    let buckets =
+      List.sort Float.compare buckets
+      |> List.map (fun thresh -> thresh, A.make 0)
+    in
+    let buckets = Array.of_list @@ buckets @ [ infinity, A.make 0 ] in
+    let self : t = { name; descr; tags; sum = A.make 0.; buckets } in
+    reg.hists <- self :: reg.hists;
+    self
+
+  let add (self : t) n =
+    while
+      let old = A.get self.sum in
+      not (A.compare_and_set self.sum old (old +. n))
+    do
+      ()
+    done;
+    let i = ref 0 in
+    let continue = ref true in
+    while !continue && !i < Array.length self.buckets do
+      let thresh, count = self.buckets.(!i) in
+      if n <= thresh then (
+        continue := false;
+        A.incr count
+      ) else
+        incr i
+    done
+
+  let emit buf (self : t) : unit =
+    opt_iter_ (bpf buf "# HELP %s %s\n" self.name) self.descr;
+    bpf buf "# TYPE %s histogram\n" self.name;
+
+    let count = ref 0 in
+    for i = 0 to Array.length self.buckets - 1 do
+      let thresh, buck_count = self.buckets.(i) in
+      count := !count + A.get buck_count;
+
+      let name =
+        if thresh = infinity then
+          "+Inf"
+        else
+          string_of_float thresh
+      in
+      bpf buf "%s%a %d\n" self.name emit_tags_
+        (("le", name) :: self.tags)
+        !count
+    done;
+    bpf buf "%s_count%a %d\n" self.name emit_tags_ self.tags !count;
+    bpf buf "%s_sum%a %.3f\n" self.name emit_tags_ self.tags (A.get self.sum);
+    ()
+end
+
 module Registry = struct
   type t = registry
 
-  let create () : t = { counters = []; gauges = [] }
+  let create () : t = { counters = []; gauges = []; hists = [] }
 
   let emit (buf : Buffer.t) (self : t) : unit =
     List.iter (Gauge.emit buf) self.gauges;
     List.iter (Counter.emit buf) self.counters;
+    List.iter (Histogram.emit buf) self.hists;
     ()
 
   let emit_str (self : t) : string =
@@ -100,12 +170,22 @@ let http_middleware (reg : Registry.t) : H.Middleware.t =
   let c_err =
     Counter.create reg "tiny_httpd_errors" ~descr:"number of HTTP errors"
   in
+  let h_latency =
+    Histogram.create reg "tiny_httpd_latency" ~descr:"latency of HTTP responses"
+      ~buckets:[ 0.001; 0.01; 0.1; 0.5; 1.; 5.; 10. ]
+  in
 
   fun h : H.Middleware.handler ->
     fun req ~resp : unit ->
+     let start = Time_.now_us () in
      Counter.incr c_req;
      h req ~resp:(fun (response : H.Response.t) ->
          let code = response.code in
+
+         let elapsed_us = Time_.now_us () -. start in
+         let elapsed_s = elapsed_us /. 1e6 in
+         Histogram.add h_latency elapsed_s;
+
          if code < 200 || code >= 300 then Counter.incr c_err;
          resp response)
 
diff --git a/src/prometheus/tiny_httpd_prometheus.mli b/src/prometheus/tiny_httpd_prometheus.mli
index a2104af1..e82a5e13 100644
--- a/src/prometheus/tiny_httpd_prometheus.mli
+++ b/src/prometheus/tiny_httpd_prometheus.mli
@@ -43,6 +43,21 @@ module Gauge : sig
   val decr_by : t -> int -> unit
 end
 
+module Histogram : sig
+  type t
+  (** Histogram *)
+
+  val create :
+    Registry.t ->
+    ?tags:tags ->
+    ?descr:string ->
+    buckets:float list ->
+    string ->
+    t
+
+  val add : t -> float -> unit
+end
+
 (* TODO:
    module Histogram : sig
    end
diff --git a/tests/prometheus/t_prom.expected b/tests/prometheus/t_prom.expected
index 590b6c82..0305272e 100644
--- a/tests/prometheus/t_prom.expected
+++ b/tests/prometheus/t_prom.expected
@@ -7,6 +7,16 @@ yolo_gauge{level="max"} 2525
 t_c2 1
 # TYPE t_c1 counter
 t_c1 42
+# HELP latency latency
+# TYPE latency histogram
+latency{le="0.01"} 2
+latency{le="0.1"} 4
+latency{le="0.5"} 7
+latency{le="1."} 8
+latency{le="10."} 9
+latency{le="+Inf"} 10
+latency_count 10
+latency_sum 31.530
 
 ```
 ==== second try====
@@ -18,5 +28,15 @@ yolo_gauge{level="max"} 42000
 t_c2 2
 # TYPE t_c1 counter
 t_c1 53
+# HELP latency latency
+# TYPE latency histogram
+latency{le="0.01"} 2
+latency{le="0.1"} 4
+latency{le="0.5"} 8
+latency{le="1."} 9
+latency{le="10."} 10
+latency{le="+Inf"} 12
+latency_count 12
+latency_sum 54.930
 
 ```
diff --git a/tests/prometheus/t_prom.ml b/tests/prometheus/t_prom.ml
index f8144b7b..04b1639e 100644
--- a/tests/prometheus/t_prom.ml
+++ b/tests/prometheus/t_prom.ml
@@ -6,11 +6,26 @@ let c1 = P.Counter.create reg "t_c1"
 let c2 = P.Counter.create reg "t_c2" ~descr:"more awesome than c1"
 let g1 = P.Gauge.create reg ~tags:[ "level", "max" ] "yolo_gauge"
 
+let h1 =
+  P.Histogram.create reg ~descr:"latency"
+    ~buckets:[ 0.01; 0.1; 0.5; 1.; 10. ]
+    "latency"
+
 let () =
   print_endline "==== first try ====";
   P.Counter.incr_by c1 42;
   P.Counter.incr c2;
   P.Gauge.set g1 2525;
+  P.Histogram.add h1 0.2;
+  P.Histogram.add h1 0.003;
+  P.Histogram.add h1 0.002;
+  P.Histogram.add h1 0.025;
+  P.Histogram.add h1 0.9;
+  P.Histogram.add h1 7.4;
+  P.Histogram.add h1 22.2;
+  P.Histogram.add h1 0.3;
+  P.Histogram.add h1 0.4;
+  P.Histogram.add h1 0.1;
 
   pf "```\n%s\n```\n" @@ P.Registry.emit_str reg
 
@@ -19,5 +34,7 @@ let () =
   P.Counter.incr_by c1 11;
   P.Counter.incr c2;
   P.Gauge.set g1 42_000;
+  P.Histogram.add h1 23.2;
+  P.Histogram.add h1 0.2;
 
   pf "```\n%s\n```\n" @@ P.Registry.emit_str reg
diff --git a/tiny_httpd.opam b/tiny_httpd.opam
index 3cda25ee..6a2d337b 100644
--- a/tiny_httpd.opam
+++ b/tiny_httpd.opam
@@ -21,6 +21,9 @@ depends: [
   "ptime" {with-test}
   "qcheck-core" {>= "0.9" & with-test}
 ]
+depopts: [
+  "mtime" {>= "2.0"}
+]
 build: [
   ["dune" "subst"] {dev}
   [