Preserve durable trace context in SQL history

src/DurableTask.SqlServer/SqlUtils.cs

@@ -24,6 +24,10 @@ static class SqlUtils
     {
         static readonly Random random = new Random();
         static readonly char[] TraceContextSeparators = new char[] { '\n' };
+        const string TraceContextTraceStatePrefix = "@tracestate=";
+        const string TraceContextIdPrefix = "@id=";
+        const string TraceContextSpanIdPrefix = "@spanid=";
+        const string TraceContextClientSpanIdPrefix = "@clientspanid=";
         const int MaxTagsPayloadSize = 8000;
 
         public static string? GetStringOrNull(this DbDataReader reader, int columnIndex)
@@ -135,6 +139,7 @@ public static HistoryEvent GetHistoryEvent(this DbDataReader reader, bool isOrch
                     {
                         Input = GetPayloadText(reader),
                         InstanceId = "", // Placeholder - shouldn't technically be needed (adding it requires a SQL schema change)
+                        ClientSpanId = GetSubOrchestrationClientSpanId(reader),
                         Name = GetName(reader),
                         Version = null,
                     };
@@ -441,6 +446,19 @@ static DateTime GetUtcDateTime(DbDataReader reader, int ordinal)
 
         internal static SqlString GetTraceContext(HistoryEvent e)
         {
+            if (e is SubOrchestrationInstanceCreatedEvent subOrchestrationEvent)
+            {
+                if (string.IsNullOrEmpty(subOrchestrationEvent.ClientSpanId))
+                {
+                    return SqlString.Null;
+                }
+
+                // Reserve line 1 for traceparent (empty here) so all TraceContext payloads share
+                // a single parsing contract: parts[0] is always traceparent (possibly empty),
+                // and subsequent lines carry typed @key=value fields like @clientspanid=.
+                return new SqlString($"\n{TraceContextClientSpanIdPrefix}{subOrchestrationEvent.ClientSpanId}");
+            }
+
             if (e is not ISupportsDurableTraceContext eventWithTraceContext ||
                 eventWithTraceContext.ParentTraceContext == null)
             {
@@ -452,15 +470,46 @@ internal static SqlString GetTraceContext(HistoryEvent e)
             // We prefer a simple format instead of JSON because external callers may interact with this
             // data and we don't want to expose them to some internal JSON serialization format.
             var sb = new StringBuilder(traceContext.TraceParent, capacity: 800);
-            if (!string.IsNullOrEmpty(traceContext.TraceState))
+            if (!string.IsNullOrEmpty(traceContext.Id) || !string.IsNullOrEmpty(traceContext.SpanId))
+            {
+                if (!string.IsNullOrEmpty(traceContext.TraceState))
+                {
+                    sb.Append('\n').Append(TraceContextTraceStatePrefix).Append(traceContext.TraceState);
+                }
+
+                if (!string.IsNullOrEmpty(traceContext.Id))
+                {
+                    sb.Append('\n').Append(TraceContextIdPrefix).Append(traceContext.Id);
+                }
+
+                if (!string.IsNullOrEmpty(traceContext.SpanId))
+                {
+                    sb.Append('\n').Append(TraceContextSpanIdPrefix).Append(traceContext.SpanId);
+                }
+            }
+            else if (!string.IsNullOrEmpty(traceContext.TraceState))
             {
                 sb.Append('\n').Append(traceContext.TraceState);
             }
 
             return sb.ToString();
         }
 
-        static DistributedTraceContext? GetTraceContext(DbDataReader reader)
+        /// <summary>
+        /// Parsed result of a TraceContext column payload. Centralizes the on-the-wire format
+        /// (line 1 = traceparent, subsequent lines = typed @key=value fields) so all callers
+        /// share the same parsing contract.
+        /// </summary>
+        struct ParsedTraceContext
+        {
+            public string? TraceParent { get; set; }
+            public string? TraceState { get; set; }
+            public string? Id { get; set; }
+            public string? SpanId { get; set; }
+            public string? ClientSpanId { get; set; }
+        }
+
+        static ParsedTraceContext? ParseTraceContext(DbDataReader reader)
         {
             int ordinal = reader.GetOrdinal("TraceContext");
             if (reader.IsDBNull(ordinal))
@@ -474,18 +523,95 @@ internal static SqlString GetTraceContext(HistoryEvent e)
                 return null;
             }
 
-            string[] parts = text.Split(TraceContextSeparators, count: 2, StringSplitOptions.RemoveEmptyEntries);
-            var traceContext = new DistributedTraceContext(traceParent: parts[0]);
+            string[] parts = text.Split(TraceContextSeparators, StringSplitOptions.None);
+
+            string? traceParent = null;
+            string? traceState = null;
+            string? id = null;
+            string? spanId = null;
+            string? clientSpanId = null;
 
-            if (parts.Length > 1)
+            // Line 1 is reserved for traceparent. Older histories may have written a typed
+            // "@key=value" prefix on line 1 (legacy sub-orchestration payload). Detect that
+            // case by checking for the @ sentinel; otherwise treat parts[0] as traceparent.
+            int startIndex;
+            if (!string.IsNullOrEmpty(parts[0]) && parts[0][0] != '@')
             {
-                traceContext.TraceState = parts[1];
+                traceParent = parts[0];
+                startIndex = 1;
             }
+            else
+            {
+                startIndex = 0;
+            }
+
+            for (int i = startIndex; i < parts.Length; i++)
+            {
+                string part = parts[i];
+                if (string.IsNullOrEmpty(part))
+                {
+                    continue;
+                }
 
-            traceContext.ActivityStartTime = GetTimestamp(reader);
+                if (part.StartsWith(TraceContextTraceStatePrefix, StringComparison.Ordinal))
+                {
+                    traceState = part.Substring(TraceContextTraceStatePrefix.Length);
+                }
+                else if (part.StartsWith(TraceContextIdPrefix, StringComparison.Ordinal))
+                {
+                    id = part.Substring(TraceContextIdPrefix.Length);
+                }
+                else if (part.StartsWith(TraceContextSpanIdPrefix, StringComparison.Ordinal))
+                {
+                    spanId = part.Substring(TraceContextSpanIdPrefix.Length);
+                }
+                else if (part.StartsWith(TraceContextClientSpanIdPrefix, StringComparison.Ordinal))
+                {
+                    clientSpanId = part.Substring(TraceContextClientSpanIdPrefix.Length);
+                }
+                else if (traceState == null && i > 0)
+                {
+                    // Preserve the legacy format, where the optional second line stored only tracestate.
+                    traceState = part;
+                }
+            }
+
+            return new ParsedTraceContext
+            {
+                TraceParent = traceParent,
+                TraceState = traceState,
+                Id = id,
+                SpanId = spanId,
+                ClientSpanId = clientSpanId,
+            };
+        }
+
+        static DistributedTraceContext? GetTraceContext(DbDataReader reader)
+        {
+            ParsedTraceContext? parsed = ParseTraceContext(reader);
+            if (parsed == null || string.IsNullOrEmpty(parsed.Value.TraceParent))
+            {
+                // No traceparent means this row carries only sub-orchestration-specific data
+                // (e.g. @clientspanid=...) which is not a DistributedTraceContext.
+                return null;
+            }
+
+            ParsedTraceContext value = parsed.Value;
+            var traceContext = new DistributedTraceContext(traceParent: value.TraceParent!)
+            {
+                TraceState = value.TraceState,
+                Id = value.Id,
+                SpanId = value.SpanId,
+                ActivityStartTime = GetTimestamp(reader),
+            };
             return traceContext;
         }
 
+        static string? GetSubOrchestrationClientSpanId(DbDataReader reader)
+        {
+            return ParseTraceContext(reader)?.ClientSpanId;
+        }
+
         internal static IDictionary<string, string>? GetTags(DbDataReader reader)
         {
             int ordinal = reader.GetOrdinal("Tags");