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

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

@@ -90,20 +90,23 @@ ITrace StartChild(
 
         /// <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>
         void AddDatum(string key, TraceDatum 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>
         void AddDatum(string key, object 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>

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

@@ -124,18 +124,27 @@ public static Trace GetRootTrace(
 
         public void AddDatum(string key, TraceDatum traceDatum)
         {
-            this.data.Value.Add(key, traceDatum);
-            this.Summary.UpdateRegionContacted(traceDatum);
+            lock (this.data)
+            {
+                this.data.Value.Add(key, traceDatum);
+                this.Summary.UpdateRegionContacted(traceDatum);
+            }
         }
 
         public void AddDatum(string key, object value)
         {
-            this.data.Value.Add(key, value);
+            lock (this.data)
+            {
+                this.data.Value.Add(key, value);
+            }
         }
 
         public void AddOrUpdateDatum(string key, object value)
         {
-            this.data.Value[key] = value;
+            lock (this.data)
+            {
+                this.data.Value[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);
+        }
     }
 }