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

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

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Tracing
 {
     using System;
+    using System.Collections.Concurrent;
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.Linq;
@@ -16,7 +17,7 @@ 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 readonly ConcurrentDictionary<string, object> data;
         private ValueStopwatch stopwatch;
 
         private Trace(
@@ -34,7 +35,7 @@ private Trace(
             this.Component = component;
             this.Parent = parent;
             this.children = new List<ITrace>();
-            this.data = new Lazy<Dictionary<string, object>>();
+            this.data = new ConcurrentDictionary<string, object>();
             this.Summary = summary ?? throw new ArgumentNullException(nameof(summary));
         }
 
@@ -56,7 +57,7 @@ private Trace(
 
         public IReadOnlyList<ITrace> Children => this.children;
 
-        public IReadOnlyDictionary<string, object> Data => this.data.IsValueCreated ? this.data.Value : Trace.EmptyDictionary;
+        public IReadOnlyDictionary<string, object> Data => this.data.Count > 0 ? this.data : Trace.EmptyDictionary;
 
         public void Dispose()
         {
@@ -122,20 +123,38 @@ public static Trace GetRootTrace(
                 summary: new TraceSummary());
         }
 
+        /// <summary>
+        /// Adds a datum to the this trace instance.
+        /// This method is thread-safe.
+        /// </summary>
+        /// <param name="key">The key to associate the datum.</param>
+        /// <param name="traceDatum">The datum itself.</param>
         public void AddDatum(string key, TraceDatum traceDatum)
         {
-            this.data.Value.Add(key, traceDatum);
+            this.data.TryAdd(key, traceDatum);
             this.Summary.UpdateRegionContacted(traceDatum);
         }
 
+        /// <summary>
+        /// Adds a datum to the this trace instance.
+        /// This method is thread-safe.
+        /// </summary>
+        /// <param name="key">The key to associate the datum.</param>
+        /// <param name="value">The datum itself.</param>
         public void AddDatum(string key, object value)
         {
-            this.data.Value.Add(key, value);
+            this.data.TryAdd(key, value);
         }
 
+        /// <summary>
+        /// Updates the given datum in this trace instance if exists, otherwise Add
+        /// This method is thread-safe.
+        /// </summary>
+        /// <param name="key">The key to associate the datum.</param>
+        /// <param name="value">The datum itself.</param>
         public void AddOrUpdateDatum(string key, object value)
         {
-            this.data.Value[key] = value;
+            this.data[key] = value;
         }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceTests.cs

@@ -4,6 +4,7 @@
     using System.Collections.Generic;
     using System.Linq;
     using System.Reflection;
+    using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Diagnostics;
     using Microsoft.Azure.Cosmos.Json;
     using Microsoft.Azure.Cosmos.Tracing;
@@ -134,5 +135,50 @@ public void ValidateStoreResultSerialization()
 
             Assert.AreEqual(0, storeResultProperties.Count, $"Json is missing properties: {string.Join(';', storeResultProperties)}");
         }
+        [TestMethod]
+        public void TestAddOrUpdateDatumThreadSafety()
+        {
+            Trace trace = Trace.GetRootTrace("ThreadSafetyTest");
+
+            // Create multiple threads to access the dictionary concurrently
+            const int numThreads = 10;
+            const int operationsPerThread = 100;
+
+            // Use a list to keep track of the tasks
+            List<Task> tasks = new List<Task>();
+
+            for (int i = 0; i < numThreads; i++)
+            {
+                int threadIndex = i;
+                tasks.Add(Task.Run(() =>
+                {
+                    for (int j = 0; j < operationsPerThread; j++)
+                    {
+                        string key = $"key_{threadIndex}_{j}";
+                        object value = j;
+
+                        // Perform operations that would previously cause thread safety issues
+                        trace.AddOrUpdateDatum(key, value);
+
+                        // Also test AddDatum
+                        try
+                        {
+                            trace.AddDatum($"add_{threadIndex}_{j}", value);
+                        }
+                        catch (ArgumentException)
+                        {
+                            // Ignore key already exists exceptions which may occur
+                            // when threads try to add the same key
+                        }
+                    }
+                }));
+            }
+
+            // Wait for all tasks to complete
+            Task.WaitAll(tasks.ToArray());
+
+            // Verify the data dictionary has entries
+            Assert.IsTrue(trace.Data.Count > 0);
+        }
     }
 }