Merge pull request #75 from tatchi/add-set-span-status

add Scope.set_span_status

src/core/opentelemetry.ml

@@ -804,6 +804,24 @@ end = struct
       ?attrs ()
 end
 
+module Span_status : sig
+  open Proto.Trace
+
+  type t = status
+
+  type code = status_status_code
+
+  val make : message:string -> code:code -> t
+end = struct
+  open Proto.Trace
+
+  type t = status
+
+  type code = status_status_code
+
+  let make ~message ~code = { message; code }
+end
+
 (** {2 Scopes} *)
 
 (** Scopes.
@@ -825,36 +843,67 @@ module Scope : sig
 
   val links : t -> Span_link.t list
 
+  val status : t -> Span_status.t option
+
   val make :
     trace_id:Trace_id.t ->
     span_id:Span_id.t ->
     ?events:Event.t list ->
     ?attrs:key_value list ->
     ?links:Span_link.t list ->
+    ?status:Span_status.t ->
     unit ->
     t
 
   val to_span_ctx : t -> Span_ctx.t
+  (** Turn the scope into a span context *)
 
   val add_event : t -> (unit -> Event.t) -> unit
+  (** Add an event to the scope. It will be aggregated into the span.
+
+  Note that this takes a function that produces an event, and will only
+  call it if there is an instrumentation backend. *)
 
   val record_exception : t -> exn -> Printexc.raw_backtrace -> unit
 
   val add_attrs : t -> (unit -> key_value list) -> unit
+  (** Add attributes to the scope. It will be aggregated into the span.
+
+  Note that this takes a function that produces attributes, and will only
+  call it if there is an instrumentation backend. *)
 
   val add_links : t -> (unit -> Span_link.t list) -> unit
+  (** Add links to the scope. It will be aggregated into the span.
+
+  Note that this takes a function that produces links, and will only
+  call it if there is an instrumentation backend. *)
+
+  val set_status : t -> Span_status.t -> unit
+  (** set the span status.
+
+  Note that this function will be
+  called only if there is an instrumentation backend. *)
 
   val ambient_scope_key : t Ambient_context.key
+  (** The opaque key necessary to access/set the ambient scope with
+  {!Ambient_context}. *)
 
   val get_ambient_scope : ?scope:t -> unit -> t option
+  (** Obtain current scope from {!Ambient_context}, if available. *)
 
   val with_ambient_scope : t -> (unit -> 'a) -> 'a
+  (** [with_ambient_scope sc thunk] calls [thunk()] in a context where [sc] is
+  the (thread|continuation)-local scope, then reverts to the previous local
+  scope, if any.
+
+  @see <https://github.com/ELLIOTTCABLE/ocaml-ambient-context> ambient-context docs *)
 end = struct
   type item_list =
     | Nil
     | Ev of Event.t * item_list
     | Attr of key_value * item_list
     | Span_link of Span_link.t * item_list
+    | Span_status of Span_status.t * item_list
 
   type t = {
     trace_id: Trace_id.t;
@@ -866,7 +915,7 @@ end = struct
     let rec loop acc = function
       | Nil -> acc
       | Attr (attr, l) -> loop (attr :: acc) l
-      | Ev (_, l) | Span_link (_, l) -> loop acc l
+      | Ev (_, l) | Span_link (_, l) | Span_status (_, l) -> loop acc l
     in
     loop [] scope.items
 
@@ -874,7 +923,7 @@ end = struct
     let rec loop acc = function
       | Nil -> acc
       | Ev (event, l) -> loop (event :: acc) l
-      | Attr (_, l) | Span_link (_, l) -> loop acc l
+      | Attr (_, l) | Span_link (_, l) | Span_status (_, l) -> loop acc l
     in
     loop [] scope.items
 
@@ -882,14 +931,27 @@ end = struct
     let rec loop acc = function
       | Nil -> acc
       | Span_link (span_link, l) -> loop (span_link :: acc) l
-      | Ev (_, l) | Attr (_, l) -> loop acc l
+      | Ev (_, l) | Attr (_, l) | Span_status (_, l) -> loop acc l
     in
     loop [] scope.items
 
-  let make ~trace_id ~span_id ?(events = []) ?(attrs = []) ?(links = []) () : t
-      =
+  let status scope =
+    let rec loop acc = function
+      | Nil -> acc
+      | Span_status (status, _) -> Some status
+      | Ev (_, l) | Attr (_, l) | Span_link (_, l) -> loop acc l
+    in
+    loop None scope.items
+
+  let make ~trace_id ~span_id ?(events = []) ?(attrs = []) ?(links = []) ?status
+      () : t =
     let items =
-      let items = List.fold_left (fun acc ev -> Ev (ev, acc)) Nil events in
+      let items =
+        match status with
+        | None -> Nil
+        | Some status -> Span_status (status, Nil)
+      in
+      let items = List.fold_left (fun acc ev -> Ev (ev, acc)) items events in
       let items =
         List.fold_left (fun acc attr -> Attr (attr, acc)) items attrs
       in
@@ -897,14 +959,9 @@ end = struct
     in
     { trace_id; span_id; items }
 
-  (** Turn the scope into a span context *)
   let[@inline] to_span_ctx (self : t) : Span_ctx.t =
     Span_ctx.make ~trace_id:self.trace_id ~parent_id:self.span_id ()
 
-  (** Add an event to the scope. It will be aggregated into the span.
-
-      Note that this takes a function that produces an event, and will only
-      call it if there is an instrumentation backend. *)
   let[@inline] add_event (scope : t) (ev : unit -> Event.t) : unit =
     if Collector.has_backend () then scope.items <- Ev (ev (), scope.items)
 
@@ -923,19 +980,11 @@ end = struct
       scope.items <- Ev (ev, scope.items)
     )
 
-  (** Add attributes to the scope. It will be aggregated into the span.
-
-      Note that this takes a function that produces attributes, and will only
-      call it if there is an instrumentation backend. *)
   let[@inline] add_attrs (scope : t) (attrs : unit -> key_value list) : unit =
     if Collector.has_backend () then
       scope.items <-
         List.fold_left (fun acc attr -> Attr (attr, acc)) scope.items (attrs ())
 
-  (** Add links to the scope. It will be aggregated into the span.
-
-      Note that this takes a function that produces links, and will only
-      call it if there is an instrumentation backend. *)
   let[@inline] add_links (scope : t) (links : unit -> Span_link.t list) : unit =
     if Collector.has_backend () then
       scope.items <-
@@ -943,21 +992,18 @@ end = struct
           (fun acc link -> Span_link (link, acc))
           scope.items (links ())
 
-  (** The opaque key necessary to access/set the ambient scope with
-      {!Ambient_context}. *)
+  let set_status (scope : t) (status : Span_status.t) : unit =
+    if Collector.has_backend () then (
+      scope.items <- Span_status (status, scope.items)
+    )
+
   let ambient_scope_key : t Ambient_context.key = Ambient_context.create_key ()
 
-  (** Obtain current scope from {!Ambient_context}, if available. *)
   let get_ambient_scope ?scope () : t option =
     match scope with
     | Some _ -> scope
     | None -> Ambient_context.get ambient_scope_key
 
-  (** [with_ambient_scope sc thunk] calls [thunk()] in a context where [sc] is
-      the (thread|continuation)-local scope, then reverts to the previous local
-      scope, if any.
-
-      @see <https://github.com/ELLIOTTCABLE/ocaml-ambient-context> ambient-context docs *)
   let[@inline] with_ambient_scope (sc : t) (f : unit -> 'a) : 'a =
     Ambient_context.with_binding ambient_scope_key sc (fun _ -> f ())
 end
@@ -985,16 +1031,6 @@ module Span : sig
     | Span_kind_producer
     | Span_kind_consumer
 
-  type nonrec status_code = status_status_code =
-    | Status_code_unset
-    | Status_code_ok
-    | Status_code_error
-
-  type nonrec status = status = {
-    message: string;
-    code: status_code;
-  }
-
   val id : t -> Span_id.t
 
   type key_value =
@@ -1049,16 +1085,6 @@ end = struct
       | `None
       ]
 
-  type nonrec status_code = status_status_code =
-    | Status_code_unset
-    | Status_code_ok
-    | Status_code_error
-
-  type nonrec status = status = {
-    message: string;
-    code: status_code;
-  }
-
   let id self = Span_id.of_bytes self.span_id
 
   let create ?(kind = !Globals.default_span_kind) ?(id = Span_id.create ())
@@ -1143,13 +1169,22 @@ module Trace = struct
     (* called once we're done, to emit a span *)
     let finally res =
       let status =
-        match res with
-        | Ok () -> default_status ~code:Status_code_ok ()
-        | Error (e, bt) ->
-          (* add backtrace *)
-          Scope.record_exception scope e bt;
-          default_status ~code:Status_code_error ~message:(Printexc.to_string e)
-            ()
+        match Scope.status scope with
+        | Some status -> Some status
+        | None ->
+          (match res with
+          | Ok () ->
+            (* By default, all spans are Unset, which means a span completed without error.
+               The Ok status is reserved for when you need to explicitly mark a span as successful
+               rather than stick with the default of Unset (i.e., “without error”).
+
+                https://opentelemetry.io/docs/languages/go/instrumentation/#set-span-status *)
+            None
+          | Error (e, bt) ->
+            Scope.record_exception scope e bt;
+            Some
+              (default_status ~code:Status_code_error
+                 ~message:(Printexc.to_string e) ()))
       in
       let span, _ =
         (* TODO: should the attrs passed to with_ go on the Span
@@ -1159,7 +1194,7 @@ module Trace = struct
           ~id:span_id ?trace_state ~attrs:(Scope.attrs scope)
           ~events:(Scope.events scope) ~start_time
           ~end_time:(Timestamp_ns.now_unix_ns ())
-          ~status name
+          ?status name
       in
       emit ?service_name [ span ]
     in

src/integrations/cohttp/opentelemetry_cohttp_lwt.ml

@@ -130,7 +130,7 @@ end = struct
         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);
+        Otel.Scope.add_attrs scope (fun () -> attrs_of_response res);
         Lwt.return (res, body))
 
   let with_ ?trace_state ?service_name ?attrs
@@ -190,7 +190,7 @@ let client ?(scope : Otel.Scope.t option) (module C : Cohttp_lwt.S.Client) =
         ~attrs (fun scope ->
           let headers = add_traceparent scope headers in
           let* res, body = C.call ?ctx ~headers ?body ?chunked meth uri in
-          Otel.Trace.add_attrs scope (fun () ->
+          Otel.Scope.add_attrs scope (fun () ->
               let code = Response.status res in
               let code = Code.code_of_status code in
               [ "http.status_code", `Int code ]);
@@ -220,7 +220,7 @@ let client ?(scope : Otel.Scope.t option) (module C : Cohttp_lwt.S.Client) =
         ~attrs (fun scope ->
           let headers = add_traceparent scope headers in
           let* res, body = C.post_form ?ctx ~headers ~params uri in
-          Otel.Trace.add_attrs scope (fun () ->
+          Otel.Scope.add_attrs scope (fun () ->
               let code = Response.status res in
               let code = Code.code_of_status code in
               [ "http.status_code", `Int code ]);

src/trace/opentelemetry_trace.ml

@@ -147,10 +147,10 @@ module Internal = struct
     let end_time = Timestamp_ns.now_unix_ns () in
     let kind, attrs = otel_attrs_of_otrace_data (Scope.attrs scope) in
 
-    let status : Span.status =
+    let status : Span_status.t =
       match List.assoc_opt Well_known.status_error_key attrs with
-      | Some (`String message) -> { message; code = Span.Status_code_error }
-      | _ -> { message = ""; code = Span.Status_code_ok }
+      | Some (`String message) -> { message; code = Status_code_error }
+      | _ -> { message = ""; code = Status_code_ok }
     in
 
     let attrs =

tests/bin/emit1.ml

@@ -72,7 +72,8 @@ let run_job () =
         if j = 4 && !i mod 13 = 0 then failwith "oh no";
 
         (* simulate a failure *)
-        T.Trace.add_event scope (fun () -> T.Event.make "done with alloc")
+        Opentelemetry.Scope.add_event scope (fun () ->
+            T.Event.make "done with alloc")
       with Failure _ -> ()
     done
   done

tests/bin/emit1_cohttp.ml

@@ -67,7 +67,8 @@ let run_job () : unit Lwt.t =
         if j = 4 && !i mod 13 = 0 then failwith "oh no";
 
         (* simulate a failure *)
-        T.Trace.add_event scope (fun () -> T.Event.make "done with alloc");
+        Opentelemetry.Scope.add_event scope (fun () ->
+            T.Event.make "done with alloc");
         Lwt.return ()
       with Failure _ -> Lwt.return ()
     done