Merge pull request #9 from AestheticIntegration/matt/cohttp

Integration with cohttp-lwt

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"))

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"

src/integrations/cohttp/README.md

@@ -0,0 +1,11 @@
+# Opentelemetry tracing for Cohttp_lwt servers
+
+Wrap your server callback with `Opentelemetry_cohttp_lwt.Server.trace`:
+
+```ocaml
+let my_server callback =
+  let callback =
+    Opentelemetry_cohttp_lwt.Server.trace ~service_name:"my-service" callback in
+  Cohttp_lwt_unix.Server.create ~mode:(`TCP (`Port 8080))
+    (Server.make () ~callback)
+```

src/integrations/cohttp/dune

@@ -0,0 +1,4 @@
+(library
+ (name opentelemetry_cohttp_lwt)
+ (public_name opentelemetry-cohttp-lwt)
+ (libraries cohttp-lwt opentelemetry opentelemetry-lwt))

src/integrations/cohttp/opentelemetry_cohttp_lwt.ml

@@ -0,0 +1,147 @@
+module Otel = Opentelemetry
+module Otel_lwt = Opentelemetry_lwt
+open Cohttp
+open Cohttp_lwt
+
+module Server : sig
+  (** Trace requests to a Cohttp server.
+
+      Use it like this:
+
+          let my_server callback =
+            let callback_traced =
+              Opentelemetry_cohttp_lwt.Server.trace
+                ~service_name:"my-service"
+                (fun _scope -> callback)
+            in
+            Cohttp_lwt_unix.Server.create
+              ~mode:(`TCP (`Port 8080))
+              (Server.make () ~callback:callback_traced)
+   *)
+  val trace :
+    service_name:string ->
+    ?attrs:Otel.Span.key_value list ->
+    ('conn -> Request.t -> 'body -> (Response.t * 'body) Lwt.t) ->
+    'conn -> Request.t -> 'body -> (Response.t * 'body) Lwt.t
+
+  (** Trace a new internal span.
+
+      Identical to [Opentelemetry_lwt.Trace.with_], but fetches/stores the trace
+      scope in the [x-ocaml-otel-traceparent] header in the request for
+      convenience.
+   *)
+  val with_:
+    ?trace_state:string ->
+    ?service_name:string ->
+    ?attrs:Otel.Span.key_value list ->
+    ?kind:Otel.Span.kind ->
+    ?links:(Otel.Trace_id.t * Otel.Span_id.t * string) list ->
+    string ->
+    Request.t ->
+    (Request.t -> 'a Lwt.t) ->
+    'a Lwt.t
+
+  (** Get the tracing scope from the custom [x-ocaml-otel-traceparent] header
+      added by [trace] and [with_].
+   *)
+  val get_trace_context : ?from:[`Internal | `External] -> Request.t -> Otel.Trace.scope option
+
+  (** Set the tracing scope in the custom [x-ocaml-otel-traceparent] header used
+      by [trace] and [with_].
+   *)
+  val set_trace_context : Otel.Trace.scope -> Request.t -> Request.t
+
+  (** Strip the custom [x-ocaml-otel-traceparent] header added by [trace] and
+      [with_].
+   *)
+  val remove_trace_context : Request.t -> Request.t
+end = struct
+  let attrs_of_request (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 attrs_of_response (res : Response.t) =
+    let code = Response.status res in
+    let code = Code.code_of_status code in
+    [ ("http.status_code", `Int code) ]
+
+  let header_x_ocaml_otel_traceparent = "x-ocaml-otel-traceparent"
+
+  let set_trace_context (scope : Otel.Trace.scope) req =
+    let module Traceparent = Otel.Trace_context.Traceparent in
+    let headers =
+      Header.add (Request.headers req) header_x_ocaml_otel_traceparent
+        (Traceparent.to_value ~trace_id:scope.trace_id ~parent_id:scope.span_id ())
+    in
+    { req with headers }
+
+  let get_trace_context ?(from=`Internal) req =
+    let module Traceparent = Otel.Trace_context.Traceparent in
+    let name =
+      match from with
+      | `Internal -> header_x_ocaml_otel_traceparent
+      | `External -> Traceparent.name
+    in
+    match Header.get (Request.headers req) name with
+    | None -> None
+    | Some v ->
+       (match Traceparent.of_value v with
+        | Ok (trace_id, parent_id) ->
+           (Some Otel.Trace.{ trace_id; span_id = parent_id; events = []; attrs = []})
+        | Error _ -> None)
+
+  let remove_trace_context req =
+    let headers = Header.remove (Request.headers req) header_x_ocaml_otel_traceparent in
+    { req with headers }
+
+  let trace ~service_name ?(attrs=[]) callback =
+    fun conn req body ->
+    let scope = get_trace_context ~from:`External req in
+    Otel_lwt.Trace.with_
+      ~service_name
+      "request"
+      ~kind:Span_kind_server
+      ?trace_id:(Option.map (fun scope -> scope.Otel.Trace.trace_id) scope)
+      ?parent:(Option.map (fun scope -> scope.Otel.Trace.span_id) scope)
+      ~attrs:(attrs @ attrs_of_request req)
+      (fun scope ->
+        let open Lwt.Syntax in
+        let req = set_trace_context scope req in
+        let* (res, body) = callback conn req body in
+        Otel.Trace.add_attrs scope (fun () -> attrs_of_response res) ;
+        Lwt.return (res, body) )
+
+  let with_ ?trace_state ?service_name ?attrs ?(kind=Otel.Span.Span_kind_internal) ?links name req (f : Request.t -> 'a Lwt.t) =
+    let scope = get_trace_context ~from:`Internal req in
+    Otel_lwt.Trace.with_
+      ?trace_state
+      ?service_name
+      ?attrs
+      ~kind
+      ?trace_id:(Option.map (fun scope -> scope.Otel.Trace.trace_id) scope)
+      ?parent:(Option.map (fun scope -> scope.Otel.Trace.span_id) scope)
+      ?links
+      name
+      (fun scope ->
+        let open Lwt.Syntax in
+        let req = set_trace_context scope req in
+        f req)
+end