CP-308811: Add an option to limit the span depth in tracing

ocaml/libs/tracing/tracing.ml

@@ -222,6 +222,8 @@ module TraceContext = struct
 
   let empty = {traceparent= None; baggage= None}
 
+  let depth_key = "span.depth"
+
   let with_traceparent traceparent ctx = {ctx with traceparent}
 
   let with_baggage baggage ctx = {ctx with baggage}
@@ -230,6 +232,20 @@ module TraceContext = struct
 
   let baggage_of ctx = ctx.baggage
 
+  let baggage_depth_of ctx =
+    Option.bind (baggage_of ctx) (List.assoc_opt depth_key)
+    |> Option.value ~default:"1"
+    |> int_of_string
+
+  let update_with_baggage k v ctx =
+    let new_baggage =
+      baggage_of ctx
+      |> Option.value ~default:[]
+      |> List.remove_assoc k
+      |> List.cons (k, v)
+    in
+    with_baggage (Some new_baggage) ctx
+
   let parse input =
     let open Astring.String in
     let trim_pair (key, value) = (trim key, trim value) in
@@ -322,22 +338,36 @@ module Span = struct
 
   let start ?(attributes = Attributes.empty)
       ?(trace_context : TraceContext.t option) ~name ~parent ~span_kind () =
-    let trace_id, extra_context =
+    let trace_id, extra_context, depth =
       match parent with
       | None ->
-          (Trace_id.make (), TraceContext.empty)
+          (Trace_id.make (), TraceContext.empty, 1)
       | Some span_parent ->
-          (span_parent.context.trace_id, span_parent.context.trace_context)
+          ( span_parent.context.trace_id
+          , span_parent.context.trace_context
+          , TraceContext.baggage_depth_of span_parent.context.trace_context + 1
+          )
     in
     let span_id = Span_id.make () in
+    let extra_context_with_depth =
+      TraceContext.(
+        update_with_baggage depth_key (string_of_int depth) extra_context
+      )
+    in
     let context : SpanContext.t =
-      {trace_id; span_id; trace_context= extra_context}
+      {trace_id; span_id; trace_context= extra_context_with_depth}
     in
     let context =
-      (* If trace_context is provided to the call, override any inherited trace context. *)
-      trace_context
-      |> Option.fold ~none:context
-           ~some:(Fun.flip SpanContext.with_trace_context context)
+      (* If trace_context is provided to the call, override any inherited trace
+         context except span.depth which should still be maintained. *)
+      match trace_context with
+      | Some tc ->
+          let tc_with_depth =
+            TraceContext.(update_with_baggage depth_key (string_of_int depth) tc)
+          in
+          SpanContext.with_trace_context tc_with_depth context
+      | None ->
+          context
     in
     (* Using gettimeofday over Mtime as it is better for sharing timestamps between the systems *)
     let begin_time = Unix.gettimeofday () in
@@ -473,6 +503,11 @@ module Spans = struct
 
   let set_max_traces x = Atomic.set max_traces x
 
+  (* Default is much larger than the largest current traces, so effectively off *)
+  let max_depth = Atomic.make 100
+
+  let set_max_depth x = Atomic.set max_depth x
+
   let finished_spans = Atomic.make ([], 0)
 
   let span_hashtbl_is_empty () = TraceMap.is_empty (Atomic.get spans)
@@ -713,12 +748,18 @@ module Tracer = struct
   let get_tracer ~name:_ = TracerProvider.get_current ()
 
   let span_of_span_context context name : Span.t =
+    let tc = SpanContext.context_of_span_context context in
+    let new_depth = TraceContext.baggage_depth_of tc in
+    let new_tc =
+      TraceContext.(update_with_baggage depth_key (string_of_int new_depth) tc)
+    in
+    let context = SpanContext.with_trace_context new_tc context in
     {
       context
     ; status= {status_code= Status.Unset; _description= None}
     ; name
     ; parent= None
-    ; span_kind= SpanKind.Client (* This will be the span of the client call*)
+    ; span_kind= SpanKind.Client (* This will be the span of the client call *)
     ; begin_time= Unix.gettimeofday ()
     ; end_time= None
     ; links= []
@@ -730,10 +771,32 @@ module Tracer = struct
       ?(span_kind = SpanKind.Internal) ~name ~parent () :
       (Span.t option, exn) result =
     let open TracerProvider in
-    (* Do not start span if the TracerProvider is disabled*)
+    let parent_depth =
+      Option.fold ~none:1
+        ~some:(fun parent ->
+          parent.Span.context
+          |> SpanContext.context_of_span_context
+          |> TraceContext.baggage_depth_of
+        )
+        parent
+    in
+    (* Do not start span if the TracerProvider is disabled *)
     if not t.enabled then
+      ok_none (* Do not start span if the max depth has been reached *)
+    else if parent_depth >= Atomic.get Spans.max_depth then (
+      let parent_trace_id =
+        Option.fold ~none:"None"
+          ~some:(fun p ->
+            p.Span.context
+            |> SpanContext.span_id_of_span_context
+            |> Span_id.to_string
+          )
+          parent
+      in
+      debug "Max_span_depth limit reached, not creating span %s (parent %s)"
+        name parent_trace_id ;
       ok_none
-    else
+    ) else
       let attributes = Attributes.merge_into t.attributes attributes in
       let span =
         Span.start ~attributes ?trace_context ~name ~parent ~span_kind ()
@@ -750,16 +813,24 @@ module Tracer = struct
           |> Spans.remove_from_spans
           |> Option.map (fun existing_span ->
                  let old_context = Span.get_context existing_span in
+                 let parent_trace_context = Span.get_trace_context parent in
+                 let new_depth =
+                   TraceContext.baggage_depth_of parent_trace_context + 1
+                 in
                  let new_context : SpanContext.t =
-                   let trace_context = span.Span.context.trace_context in
+                   let trace_context =
+                     TraceContext.(
+                       update_with_baggage depth_key (string_of_int new_depth)
+                         span.Span.context.trace_context
+                     )
+                   in
                    SpanContext.context
                      (SpanContext.trace_id_of_span_context parent.context)
                      old_context.span_id
                    |> SpanContext.with_trace_context trace_context
                  in
                  let updated_span = {existing_span with parent= Some parent} in
                  let updated_span = {updated_span with context= new_context} in
-
                  let () = Spans.add_to_spans ~span:updated_span in
                  updated_span
              )
@@ -926,7 +997,15 @@ module Propagator = struct
             let trace_context' =
               TraceContext.with_traceparent (Some traceparent) trace_context
             in
-            let carrier' = P.inject_into trace_context' carrier in
+            let new_depth =
+              TraceContext.baggage_depth_of trace_context' + 1 |> string_of_int
+            in
+            let trace_context'' =
+              TraceContext.(
+                update_with_baggage depth_key new_depth trace_context'
+              )
+            in
+            let carrier' = P.inject_into trace_context'' carrier in
             f carrier'
         | _ ->
             f carrier

ocaml/libs/tracing/tracing.mli

@@ -165,6 +165,8 @@ module Spans : sig
 
   val set_max_traces : int -> unit
 
+  val set_max_depth : int -> unit
+
   val span_count : unit -> int
 
   val since : unit -> Span.t list * int

ocaml/libs/tracing/tracing_export.ml

@@ -24,6 +24,10 @@ let export_interval = ref 30.
 
 let set_export_interval t = export_interval := t
 
+let export_chunk_size = Atomic.make 10000
+
+let set_export_chunk_size x = Atomic.set export_chunk_size x
+
 let host_id = ref "localhost"
 
 let set_host_id id = host_id := id
@@ -289,17 +293,40 @@ module Destination = struct
     with exn ->
       debug "Tracing: unable to export span : %s" (Printexc.to_string exn)
 
+  let rec span_info_chunks span_info batch_size =
+    let rec list_to_chunks_inner l n curr chunks =
+      if n = 0 then
+        if l <> [] then
+          list_to_chunks_inner l batch_size [] ((curr, batch_size) :: chunks)
+        else
+          (curr, batch_size) :: chunks
+      else
+        match l with
+        | [] ->
+            (curr, List.length curr) :: chunks
+        | h :: t ->
+            list_to_chunks_inner t (n - 1) (h :: curr) chunks
+    in
+    list_to_chunks_inner (fst span_info) batch_size [] []
+
   let flush_spans () =
     let ((_span_list, span_count) as span_info) = Spans.since () in
     let attributes = [("export.traces.count", string_of_int span_count)] in
     let@ parent =
       with_tracing ~span_kind:Server ~trace_context:TraceContext.empty
         ~parent:None ~attributes ~name:"Tracing.flush_spans"
     in
-    TracerProvider.get_tracer_providers ()
-    |> List.filter TracerProvider.get_enabled
-    |> List.concat_map TracerProvider.get_endpoints
-    |> List.iter (export_to_endpoint parent span_info)
+    let endpoints =
+      TracerProvider.get_tracer_providers ()
+      |> List.filter TracerProvider.get_enabled
+      |> List.concat_map TracerProvider.get_endpoints
+    in
+    let span_info_chunks =
+      span_info_chunks span_info (Atomic.get export_chunk_size)
+    in
+    List.iter
+      (fun s_i -> List.iter (export_to_endpoint parent s_i) endpoints)
+      span_info_chunks
 
   let delay = Delay.make ()
 

ocaml/libs/tracing/tracing_export.mli

@@ -23,6 +23,13 @@ val set_export_interval : float -> unit
     Default is every [30.] seconds.
   *)
 
+val set_export_chunk_size : int -> unit
+(** [set_export_chunk_size size] sets the maximum number of finished spans that
+    can be exported in one chunk to [size].
+
+    Default is 10000 spans.
+  *)
+
 val set_host_id : string -> unit
 (** [set_host_id id] sets the id of the host to [id]. 
 

ocaml/tests/test_cluster.ml

@@ -34,9 +34,11 @@ let test_clusterd_rpc ~__context call =
       | "Observer.init"
       | "Observer.set_trace_log_dir"
       | "Observer.set_export_interval"
+      | "Observer.set_export_chunk_size"
       | "Observer.set_host_id"
       | "Observer.set_max_traces"
       | "Observer.set_max_spans"
+      | "Observer.set_max_depth"
       | "Observer.set_max_file_size"
       | "Observer.set_compress_tracing_files" )
     , _ ) ->

ocaml/tests/test_observer.ml

@@ -305,6 +305,7 @@ let verify_json_fields_and_values ~json =
             ; ("xs.host.uuid", `String _)
             ; ("xs.host.name", `String _)
             ; ("service.name", `String _)
+            ; ("span.depth", `String _)
             ]
         )
       ; ("annotations", `List _)

ocaml/xapi-idl/cluster/cli-help.t

@@ -21,10 +21,14 @@
 
          Observer.set_endpoints [OPTION]… dbg uuid endpoints
 
+         Observer.set_export_chunk_size [OPTION]… dbg int
+
          Observer.set_export_interval [OPTION]… dbg float
 
          Observer.set_host_id [OPTION]… dbg string
 
+         Observer.set_max_depth [OPTION]… dbg int
+
          Observer.set_max_file_size [OPTION]… dbg int
 
          Observer.set_max_spans [OPTION]… dbg int

ocaml/xapi-idl/lib/observer_helpers.ml

@@ -138,6 +138,10 @@ module ObserverAPI (R : RPC) = struct
     declare "Observer.set_export_interval" []
       (dbg_p @-> float_p @-> returning unit_p err)
 
+  let set_export_chunk_size =
+    declare "Observer.set_export_chunk_size" []
+      (dbg_p @-> int_p @-> returning unit_p err)
+
   let set_max_spans =
     declare "Observer.set_max_spans" []
       (dbg_p @-> int_p @-> returning unit_p err)
@@ -146,6 +150,10 @@ module ObserverAPI (R : RPC) = struct
     declare "Observer.set_max_traces" []
       (dbg_p @-> int_p @-> returning unit_p err)
 
+  let set_max_depth =
+    declare "Observer.set_max_depth" []
+      (dbg_p @-> int_p @-> returning unit_p err)
+
   let set_max_file_size =
     declare "Observer.set_max_file_size" []
       (dbg_p @-> int_p @-> returning unit_p err)
@@ -193,10 +201,14 @@ module type Server_impl = sig
 
   val set_export_interval : context -> dbg:debug_info -> interval:float -> unit
 
+  val set_export_chunk_size : context -> dbg:debug_info -> size:int -> unit
+
   val set_max_spans : context -> dbg:debug_info -> spans:int -> unit
 
   val set_max_traces : context -> dbg:debug_info -> traces:int -> unit
 
+  val set_max_depth : context -> dbg:debug_info -> depth:int -> unit
+
   val set_max_file_size : context -> dbg:debug_info -> file_size:int -> unit
 
   val set_host_id : context -> dbg:debug_info -> host_id:string -> unit
@@ -227,8 +239,12 @@ module Server (Impl : Server_impl) () = struct
     S.set_export_interval (fun dbg interval ->
         Impl.set_export_interval () ~dbg ~interval
     ) ;
+    S.set_export_chunk_size (fun dbg size ->
+        Impl.set_export_chunk_size () ~dbg ~size
+    ) ;
     S.set_max_spans (fun dbg spans -> Impl.set_max_spans () ~dbg ~spans) ;
     S.set_max_traces (fun dbg traces -> Impl.set_max_traces () ~dbg ~traces) ;
+    S.set_max_depth (fun dbg depth -> Impl.set_max_depth () ~dbg ~depth) ;
     S.set_max_file_size (fun dbg file_size ->
         Impl.set_max_file_size () ~dbg ~file_size
     ) ;