summary: Resolves multiple issues with StackExchange.Redis v2+ profiling (#1504)

fix: Allow StackExchange.Redis v2+ profiling to start outside of a transaction. (#1501 )
fix: Fix a memory leak when using StackExchange.Redis v2+. (#1473)
chore: Add more units tests
chore: Update logic to make better use of ConcurrentDictionary features.
chore: Mark weakTransaction as nullable.

---------

Co-authored-by: Chris Ventura <[email protected]>
Co-authored-by: Marty Tippin <[email protected]>

Author: 3 people

src/Agent/NewRelic/Agent/Core/Agent.cs

@@ -63,7 +63,7 @@ public class Agent : IAgent // any changes to api, update the interface in exten
         private readonly ILogEventAggregator _logEventAggregator;
         private readonly ILogContextDataFilter _logContextDataFilter;
         private Extensions.Logging.ILogger _logger;
-        private IStackExchangeRedisCache _stackExchangeRedisCache;
+        private volatile IStackExchangeRedisCache _stackExchangeRedisCache;
 
         public Agent(ITransactionService transactionService, ITransactionTransformer transactionTransformer,
             IThreadPoolStatic threadPoolStatic, ITransactionMetricNameMaker transactionMetricNameMaker, IPathHashMaker pathHashMaker,

src/Agent/NewRelic/Agent/Core/Segments/Segment.cs

@@ -115,10 +115,16 @@ public Segment(TimeSpan relativeStartTime, TimeSpan? duration, Segment segment,
         public bool IsExternal => Data.SpanCategory == SpanCategory.Http;
 
         private string _spanId;
+        /// <summary>
+        /// Gets the SpanId for the segment.  If SpanId has not been set it will create one.
+        /// This call could potentially generate more than one Id if GET is called from multiple threads at the same time.
+        /// Current usage of this property does not do this.
+        /// </summary>
         public string SpanId
         {
             get
             {
+                // If _spanId is null and this is called rapidly from different threads, the returned value could be different for each.
                 return _spanId ?? (_spanId = GuidGenerator.GenerateNewRelicGuid());
             }
             set
@@ -135,9 +141,9 @@ public void End()
                 var endTime = _transactionSegmentState.GetRelativeTime();
                 RelativeEndTime = endTime;
 
-                if (Agent.Instance.StackExchangeRedisCache != null && Agent.Instance.StackExchangeRedisCache.Count > 0)
+                if (Agent.Instance.StackExchangeRedisCache != null)
                 {
-                    Agent.Instance.StackExchangeRedisCache.Harvest(SpanId);
+                    Agent.Instance.StackExchangeRedisCache.Harvest(this);
                 }
 
                 Finish();

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Helpers/IStackExchangeRedisCache.cs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 using System;
+using NewRelic.Agent.Api;
 
 namespace NewRelic.Agent.Extensions.Helpers
 {
@@ -11,8 +12,6 @@ namespace NewRelic.Agent.Extensions.Helpers
     /// </summary>
     public interface IStackExchangeRedisCache : IDisposable
     {
-        void Harvest(string spanId);
-
-        int Count { get; }
+        void Harvest(ISegment segment);
     }
 }

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/StackExchangeRedis2Plus/ConnectWrapperV2.cs

@@ -11,7 +11,7 @@ namespace NewRelic.Providers.Wrapper.StackExchangeRedis2Plus
 {
     public class ConnectWrapperV2 : IWrapper
     {
-        public bool IsTransactionRequired => true;
+        public bool IsTransactionRequired => false;
 
         private const string WrapperName = "stackexchangeredis-connect";
 

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/StackExchangeRedis2Plus/SessionCache.cs

@@ -19,7 +19,7 @@ public class SessionCache : IStackExchangeRedisCache
     {
         private readonly EventWaitHandle _stopHandle = new EventWaitHandle(false, EventResetMode.ManualReset);
 
-        private readonly ConcurrentDictionary<string, ProfilingSession> _sessionCache = new ConcurrentDictionary<string, ProfilingSession>();
+        private readonly ConcurrentDictionary<ISegment, ProfilingSession> _sessionCache = new ConcurrentDictionary<ISegment, ProfilingSession>();
 
         private readonly IAgent _agent;
 
@@ -34,28 +34,26 @@ public SessionCache(IAgent agent, int invocationTargetHashCode)
         }
 
         /// <summary>
-        /// Finishes a profiling session for the segment indicated by the span id and creates a child DataStoreSegment for each command in the session.
+        /// Finishes a profiling session for the segment and creates a child DataStoreSegment for each command in the session.
         /// </summary>
-        /// <param name="spanId">Span ID of the segment being finalized.</param>
-        public void Harvest(string spanId)
+        /// <param name="hostSegment">Segment being finalized.</param>
+        public void Harvest(ISegment hostSegment)
         {
             // If we can't remove the session, it doesn't exist, so do nothing and return.
-            if (!_sessionCache.TryRemove(spanId, out var sessionData))
+            if (!_sessionCache.TryRemove(hostSegment, out var sessionData))
             {
                 return;
             }
 
             // Get the transaction from the session
-            var transaction = sessionData.UserToken as ITransaction;
-
-            // We want to make sure to finish the session even if the transaction is done so that it is not orphaned.
-            var commands = sessionData.FinishProfiling();
-            if (transaction.IsFinished)
+            var weakTransaction = sessionData.UserToken as WeakReference<ITransaction>;
+            if (!(weakTransaction?.TryGetTarget(out var transaction) ?? false) || transaction.IsFinished)
             {
                 return;
             }
 
             var xTransaction = (ITransactionExperimental)transaction;
+            var commands = sessionData.FinishProfiling();
             foreach (var command in commands)
             {
                 // We need to build the relative start and stop time based on the transaction start time.
@@ -107,7 +105,7 @@ public Func<ProfilingSession> GetProfilingSession()
 
                 // Don't want to save data to a session outside of a transaction or to a NoOp - no way to clean it up easily or reliably.
                 var transaction = _agent.CurrentTransaction;
-                if (!transaction.IsValid)
+                if (transaction.IsFinished || !transaction.IsValid)
                 {
                     return null;
                 }
@@ -119,43 +117,21 @@ public Func<ProfilingSession> GetProfilingSession()
                     return null;
                 }
 
-                // Use the spanid of the segment as the key for the cache.
-                var spanId = segment.SpanId;
-                if (string.IsNullOrWhiteSpace(spanId))
-                {
-                    return null;
-                }
-
-                // During async operations, the transaction can get lost and report as NoOp so we store a reference to it in the session.
-                if (!_sessionCache.TryGetValue(spanId, out var sessionData))
-                {
-                    sessionData = new ProfilingSession(transaction);
-                    _sessionCache.TryAdd(spanId, sessionData);
-                }
-
-                return sessionData;
+                return _sessionCache.GetOrAdd(segment, GetProfilingSession);
             };
         }
 
-        public int Count
+        private ProfilingSession GetProfilingSession(ISegment segment)
         {
-            get
-            {
-                return _sessionCache.Count;
-            }
+            return new ProfilingSession(new WeakReference<ITransaction>(_agent.CurrentTransaction, false));
         }
 
         // Clean up the handles, sessions, and wipe the dictionary.
         public void Dispose()
         {
             _stopHandle.Set();
-            _stopHandle.Dispose();
-            foreach (var cachedSession in _sessionCache.Values)
-            {
-                _ = cachedSession.FinishProfiling();
-            }
-
             _sessionCache.Clear();
+            _stopHandle.Dispose();
         }
     }
 }

tests/Agent/UnitTests/CompositeTests/SegmentsTests.cs

@@ -792,6 +792,82 @@ public void SegmentEndWithExceptionCapturesErrorAttributes()
             SpanAssertions.HasAttributes(expectedSpanErrorAttributes, AttributeClassification.AgentAttributes, spanWithError);
         }
 
+        [Test]
+        public void SegmentEnd_StackExchangeRedis_EndCreatesSegment()
+        {
+            var tx = _agent.CreateTransaction(
+                isWeb: false,
+                category: "testing",
+                transactionDisplayName: "test",
+                doNotTrackAsUnitOfWork: true);
+            var segment = _agent.StartStackExchangeRedisDatastoreRequestSegmentOrThrow("set", DatastoreVendor.Redis, TimeSpan.Zero, TimeSpan.FromMilliseconds(1));
+            segment.EndStackExchangeRedis();
+            tx.End();
+
+            _compositeTestAgent.Harvest();
+
+            var spanEvents = _compositeTestAgent.SpanEvents.ToArray();
+            Assert.AreEqual(2, spanEvents.Length);
+
+            var expectedSpanErrorAttributes = new List<ExpectedAttribute>
+            {
+                new ExpectedAttribute { Key = "name", Value = "Datastore/operation/Redis/set" },
+                new ExpectedAttribute { Key = "category", Value = "datastore" },
+            };
+
+            var datastoreSpan= spanEvents[1];
+            SpanAssertions.HasAttributes(expectedSpanErrorAttributes, AttributeClassification.Intrinsics, datastoreSpan);
+        }
+
+        [Test]
+        public void SegmentEnd_StackExchangeRedis_SessionCacheIsNull_HarvestNotCalled()
+        {
+            _agent.StackExchangeRedisCache = null;
+            var tx = _agent.CreateTransaction(
+                isWeb: false,
+                category: "testing",
+                transactionDisplayName: "test",
+                doNotTrackAsUnitOfWork: true);
+            var segment = _agent.StartCustomSegmentOrThrow("parentSegment");
+
+            segment.End();
+            tx.End();
+
+            _compositeTestAgent.Harvest();
+
+            var spanEvents = _compositeTestAgent.SpanEvents.ToArray();
+            Assert.AreEqual(2, spanEvents.Length);
+        }
+
+        [Test]
+        public void SegmentEnd_StackExchangeRedis_SessionCacheNotNull_HarvestCalled()
+        {
+            _agent.StackExchangeRedisCache = new StackExchangeRedisCacheMock(_agent);
+            var tx = _agent.CreateTransaction(
+                isWeb: false,
+                category: "testing",
+                transactionDisplayName: "test",
+                doNotTrackAsUnitOfWork: true);
+            var segment = _agent.StartCustomSegmentOrThrow("parentSegment");
+
+            segment.End();
+            tx.End();
+
+            _compositeTestAgent.Harvest();
+
+            var spanEvents = _compositeTestAgent.SpanEvents.ToArray();
+            Assert.AreEqual(3, spanEvents.Length);
+
+            var expectedSpanErrorAttributes = new List<ExpectedAttribute>
+            {
+                new ExpectedAttribute { Key = "name", Value = "Datastore/operation/Redis/set" },
+                new ExpectedAttribute { Key = "category", Value = "datastore" },
+            };
+
+            var datastoreSpan = spanEvents[2];
+            SpanAssertions.HasAttributes(expectedSpanErrorAttributes, AttributeClassification.Intrinsics, datastoreSpan);
+        }
+
         #region Helper methods
 
         private static TransactionTraceSegment GetFirstSegmentOrThrow()
@@ -810,4 +886,26 @@ private static TransactionTraceSegment GetFirstSegmentOrThrow()
 
         #endregion Helper methods
     }
+
+    public class StackExchangeRedisCacheMock : NewRelic.Agent.Extensions.Helpers.IStackExchangeRedisCache
+    {
+        private IAgent _agent;
+
+        public StackExchangeRedisCacheMock(IAgent agent)
+        {
+            _agent = agent;
+        }
+
+        public void Dispose()
+        {
+            // Do nothing
+        }
+
+        public void Harvest(ISegment segment)
+        {
+            // Create a segment if callled so we can test it
+            var testSegment = _agent.StartStackExchangeRedisDatastoreRequestSegmentOrThrow("set", DatastoreVendor.Redis, TimeSpan.Zero, TimeSpan.FromMilliseconds(1));
+            testSegment.EndStackExchangeRedis();
+        }
+    }
 }