From dbf459996ab424ce0116f5fc9eda5178d302ba90 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 3 Apr 2026 21:19:27 -0400
Subject: [PATCH] add trace state to span and span_ctx

---
 src/core/span.ml          | 30 +++++++++++++++++++++++++++++-
 src/core/span.mli         | 11 +++++++++++
 src/core/span_ctx.ml      | 18 +++++++++++++-----
 src/core/span_ctx.mli     |  3 +++
 src/core/trace_context.ml | 11 +++++++++++
 5 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/src/core/span.ml b/src/core/span.ml
index 21b8de84..d045a2a0 100644
--- a/src/core/span.ml
+++ b/src/core/span.ml
@@ -89,8 +89,36 @@ let to_span_link (self : t) : Span_link.t =
          None)
     ~trace_id:self.trace_id ~span_id:self.span_id ()
 
+let[@inline] trace_state (self : t) : Trace_state.t =
+  if span_has_trace_state self then
+    Trace_state.of_w3c_string self.trace_state |> Result.value ~default:[]
+  else
+    []
+
 let[@inline] to_span_ctx (self : t) : Span_ctx.t =
-  Span_ctx.make ~trace_id:(trace_id self) ~parent_id:(id self) ()
+  Span_ctx.make ~trace_id:(trace_id self) ~parent_id:(id self)
+    ~trace_state:(trace_state self) ()
+
+let add_trace_state_attr (self : t) (key : string) (value : string) : unit =
+  if not (Trace_state.is_valid_key key) then
+    invalid_arg (Printf.sprintf "Span.add_trace_state_attr: invalid key %S" key);
+  if not (Trace_state.is_valid_value value) then
+    invalid_arg
+      (Printf.sprintf "Span.add_trace_state_attr: invalid value %S" value);
+  let cur =
+    if span_has_trace_state self then
+      self.trace_state
+    else
+      ""
+  in
+  let entry = key ^ "=" ^ value in
+  let new_ts =
+    if cur = "" then
+      entry
+    else
+      entry ^ "," ^ cur
+  in
+  span_set_trace_state self new_ts
 
 (* Note: a span must not be concurrently modified from multiple
    threads or domains. *)
diff --git a/src/core/span.mli b/src/core/span.mli
index 15f34e77..708b0fea 100644
--- a/src/core/span.mli
+++ b/src/core/span.mli
@@ -79,6 +79,17 @@ val status : t -> Span_status.t option
 
 val kind : t -> Span_kind.t option
 
+val trace_state : t -> Trace_state.t
+(** Returns the decoded trace state, or [[]] if absent or invalid. *)
+
+val add_trace_state_attr : t -> string -> string -> unit
+(** [add_trace_state_attr span key value] prepends/replaces the [key=value]
+    entry in the span's trace state.
+
+    This is not the most efficient, as it'll copy the whole string to append to
+    it.
+    @raise Invalid_argument if [key] or [value] are invalid per W3C rules. *)
+
 val to_span_link : t -> Span_link.t
 (** Turn the scope into a span link *)
 
diff --git a/src/core/span_ctx.ml b/src/core/span_ctx.ml
index 3f2c73e1..022efddb 100644
--- a/src/core/span_ctx.ml
+++ b/src/core/span_ctx.ml
@@ -2,8 +2,6 @@ open Common_
 
 (* see: https://opentelemetry.io/docs/specs/otel/trace/api/#spancontext *)
 
-(* TODO: trace state *)
-
 external int_of_bool : bool -> int = "%identity"
 
 module Flags = struct
@@ -16,17 +14,25 @@ type t = {
   trace_id: Trace_id.t;
   parent_id: Span_id.t;
   flags: int;
+  trace_state: Trace_state.t;  (** W3C trace state for distributed tracing *)
 }
 
-let dummy = { trace_id = Trace_id.dummy; parent_id = Span_id.dummy; flags = 0 }
+let dummy =
+  {
+    trace_id = Trace_id.dummy;
+    parent_id = Span_id.dummy;
+    flags = 0;
+    trace_state = [];
+  }
 
-let make ?(remote = false) ?(sampled = false) ~trace_id ~parent_id () : t =
+let make ?(remote = false) ?(sampled = false) ?(trace_state = Trace_state.empty)
+    ~trace_id ~parent_id () : t =
   let flags =
     0
     lor (int_of_bool remote lsl Flags.remote)
     lor (int_of_bool sampled lsl Flags.sampled)
   in
-  { trace_id; parent_id; flags }
+  { trace_id; parent_id; flags; trace_state }
 
 let[@inline] is_valid self =
   Trace_id.is_valid self.trace_id && Span_id.is_valid self.parent_id
@@ -39,6 +45,8 @@ let[@inline] trace_id self = self.trace_id
 
 let[@inline] parent_id self = self.parent_id
 
+let[@inline] trace_state self = self.trace_state
+
 let to_w3c_trace_context (self : t) : bytes =
   let bs = Bytes.create 55 in
   Bytes.set bs 0 '0';
diff --git a/src/core/span_ctx.mli b/src/core/span_ctx.mli
index 2970daab..355dd43d 100644
--- a/src/core/span_ctx.mli
+++ b/src/core/span_ctx.mli
@@ -9,6 +9,7 @@ type t
 val make :
   ?remote:bool ->
   ?sampled:bool ->
+  ?trace_state:Trace_state.t ->
   trace_id:Trace_id.t ->
   parent_id:Span_id.t ->
   unit ->
@@ -29,6 +30,8 @@ val parent_id : t -> Span_id.t
 
 val sampled : t -> bool
 
+val trace_state : t -> Trace_state.t
+
 val to_w3c_trace_context : t -> bytes
 
 val of_w3c_trace_context : bytes -> (t, string) result
diff --git a/src/core/trace_context.ml b/src/core/trace_context.ml
index 9c8b141d..ce831914 100644
--- a/src/core/trace_context.ml
+++ b/src/core/trace_context.ml
@@ -32,3 +32,14 @@ module Traceparent = struct
     let span_ctx = Span_ctx.make ?sampled ~trace_id ~parent_id () in
     Bytes.unsafe_to_string @@ Span_ctx.to_w3c_trace_context span_ctx
 end
+
+(** The tracestate header https://www.w3.org/TR/trace-context/#tracestate-header
+*)
+module Tracestate = struct
+  let name = "tracestate"
+
+  let of_w3c_string (s : string) : (Trace_state.t, string) result =
+    Trace_state.of_w3c_string s
+
+  let to_w3c_string (ts : Trace_state.t) : string = Trace_state.to_w3c_string ts
+end