Merge pull request #307 from luizfds/master

Add locking mechanism for Subsegments in Entity.AddException method
aws · Nov 13, 2024 · ed09306 · ed09306
2 parents 7f88bc7 + 47495bb
commit ed09306
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 1 deletion.
diff --git a/sdk/src/Core/Internal/Entities/Entity.cs b/sdk/src/Core/Internal/Entities/Entity.cs
@@ -514,7 +514,13 @@ public void AddException(Exception e)
         {
             HasFault = true;
             Cause = new Cause();
-            Cause.AddException(AWSXRayRecorder.Instance.ExceptionSerializationStrategy.DescribeException(e, Subsegments));
+            List<Subsegment> subsegmentsCopy;
+            lock (_lazySubsegments.Value)
+            {
+                subsegmentsCopy = _lazySubsegments.Value.ToList(); // Create a copy to avoid holding the lock during serialization
+            }
+
+            Cause.AddException(AWSXRayRecorder.Instance.ExceptionSerializationStrategy.DescribeException(e, subsegmentsCopy));
         }
 
         /// <summary>

diff --git a/sdk/test/UnitTests/SegmentTests.cs b/sdk/test/UnitTests/SegmentTests.cs
@@ -21,6 +21,8 @@
 using System.Linq;
 using System.Runtime.Serialization;
 using System.Runtime.Serialization.Formatters.Binary;
+using System.Threading;
+using System.Threading.Tasks;
 using Amazon.XRay.Recorder.Core.Exceptions;
 using Amazon.XRay.Recorder.Core.Internal.Entities;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -300,6 +302,48 @@ public void TestAddException()
             Assert.ReferenceEquals(e, descriptor.Exception);
         }
 
+        [TestMethod]
+        public void TestAddExceptionWithConcurrentSubsegmentAddition()
+        {
+            const int EXCEPTION_ITERATIONS = 100;
+            const int SUBSEGMENT_ITERATIONS = 100_000;
+            var segment = new Segment("test", TraceId);
+            var raceConditionOccurred = false;
+            var cancellationTokenSource = new CancellationTokenSource();
+
+            var addExceptionTask = Task.Run(() =>
+            {
+                Parallel.ForEach(Enumerable.Range(0, EXCEPTION_ITERATIONS), _ =>
+                {
+                    try
+                    {
+                        var exception = new Exception("Test Exception");
+                        segment.AddException(exception);
+                    }
+                    catch (InvalidOperationException)
+                    {
+                        raceConditionOccurred = true;
+                    }
+                });
+            });
+
+            var addSubsegmentsTask = Task.Run(() =>
+            {
+                try
+                {
+                    Parallel.ForEach(Enumerable.Range(0, SUBSEGMENT_ITERATIONS), new ParallelOptions { CancellationToken = cancellationTokenSource.Token }, _ =>
+                    {
+                        segment.AddSubsegment(new Subsegment("TestSubsegment"));
+                    });
+                }
+                catch (OperationCanceledException) { }
+            });
+
+            Task.WaitAll(addExceptionTask, addSubsegmentsTask);
+
+            Assert.IsFalse(raceConditionOccurred, "A race condition error occurred unexpectedly.");
+        }
+
         [TestMethod]
         public void TestHttpOverwriteValue()
         {