feat: Enable DynamoDB entity relationships (#2923)

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Caching/LRUCache.cs

@@ -0,0 +1,121 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace NewRelic.Agent.Extensions.Caching
+{
+    /// <summary>
+    /// A thread-safe LRU cache implementation.
+    /// </summary>
+    /// <typeparam name="TKey"></typeparam>
+    /// <typeparam name="TValue"></typeparam>
+    public class LRUCache<TKey, TValue>
+    {
+        private readonly int _capacity;
+        private readonly Dictionary<TKey, LinkedListNode<CacheItem>> _cacheMap;
+        private readonly LinkedList<CacheItem> _lruList;
+        private readonly ReaderWriterLockSlim _lock = new();
+
+        public LRUCache(int capacity)
+        {
+            if (capacity <= 0)
+            {
+                throw new ArgumentException("Capacity must be greater than zero.", nameof(capacity));
+            }
+
+            _capacity = capacity;
+            _cacheMap = new Dictionary<TKey, LinkedListNode<CacheItem>>(capacity);
+            _lruList = new LinkedList<CacheItem>();
+        }
+
+        public TValue Get(TKey key)
+        {
+            _lock.EnterUpgradeableReadLock();
+            try
+            {
+                if (_cacheMap.TryGetValue(key, out var node))
+                {
+                    // Move the accessed node to the front of the list
+                    _lock.EnterWriteLock();
+                    try
+                    {
+                        _lruList.Remove(node);
+                        _lruList.AddFirst(node);
+                    }
+                    finally
+                    {
+                        _lock.ExitWriteLock();
+                    }
+                    return node.Value.Value;
+                }
+                throw new KeyNotFoundException("The given key was not present in the cache.");
+            }
+            finally
+            {
+                _lock.ExitUpgradeableReadLock();
+            }
+        }
+
+        public void Put(TKey key, TValue value)
+        {
+            _lock.EnterWriteLock();
+            try
+            {
+                if (_cacheMap.TryGetValue(key, out var node))
+                {
+                    // Update the value and move the node to the front of the list
+                    node.Value.Value = value;
+                    _lruList.Remove(node);
+                    _lruList.AddFirst(node);
+                }
+                else
+                {
+                    if (_cacheMap.Count >= _capacity)
+                    {
+                        // Remove the least recently used item
+                        var lruNode = _lruList.Last;
+                        _cacheMap.Remove(lruNode.Value.Key);
+                        _lruList.RemoveLast();
+                    }
+
+                    // Add the new item to the cache
+                    var cacheItem = new CacheItem(key, value);
+                    var newNode = new LinkedListNode<CacheItem>(cacheItem);
+                    _lruList.AddFirst(newNode);
+                    _cacheMap[key] = newNode;
+                }
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
+        }
+        public bool ContainsKey(TKey key)
+        {
+            _lock.EnterReadLock();
+            try
+            {
+                return _cacheMap.ContainsKey(key);
+            }
+            finally
+            {
+                _lock.ExitReadLock();
+            }
+        }
+
+        private class CacheItem
+        {
+            public TKey Key { get; }
+            public TValue Value { get; set; }
+
+            public CacheItem(TKey key, TValue value)
+            {
+                Key = key;
+                Value = value;
+            }
+        }
+    }
+}

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Caching/LRUHashSet.cs

@@ -0,0 +1,114 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace NewRelic.Agent.Extensions.Caching
+{
+    /// <summary>
+    /// A thread-safe LRU HashSet implementation.
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    public class LRUHashSet<T>
+    {
+        private readonly int _capacity;
+        private readonly HashSet<T> _hashSet;
+        private readonly LinkedList<T> _lruList;
+        private readonly ReaderWriterLockSlim _lock = new();
+
+        public LRUHashSet(int capacity)
+        {
+            if (capacity <= 0)
+            {
+                throw new ArgumentException("Capacity must be greater than zero.", nameof(capacity));
+            }
+
+            _capacity = capacity;
+            _hashSet = new HashSet<T>();
+            _lruList = new LinkedList<T>();
+        }
+
+        public bool Add(T item)
+        {
+            _lock.EnterWriteLock();
+            try
+            {
+                if (_hashSet.Contains(item))
+                {
+                    // Move the accessed item to the front of the list
+                    _lruList.Remove(item);
+                    _lruList.AddFirst(item);
+                    return false;
+                }
+                else
+                {
+                    if (_hashSet.Count >= _capacity)
+                    {
+                        // Remove the least recently used item
+                        var lruItem = _lruList.Last.Value;
+                        _hashSet.Remove(lruItem);
+                        _lruList.RemoveLast();
+                    }
+
+                    // Add the new item to the set and list
+                    _hashSet.Add(item);
+                    _lruList.AddFirst(item);
+                    return true;
+                }
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
+        }
+
+        public bool Contains(T item)
+        {
+            _lock.EnterReadLock();
+            try
+            {
+                return _hashSet.Contains(item);
+            }
+            finally
+            {
+                _lock.ExitReadLock();
+            }
+        }
+
+        public bool Remove(T item)
+        {
+            _lock.EnterWriteLock();
+            try
+            {
+                if (_hashSet.Remove(item))
+                {
+                    _lruList.Remove(item);
+                    return true;
+                }
+                return false;
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
+        }
+
+        public int Count
+        {
+            get
+            {
+                _lock.EnterReadLock();
+                try
+                {
+                    return _hashSet.Count;
+                }
+                finally
+                {
+                    _lock.ExitReadLock();
+                }
+            }
+        }
+    }
+}

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Caching/WeakReferenceKey.cs

@@ -0,0 +1,50 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System;
+
+namespace NewRelic.Agent.Extensions.Caching
+{
+    /// <summary>
+    /// Creates an object that can be used as a dictionary key, which holds a WeakReference&lt;T&gt;
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    public class WeakReferenceKey<T> where T : class
+    {
+        private WeakReference<T> WeakReference { get; }
+
+        public WeakReferenceKey(T cacheKey)
+        {
+            WeakReference = new WeakReference<T>(cacheKey);
+        }
+
+        public override bool Equals(object obj)
+        {
+            if (obj is WeakReferenceKey<T> otherKey)
+            {
+                if (WeakReference.TryGetTarget(out var thisTarget) &&
+                    otherKey.WeakReference.TryGetTarget(out var otherTarget))
+                {
+                    return ReferenceEquals(thisTarget, otherTarget);
+                }
+            }
+
+            return false;
+        }
+
+        public override int GetHashCode()
+        {
+            if (WeakReference.TryGetTarget(out var target))
+            {
+                return target.GetHashCode();
+            }
+
+            return 0;
+        }
+
+        /// <summary>
+        /// Gets the value from the WeakReference or null if the target has been garbage collected.
+        /// </summary>
+        public T Value => WeakReference.TryGetTarget(out var target) ? target : null;
+    }
+}

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AmazonServiceClientWrapper.cs

@@ -4,49 +4,62 @@
 using System;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.AwsSdk;
+using NewRelic.Agent.Extensions.Caching;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 
-namespace NewRelic.Providers.Wrapper.AwsSdk
+namespace NewRelic.Providers.Wrapper.AwsSdk;
+
+public class AmazonServiceClientWrapper : IWrapper
 {
-    public class AmazonServiceClientWrapper : IWrapper
+    private const int LRUCapacity = 100;
+    // cache the account id per instance of AmazonServiceClient.Config
+    public static LRUCache<WeakReferenceKey<object>, string> AwsAccountIdByClientConfigCache = new(LRUCapacity);
+
+    // cache instances of AmazonServiceClient
+    private static readonly LRUHashSet<WeakReferenceKey<object>> AmazonServiceClientInstanceCache = new(LRUCapacity);
+
+    public bool IsTransactionRequired => false;
+
+    public CanWrapResponse CanWrap(InstrumentedMethodInfo instrumentedMethodInfo)
     {
-        /// <summary>
-        /// The AWS account id.
-        /// Parsed from the access key in the credentials of the client - or fall back to the configuration value if parsing fails.
-        /// Assumes only a single account id is used in the application.
-        /// </summary>
-        public static string AwsAccountId { get; private set; }
+        return new CanWrapResponse(instrumentedMethodInfo.RequestedWrapperName == nameof(AmazonServiceClientWrapper));
+    }
+
+    public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
+    {
+        object client = instrumentedMethodCall.MethodCall.InvocationTarget;
 
-        public bool IsTransactionRequired => false;
+        var weakReferenceKey = new WeakReferenceKey<object>(client);
+        if (AmazonServiceClientInstanceCache.Contains(weakReferenceKey)) // don't do anything if we've already seen this client instance
+            return Delegates.NoOp;
 
-        public CanWrapResponse CanWrap(InstrumentedMethodInfo instrumentedMethodInfo)
+        AmazonServiceClientInstanceCache.Add(weakReferenceKey);
+
+        string awsAccountId;
+        try
         {
-            return new CanWrapResponse(instrumentedMethodInfo.RequestedWrapperName == nameof(AmazonServiceClientWrapper));
+            // get the AWSCredentials parameter
+            dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
+
+            dynamic immutableCredentials = awsCredentials.GetCredentials();
+            string accessKey = immutableCredentials.AccessKey;
+
+            // convert the access key to an account id
+            awsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
+        }
+        catch (Exception e)
+        {
+            agent.Logger.Info($"Unable to parse AWS Account ID from AccessKey. Using AccountId from configuration instead. Exception: {e.Message}");
+            awsAccountId = agent.Configuration.AwsAccountId;
         }
 
-        public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
+        return Delegates.GetDelegateFor(onComplete: () =>
         {
-            if (AwsAccountId != null)
-                return Delegates.NoOp;
-
-            try
-            {
-                // get the AWSCredentials parameter
-                dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
-
-                dynamic immutableCredentials = awsCredentials.GetCredentials();
-                string accessKey = immutableCredentials.AccessKey;
-
-                // convert the access key to an account id
-                AwsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
-            }
-            catch (Exception e)
-            {
-                agent.Logger.Info($"Unable to parse AWS Account ID from AccessKey. Using AccountId from configuration instead. Exception: {e.Message}");
-                AwsAccountId = agent.Configuration.AwsAccountId;
-            }
+            // get the _config field from the client
+            object clientConfig = ((dynamic)client).Config;
 
-            return Delegates.NoOp;
-        }
+            // cache the account id using clientConfig as the key
+            AwsAccountIdByClientConfigCache.Put(new WeakReferenceKey<object>(clientConfig), awsAccountId);
+        });
     }
 }