From b505dcd6e9370437b840cdf3cef776370ef60584 Mon Sep 17 00:00:00 2001
From: Martijn Laarman <Mpdreamz@gmail.com>
Date: Mon, 19 Feb 2024 15:38:55 +0100
Subject: [PATCH] Fix #2248 duplicate key errors on dropped span stats update
 (#2283)

* Fix #2248 duplicate key errors on dropped span stats update

* remove erroneous additional lock inside update lambda

* dotnet format

* Ensure we still respect max count of 128
---
 src/Elastic.Apm/Model/Transaction.cs | 43 +++++++++++++---------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/src/Elastic.Apm/Model/Transaction.cs b/src/Elastic.Apm/Model/Transaction.cs
index 7421091ea..659f24c6b 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,35 @@ 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)
-				}
-			};
+			lock (_lock)
+				_droppedSpanStatsMap ??= new ConcurrentDictionary<DroppedSpanStatsKey, DroppedSpanStats>();
 		}
-		else
+
+		lock (_lock)
 		{
 			if (_droppedSpanStatsMap.Count >= 128)
 				return;
-
-			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));
-			}
+			//AddOrUpdate callbacks can run multiple times so still wrapping this in a lock
+			var key = new DroppedSpanStatsKey(serviceTargetType, serviceTargetName, outcome);
+			_droppedSpanStatsMap.AddOrUpdate(key,
+				 _ => new DroppedSpanStats(serviceTargetType, serviceTargetName, destinationServiceResource, outcome, duration),
+				 (_, stats) =>
+				 {
+					 stats.Duration ??=
+						 new DroppedSpanStats.DroppedSpanDuration { Sum = new DroppedSpanStats.DroppedSpanDuration.DroppedSpanDurationSum() };
+
+					 stats.Duration.Count++;
+					 stats.Duration.Sum.UsRaw += duration;
+					 return stats;
+				 });
 		}
 	}