diff --git a/Microsoft.Azure.Cosmos/src/Tracing/Trace.cs b/Microsoft.Azure.Cosmos/src/Tracing/Trace.cs index cb644f1ff0..83e8568165 100644 --- a/Microsoft.Azure.Cosmos/src/Tracing/Trace.cs +++ b/Microsoft.Azure.Cosmos/src/Tracing/Trace.cs @@ -6,10 +6,6 @@ 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 @@ -17,6 +13,7 @@ internal sealed class Trace : ITrace private static readonly IReadOnlyDictionary EmptyDictionary = new Dictionary(); private readonly List children; private readonly Lazy> data; + private volatile IReadOnlyDictionary materializedData = null; private ValueStopwatch stopwatch; private Trace( @@ -56,7 +53,7 @@ private Trace( public IReadOnlyList Children => this.children; - public IReadOnlyDictionary Data => this.data.IsValueCreated ? this.data.Value : Trace.EmptyDictionary; + public IReadOnlyDictionary Data => this.EnsureMaterializedData(); public void Dispose() { @@ -124,18 +121,55 @@ public static Trace GetRootTrace( public void AddDatum(string key, TraceDatum traceDatum) { - this.data.Value.Add(key, traceDatum); + lock (this.Name) + { + this.data.Value.Add(key, traceDatum); + this.materializedData = null; + } + this.Summary.UpdateRegionContacted(traceDatum); } public void AddDatum(string key, object value) { - this.data.Value.Add(key, value); + lock (this.Name) + { + this.data.Value.Add(key, value); + this.materializedData = null; + } } public void AddOrUpdateDatum(string key, object value) { - this.data.Value[key] = value; + lock (this.Name) + { + this.data.Value[key] = value; + this.materializedData = null; + } + } + + private IReadOnlyDictionary EnsureMaterializedData() + { + IReadOnlyDictionary snapshot = this.materializedData; + if (snapshot != null) + { + return snapshot; + } + + lock (this.Name) + { + if (snapshot != null) + { + return snapshot; + } + + if (!this.data.IsValueCreated) + { + return this.materializedData = EmptyDictionary; + } + + return this.materializedData = new Dictionary(this.data.Value); + } } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceTests.cs index 07472a3a44..609110b7f1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceTests.cs +++ b/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 tasks = new List(); + + 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); + } } } \ No newline at end of file