Trace: Fixes thread safety issue in Trace class causing high CPU usage and InvalidOperationException

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosTraceDiagnostics.cs

@@ -40,6 +40,11 @@ public CosmosTraceDiagnostics(ITrace trace)
 
         public override string ToString()
         {
+            if (this.Value is Tracing.Trace rootConcreteTrace)
+            {
+                rootConcreteTrace.SetWalkingStateRecursively();
+            }
+
             return this.ToJsonString();
         }
 
@@ -50,16 +55,31 @@ public override TimeSpan GetClientElapsedTime()
 
         public override IReadOnlyList<(string regionName, Uri uri)> GetContactedRegions()
         {
+            if (this.Value is Tracing.Trace rootConcreteTrace)
+            {
+                rootConcreteTrace.SetWalkingStateRecursively();
+            }
+
             return this.Value?.Summary?.RegionsContacted;
         }
 
         public override ServerSideCumulativeMetrics GetQueryMetrics()
         {
+            if (this.Value is Tracing.Trace rootConcreteTrace)
+            {
+                rootConcreteTrace.SetWalkingStateRecursively();
+            }
+
             return this.accumulatedMetrics.Value;
         }
 
         internal bool IsGoneExceptionHit()
         {
+            if (this.Value is Tracing.Trace rootConcreteTrace)
+            {
+                rootConcreteTrace.SetWalkingStateRecursively();
+            }
+
             return this.WalkTraceTreeForGoneException(this.Value);
         }
 

Microsoft.Azure.Cosmos/src/Pagination/ParallelPrefetch.Testing.cs

@@ -57,6 +57,8 @@ public int AwaitedTasks
 
             IReadOnlyDictionary<string, object> ITrace.Data => this.innerTrace.Data;
 
+            bool ITrace.IsBeingWalked => this.innerTrace.IsBeingWalked;
+
             public ParallelPrefetchTestConfig(
                 ArrayPool<IPrefetcher> prefetcherPool,
                 ArrayPool<Task> taskPool,
@@ -116,6 +118,11 @@ void IDisposable.Dispose()
             {
                 this.innerTrace.Dispose();
             }
+
+            bool ITrace.TryGetDatum(string key, out object datum)
+            {
+                return this.innerTrace.TryGetDatum(key, out datum);
+            }
         }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsTraceExtractor.cs

@@ -17,6 +17,9 @@ public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideM
                 return;
             }
 
+            // Assert that walking state is set
+            System.Diagnostics.Debug.Assert(currentTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
+
             foreach (object datum in currentTrace.Data.Values)
             {
                 if (datum is QueryMetricsTraceDatum queryMetricsTraceDatum)
@@ -41,6 +44,9 @@ private static void WalkTraceTreeForPartitionInfo(ITrace currentTrace, ServerSid
                 return;
             }
 
+            // Assert that walking state is set
+            System.Diagnostics.Debug.Assert(currentTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
+
             foreach (Object datum in currentTrace.Data.Values)
             {
                 if (datum is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs

@@ -180,7 +180,7 @@ public override async Task<ResponseMessage> ReadNextAsync(ITrace trace, Cancella
 
             // If Correlated Id already exists and is different, add a new one in comma separated list
             // Scenario: A new iterator is created with same ContinuationToken and Trace 
-            if (trace.Data.TryGetValue(QueryIterator.CorrelatedActivityIdKeyName, out object correlatedActivityIds))
+            if (trace.TryGetDatum(QueryIterator.CorrelatedActivityIdKeyName, out object correlatedActivityIds))
             {
                 List<string> correlatedIdList = correlatedActivityIds.ToString().Split(',').ToList();
                 if (!correlatedIdList.Contains(this.correlatedActivityId.ToString()))

Microsoft.Azure.Cosmos/src/Tracing/ITrace.cs

@@ -68,6 +68,11 @@ interface ITrace : IDisposable
         /// </summary>
         IReadOnlyDictionary<string, object> Data { get; }
 
+        /// <summary>
+        /// Gets a value indicating whether this trace is currently being walked.
+        /// </summary>
+        bool IsBeingWalked { get; }
+
         /// <summary>
         /// Starts a Trace and adds it as a child to this instance.
         /// </summary>
@@ -115,5 +120,12 @@ ITrace StartChild(
         /// <param name="trace">Existing trace.</param>
         void AddChild(ITrace trace);
 
+        /// <summary>
+        /// Tries to get a specific datum - it is safe to call this even before IsBeingWalked is set.
+        /// </summary>
+        /// <param name="key">The key to identify the datum.</param>
+        /// <param name="datum">The datum itself.</param>
+        /// <returns>A flag indicating whether the datum with this key exists.</returns>
+        bool TryGetDatum(string key, out object datum);
     }
 }

Microsoft.Azure.Cosmos/src/Tracing/NoOpTrace.cs

@@ -40,6 +40,8 @@ private NoOpTrace()
 
         public IReadOnlyDictionary<string, object> Data => NoOpData;
 
+        public bool IsBeingWalked => true; // this needs to return true to allow materialization of NoOpTrace
+
         public void Dispose()
         {
             // NoOp
@@ -86,5 +88,11 @@ public void UpdateRegionContacted(TraceDatum traceDatum)
         {
             // NoOp
         }
+
+        bool ITrace.TryGetDatum(string key, out object datum)
+        {
+            datum = null;
+            return false;
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Tracing/Trace.cs

@@ -7,17 +7,15 @@ namespace Microsoft.Azure.Cosmos.Tracing
     using System;
     using System.Collections.Generic;
     using System.Diagnostics;
-    using System.Linq;
-    using System.Runtime.CompilerServices;
-    using Microsoft.Azure.Cosmos.Tracing.TraceData;
     using Microsoft.Azure.Documents;
 
     internal sealed class Trace : ITrace
     {
         private static readonly IReadOnlyDictionary<string, object> EmptyDictionary = new Dictionary<string, object>();
-        private readonly List<ITrace> children;
-        private readonly Lazy<Dictionary<string, object>> data;
+        private volatile List<ITrace> children;
+        private volatile Dictionary<string, object> data;
         private ValueStopwatch stopwatch;
+        private volatile bool isBeingWalked;
 
         private Trace(
             string name,
@@ -34,7 +32,7 @@ private Trace(
             this.Component = component;
             this.Parent = parent;
             this.children = new List<ITrace>();
-            this.data = new Lazy<Dictionary<string, object>>();
+            this.data = null;
             this.Summary = summary ?? throw new ArgumentNullException(nameof(summary));
         }
 
@@ -54,9 +52,37 @@ private Trace(
 
         public ITrace Parent { get; }
 
-        public IReadOnlyList<ITrace> Children => this.children;
+        // NOTE: no lock necessary here only because this.children is volatile
+        // and every reference to it is immutable when isBeingWalked == true
+        // and isBeingWalked is guaranteed to be set to true before this
+        // Property is called
+        public IReadOnlyList<ITrace> Children
+        {
+            get
+            {
+                // Assert that walking state is set
+                Debug.Assert(this.isBeingWalked, "SetWalkingStateRecursively should be set to true");
+
+                return this.children;
+            }
+        }
 
-        public IReadOnlyDictionary<string, object> Data => this.data.IsValueCreated ? this.data.Value : Trace.EmptyDictionary;
+        // NOTE: no lock necessary here only because this.data is volatile
+        // and every reference to it is immutable when isBeingWalked == true
+        // and isBeingWalked is guaranteed to be set to true before this
+        // Property is called
+        public IReadOnlyDictionary<string, object> Data
+        {
+            get
+            {
+                // Assert that walking state is set
+                Debug.Assert(this.isBeingWalked, "SetWalkingStateRecursively should be set to true");
+
+                return this.data ?? Trace.EmptyDictionary;
+            }
+        }
+
+        public bool IsBeingWalked => this.isBeingWalked;
 
         public void Dispose()
         {
@@ -90,14 +116,30 @@ public ITrace StartChild(
                 summary: this.Summary);
 
             this.AddChild(child);
+
             return child;
         }
 
         public void AddChild(ITrace child)
         {
-            lock (this.children)
+            lock (this.Name)
             {
-                this.children.Add(child);
+                if (!this.isBeingWalked)
+                {
+                    this.children.Add(child);
+
+                    return;
+                }
+
+                if (child is Trace traceChild)
+                {
+                    traceChild.SetWalkingStateRecursively();
+                }
+
+                List<ITrace> writableSnapshot = new List<ITrace>(this.children.Count + 1);
+                writableSnapshot.AddRange(this.children);
+                writableSnapshot.Add(child);
+                this.children = writableSnapshot;
             }
         }
 
@@ -124,18 +166,89 @@ public static Trace GetRootTrace(
 
         public void AddDatum(string key, TraceDatum traceDatum)
         {
-            this.data.Value.Add(key, traceDatum);
             this.Summary.UpdateRegionContacted(traceDatum);
+            this.AddDatum(key, traceDatum as Object);
         }
 
         public void AddDatum(string key, object value)
         {
-            this.data.Value.Add(key, value);
+            lock (this.Name)
+            {
+                this.data ??= new Dictionary<string, object>();
+
+                if (!this.isBeingWalked)
+                {
+                    // If materialization has not started yet no cloning is needed
+                    this.data.Add(key, value);
+                    return;
+                }
+
+                this.data = new Dictionary<string, object>(this.data)
+                {
+                    { key, value }
+                };
+            }
         }
 
         public void AddOrUpdateDatum(string key, object value)
         {
-            this.data.Value[key] = value;
+            lock (this.Name)
+            {
+                this.data ??= new Dictionary<string, object>();
+
+                if (!this.isBeingWalked)
+                {
+                    // If materialization has not started yet no cloning is needed
+                    this.data[key] = value;
+                    return; // Ignore modifications while being walked
+                }
+
+                this.data = new Dictionary<string, object>(this.data)
+                {
+                    [key] = value
+                };
+            }
+        }
+
+        internal void SetWalkingStateRecursively()
+        {
+            if (this.isBeingWalked)
+            {
+                return; // Already set, return early
+            }
+
+            lock (this.Name)
+            {
+                if (this.isBeingWalked)
+                {
+                    return; // Already set, return early
+                }
+
+                foreach (ITrace child in this.children)
+                {
+                    if (child is Trace concreteChild)
+                    {
+                        concreteChild.SetWalkingStateRecursively();
+                    }
+                }
+
+                // Set the walking state for this trace after processing children
+                this.isBeingWalked = true;
+            }
+        }
+
+        bool ITrace.TryGetDatum(string key, out object datum)
+        {
+            lock (this.Name)
+            {
+                if (this.data == null)
+                {
+                    datum = null;
+                    return false;
+                }
+
+                return this.data.TryGetValue(key, out datum);
+            }
         }
     }
 }

Microsoft.Azure.Cosmos/src/Tracing/TraceData/SummaryDiagnostics.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Tracing.TraceData
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Globalization;
     using System.Linq;
     using Microsoft.Azure.Cosmos.Json;
@@ -22,7 +23,11 @@ public SummaryDiagnostics(ITrace trace)
                 = new Lazy<Dictionary<(int, int), int>>(() => new Dictionary<(int, int), int>());
             this.AllRegionsContacted 
                 = new Lazy<HashSet<Uri>>(() => new HashSet<Uri>());
-
+
+            if (trace is Tracing.Trace rootConcreteTrace)
+            {
+                rootConcreteTrace.SetWalkingStateRecursively();
+            }
             this.CollectSummaryFromTraceTree(trace);
         }
 
@@ -38,6 +43,9 @@ public SummaryDiagnostics(ITrace trace)
 
         private void CollectSummaryFromTraceTree(ITrace currentTrace)
         {
+            // Assert that walking state is set
+            Debug.Assert(currentTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
+
             foreach (object datums in currentTrace.Data.Values)
             {
                 if (datums is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Tracing/EndToEndTraceWriterBaselineTests.cs

@@ -1865,6 +1865,8 @@ public TraceForBaselineTesting(
 
             public IReadOnlyDictionary<string, object> Data => this.data;
 
+            public bool IsBeingWalked => true; // needs to return true to allow materialization
+
             public IReadOnlyList<(string, Uri)> RegionsContacted => new List<(string, Uri)>();
 
             public void AddDatum(string key, TraceDatum traceDatum)
@@ -1924,6 +1926,11 @@ public void AddOrUpdateDatum(string key, object value)
 
                 this.data[key] = "Redacted To Not Change The Baselines From Run To Run";
             }
+
+            bool ITrace.TryGetDatum(string key, out object datum)
+            {
+                return this.data.TryGetValue(key, out datum);
+            }
         }
 
         private sealed class RequestHandlerSleepHelper : RequestHandler