fix: remove usage of non-thread safe HashSet in AwsSdk pipeline wrapp…

…ers (#2855)
newrelic · Oct 28, 2024 · 87cace0 · 87cace0
1 parent e77683b
commit 87cace0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using NewRelic.Agent.Api;
+using NewRelic.Agent.Extensions.Collections;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 
 namespace NewRelic.Providers.Wrapper.AwsSdk
@@ -12,7 +13,7 @@ public class AwsSdkPipelineWrapper : IWrapper
         public bool IsTransactionRequired => true;
 
         private const string WrapperName = "AwsSdkPipelineWrapper";
-        private static HashSet<string> _unsupportedRequestTypes = new();
+        private static ConcurrentHashSet<string> _unsupportedRequestTypes = new();
 
         public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
         {
@@ -54,8 +55,11 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
                 return SQSRequestHandler.HandleSQSRequest(instrumentedMethodCall, agent, transaction, request, isAsync, executionContext);
             }
 
-            if (_unsupportedRequestTypes.Add(requestType)) // log once per unsupported request type
+            if (!_unsupportedRequestTypes.Contains(requestType))  // log once per unsupported request type
+            {                
                 agent.Logger.Debug($"AwsSdkPipelineWrapper: Unsupported request type: {requestType}. Returning NoOp delegate.");
+                _unsupportedRequestTypes.Add(requestType);
+            }
 
             return Delegates.NoOp;
         }

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SQSRequestHandler.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SQSRequestHandler.cs
@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.AwsSdk;
+using NewRelic.Agent.Extensions.Collections;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 using NewRelic.Reflection;
 
@@ -15,7 +16,7 @@ namespace NewRelic.Providers.Wrapper.AwsSdk
     internal static class SQSRequestHandler
     {
         private static readonly ConcurrentDictionary<Type, Func<object, object>> _getRequestResponseFromGeneric = new();
-        private static readonly HashSet<string> _unsupportedSQSRequestTypes = [];
+        private static readonly ConcurrentHashSet<string> _unsupportedSQSRequestTypes = [];
 
         public static AfterWrappedMethodDelegate HandleSQSRequest(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction, dynamic request, bool isAsync, dynamic executionContext)
         {
@@ -35,8 +36,11 @@ public static AfterWrappedMethodDelegate HandleSQSRequest(InstrumentedMethodCall
                     action = MessageBrokerAction.Purge;
                     break;
                 default:
-                    if (_unsupportedSQSRequestTypes.Add(requestType)) // log once per unsupported request type
+                    if (!_unsupportedSQSRequestTypes.Contains(requestType))  // log once per unsupported request type
+                    {
                         agent.Logger.Debug($"AwsSdkPipelineWrapper: SQS Request type {requestType} is not supported. Returning NoOp delegate.");
+                        _unsupportedSQSRequestTypes.Add(requestType);
+                    }
 
                     return Delegates.NoOp;
             }