From c5a1c3224eccbcf8ea6bb94a17db66550af5ae6d Mon Sep 17 00:00:00 2001
From: Martijn Laarman <Mpdreamz@gmail.com>
Date: Wed, 7 Feb 2024 11:42:37 +0100
Subject: [PATCH] Fix #2248 duplicate key errors on dropped span stats update

---
 src/Elastic.Apm/Model/Transaction.cs | 49 +++++++++++++---------------
 1 file changed, 22 insertions(+), 27 deletions(-)

diff --git a/src/Elastic.Apm/Model/Transaction.cs b/src/Elastic.Apm/Model/Transaction.cs
index 7421091ea..381886634 100644
--- a/src/Elastic.Apm/Model/Transaction.cs
+++ b/src/Elastic.Apm/Model/Transaction.cs
@@ -357,7 +357,7 @@ private void CheckAndCaptureBaggage()
 	/// <summary>
 	/// Internal dictionary to keep track of and look up dropped span stats.
 	/// </summary>
-	private Dictionary<DroppedSpanStatsKey, DroppedSpanStats> _droppedSpanStatsMap;
+	private ConcurrentDictionary<DroppedSpanStatsKey, DroppedSpanStats> _droppedSpanStatsMap;
 
 	private bool _isEnded;
 
@@ -552,38 +552,33 @@ private Activity StartActivity(bool shouldRestartTrace)
 		return activity;
 	}
 
+	private readonly object _lock = new();
 	internal void UpdateDroppedSpanStats(string serviceTargetType, string serviceTargetName, string destinationServiceResource, Outcome outcome,
 		double duration
 	)
 	{
+		//lock the lazy initialization of the dictionary
 		if (_droppedSpanStatsMap == null)
-		{
-			_droppedSpanStatsMap = new Dictionary<DroppedSpanStatsKey, DroppedSpanStats>
-			{
-				{
-					new DroppedSpanStatsKey(serviceTargetType, serviceTargetName, outcome),
-					new DroppedSpanStats(serviceTargetType, serviceTargetName, destinationServiceResource, outcome, duration)
-				}
-			};
-		}
-		else
-		{
-			if (_droppedSpanStatsMap.Count >= 128)
-				return;
+			lock (_lock) _droppedSpanStatsMap ??= new ConcurrentDictionary<DroppedSpanStatsKey, DroppedSpanStats>();
 
-			if (_droppedSpanStatsMap.TryGetValue(new DroppedSpanStatsKey(serviceTargetType, serviceTargetName, outcome), out var item))
-			{
-				item.Duration ??=
-					new DroppedSpanStats.DroppedSpanDuration { Sum = new DroppedSpanStats.DroppedSpanDuration.DroppedSpanDurationSum() };
-
-				item.Duration.Count++;
-				item.Duration.Sum.UsRaw += duration;
-			}
-			else
-			{
-				_droppedSpanStatsMap.Add(new DroppedSpanStatsKey(serviceTargetType, serviceTargetName, outcome),
-					new DroppedSpanStats(serviceTargetType, serviceTargetName, destinationServiceResource, outcome, duration));
-			}
+		lock (_lock)
+		{
+			//AddOrUpdate callbacks can run multiple times so still wrapping this in a lock
+			var key = new DroppedSpanStatsKey(serviceTargetType, serviceTargetName, outcome);
+			_droppedSpanStatsMap.AddOrUpdate(key,
+				 key => new DroppedSpanStats(serviceTargetType, serviceTargetName, destinationServiceResource, outcome, duration),
+				 (_, value) =>
+				 {
+					 lock (_lock)
+					 {
+						  value.Duration ??=
+							   new DroppedSpanStats.DroppedSpanDuration { Sum = new DroppedSpanStats.DroppedSpanDuration.DroppedSpanDurationSum() };
+
+						  value.Duration.Count++;
+						  value.Duration.Sum.UsRaw += duration;
+						  return value;
+					 }
+				 });
 		}
 	}