From a1115661b4bf985ccc144b153da4aa7dca0e5cb3 Mon Sep 17 00:00:00 2001
From: Matt Bray <mattjbray@gmail.com>
Date: Thu, 24 Mar 2022 13:46:06 +0000
Subject: [PATCH] wip: cohttp integration

---
 dune-project                                  | 11 ++++
 opentelemetry-cohttp-lwt.opam                 | 33 +++++++++++
 src/integrations/cohttp/README.md             | 11 ++++
 src/integrations/cohttp/dune                  |  4 ++
 .../cohttp/opentelemetry_cohttp_lwt.ml        | 59 +++++++++++++++++++
 5 files changed, 118 insertions(+)
 create mode 100644 opentelemetry-cohttp-lwt.opam
 create mode 100644 src/integrations/cohttp/README.md
 create mode 100644 src/integrations/cohttp/dune
 create mode 100644 src/integrations/cohttp/opentelemetry_cohttp_lwt.ml

diff --git a/dune-project b/dune-project
index 21e3f5a2..8d819808 100644
--- a/dune-project
+++ b/dune-project
@@ -45,3 +45,14 @@
    (odoc :with-doc)
    ocurl)
   (synopsis "Collector client for opentelemetry, using http + ocurl"))
+
+(package
+ (name opentelemetry-cohttp-lwt)
+ (depends
+  (ocaml (>= "4.08"))
+  (dune (>= "2.3"))
+  (opentelemetry (= :version))
+  (opentelemetry-lwt (= :version))
+  (odoc :with-doc)
+  cohttp-lwt)
+ (synopsis "Opentelemetry tracing for Cohttp HTTP servers"))
diff --git a/opentelemetry-cohttp-lwt.opam b/opentelemetry-cohttp-lwt.opam
new file mode 100644
index 00000000..af591c09
--- /dev/null
+++ b/opentelemetry-cohttp-lwt.opam
@@ -0,0 +1,33 @@
+# This file is generated by dune, edit dune-project instead
+opam-version: "2.0"
+synopsis: "Opentelemetry tracing for Cohttp HTTP servers"
+maintainer: ["the Imandra team"]
+authors: ["the Imandra team"]
+license: "MIT"
+homepage: "https://github.com/aestheticintegration/ocaml-opentelemetry"
+bug-reports:
+  "https://github.com/aestheticintegration/ocaml-opentelemetry/issues"
+depends: [
+  "ocaml" {>= "4.08"}
+  "dune" {>= "2.3"}
+  "opentelemetry" {= version}
+  "opentelemetry-lwt" {= version}
+  "odoc" {with-doc}
+  "cohttp-lwt"
+]
+build: [
+  ["dune" "subst"] {pinned}
+  [
+    "dune"
+    "build"
+    "-p"
+    name
+    "-j"
+    jobs
+    "@install"
+    "@runtest" {with-test}
+    "@doc" {with-doc}
+  ]
+]
+dev-repo:
+  "git+https://github.com/aestheticintegration/ocaml-opentelemetry.git"
diff --git a/src/integrations/cohttp/README.md b/src/integrations/cohttp/README.md
new file mode 100644
index 00000000..4b500cde
--- /dev/null
+++ b/src/integrations/cohttp/README.md
@@ -0,0 +1,11 @@
+# Opentelemetry tracing for Cohttp_lwt servers
+
+Wrap your server callback with `Opentelemetry_cohttp_lwt.trace`:
+
+```ocaml
+let my_server callback =
+  Cohttp_lwt_unix.Server.create ~mode:(`TCP (`Port 8080))
+    (Server.make 
+      ~callback:(Opentelemetry_cohttp_lwt.trace ~service_name:"my-service" callback)
+      ())
+```
diff --git a/src/integrations/cohttp/dune b/src/integrations/cohttp/dune
new file mode 100644
index 00000000..16a92278
--- /dev/null
+++ b/src/integrations/cohttp/dune
@@ -0,0 +1,4 @@
+(library
+ (name opentelemetry_cohttp_lwt)
+ (public_name opentelemetry-cohttp-lwt)
+ (libraries cohttp-lwt opentelemetry opentelemetry-lwt))
diff --git a/src/integrations/cohttp/opentelemetry_cohttp_lwt.ml b/src/integrations/cohttp/opentelemetry_cohttp_lwt.ml
new file mode 100644
index 00000000..cb2f9d40
--- /dev/null
+++ b/src/integrations/cohttp/opentelemetry_cohttp_lwt.ml
@@ -0,0 +1,59 @@
+open Cohttp
+open Cohttp_lwt
+
+type ('conn, 'body) callback =
+  'conn (* Cohttp_lwt_unix.Server.conn *)
+  -> Request.t
+  -> 'body (* Cohttp_lwt.Body.t *)
+  -> (Response.t * 'body) Lwt.t
+
+let span_attrs (req : Request.t) =
+  let meth = req |> Request.meth |> Code.string_of_method in
+  let referer = Header.get (Request.headers req) "referer" in
+  let host = Header.get (Request.headers req) "host" in
+  let ua = Header.get (Request.headers req) "user-agent" in
+  let uri = Request.uri req in
+  List.concat
+    [ [ ("http.method", `String meth) ]
+    ; (match host with None -> [] | Some h -> [ ("http.host", `String h) ])
+    ; [ ("http.url", `String (Uri.to_string uri)) ]
+    ; ( match ua with
+        | None ->
+           []
+        | Some ua ->
+           [ ("http.user_agent", `String ua) ] )
+    ; ( match referer with
+        | None ->
+           []
+        | Some r ->
+           [ ("http.request.header.referer", `String r) ] )
+    ]
+
+
+let trace ~service_name (callback : ('conn, 'body) callback ) : ('conn, 'body) callback =
+  fun conn req body ->
+  let trace_id =
+    Header.get (Request.headers req) "trace-id"
+    |> Option.map (fun s ->
+           s |> Bytes.of_string |> Opentelemetry.Trace_id.of_bytes )
+  in
+  let parent_id =
+    Header.get (Request.headers req) "parent-id"
+    |> Option.map (fun s ->
+           s |> Bytes.of_string |> Opentelemetry.Span_id.of_bytes )
+  in
+  let open Lwt.Syntax in
+  Opentelemetry_lwt.Trace.with_
+    ~service_name
+    "request"
+    ~kind:Span_kind_server
+    ~attrs:(span_attrs req)
+    ?parent:parent_id
+    ?trace_id
+    (fun scope ->
+      let* (res, body) = callback conn req body in
+      Opentelemetry.Trace.add_attrs scope (fun () ->
+          let code = Response.status res in
+          let code = Code.code_of_status code in
+          [ ("http.status_code", `Int code) ]) ;
+      Lwt.return (res, body) )