Skip to content

Commit f70b64a

Browse files
tippmar-nrtippmar-nr
authored
feat: Enable DynamoDB entity relationships (#2923)
1 parent c55e43c commit f70b64a

File tree

20 files changed

+1287
-190
lines changed

20 files changed

+1287
-190
lines changed

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

+121
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System;
5+
using System.Collections.Generic;
6+
using System.Threading;
7+
8+
namespace NewRelic.Agent.Extensions.Caching
9+
{
10+
/// <summary>
11+
/// A thread-safe LRU cache implementation.
12+
/// </summary>
13+
/// <typeparam name="TKey"></typeparam>
14+
/// <typeparam name="TValue"></typeparam>
15+
public class LRUCache<TKey, TValue>
16+
{
17+
private readonly int _capacity;
18+
private readonly Dictionary<TKey, LinkedListNode<CacheItem>> _cacheMap;
19+
private readonly LinkedList<CacheItem> _lruList;
20+
private readonly ReaderWriterLockSlim _lock = new();
21+
22+
public LRUCache(int capacity)
23+
{
24+
if (capacity <= 0)
25+
{
26+
throw new ArgumentException("Capacity must be greater than zero.", nameof(capacity));
27+
}
28+
29+
_capacity = capacity;
30+
_cacheMap = new Dictionary<TKey, LinkedListNode<CacheItem>>(capacity);
31+
_lruList = new LinkedList<CacheItem>();
32+
}
33+
34+
public TValue Get(TKey key)
35+
{
36+
_lock.EnterUpgradeableReadLock();
37+
try
38+
{
39+
if (_cacheMap.TryGetValue(key, out var node))
40+
{
41+
// Move the accessed node to the front of the list
42+
_lock.EnterWriteLock();
43+
try
44+
{
45+
_lruList.Remove(node);
46+
_lruList.AddFirst(node);
47+
}
48+
finally
49+
{
50+
_lock.ExitWriteLock();
51+
}
52+
return node.Value.Value;
53+
}
54+
throw new KeyNotFoundException("The given key was not present in the cache.");
55+
}
56+
finally
57+
{
58+
_lock.ExitUpgradeableReadLock();
59+
}
60+
}
61+
62+
public void Put(TKey key, TValue value)
63+
{
64+
_lock.EnterWriteLock();
65+
try
66+
{
67+
if (_cacheMap.TryGetValue(key, out var node))
68+
{
69+
// Update the value and move the node to the front of the list
70+
node.Value.Value = value;
71+
_lruList.Remove(node);
72+
_lruList.AddFirst(node);
73+
}
74+
else
75+
{
76+
if (_cacheMap.Count >= _capacity)
77+
{
78+
// Remove the least recently used item
79+
var lruNode = _lruList.Last;
80+
_cacheMap.Remove(lruNode.Value.Key);
81+
_lruList.RemoveLast();
82+
}
83+
84+
// Add the new item to the cache
85+
var cacheItem = new CacheItem(key, value);
86+
var newNode = new LinkedListNode<CacheItem>(cacheItem);
87+
_lruList.AddFirst(newNode);
88+
_cacheMap[key] = newNode;
89+
}
90+
}
91+
finally
92+
{
93+
_lock.ExitWriteLock();
94+
}
95+
}
96+
public bool ContainsKey(TKey key)
97+
{
98+
_lock.EnterReadLock();
99+
try
100+
{
101+
return _cacheMap.ContainsKey(key);
102+
}
103+
finally
104+
{
105+
_lock.ExitReadLock();
106+
}
107+
}
108+
109+
private class CacheItem
110+
{
111+
public TKey Key { get; }
112+
public TValue Value { get; set; }
113+
114+
public CacheItem(TKey key, TValue value)
115+
{
116+
Key = key;
117+
Value = value;
118+
}
119+
}
120+
}
121+
}

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

+114
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System;
5+
using System.Collections.Generic;
6+
using System.Threading;
7+
8+
namespace NewRelic.Agent.Extensions.Caching
9+
{
10+
/// <summary>
11+
/// A thread-safe LRU HashSet implementation.
12+
/// </summary>
13+
/// <typeparam name="T"></typeparam>
14+
public class LRUHashSet<T>
15+
{
16+
private readonly int _capacity;
17+
private readonly HashSet<T> _hashSet;
18+
private readonly LinkedList<T> _lruList;
19+
private readonly ReaderWriterLockSlim _lock = new();
20+
21+
public LRUHashSet(int capacity)
22+
{
23+
if (capacity <= 0)
24+
{
25+
throw new ArgumentException("Capacity must be greater than zero.", nameof(capacity));
26+
}
27+
28+
_capacity = capacity;
29+
_hashSet = new HashSet<T>();
30+
_lruList = new LinkedList<T>();
31+
}
32+
33+
public bool Add(T item)
34+
{
35+
_lock.EnterWriteLock();
36+
try
37+
{
38+
if (_hashSet.Contains(item))
39+
{
40+
// Move the accessed item to the front of the list
41+
_lruList.Remove(item);
42+
_lruList.AddFirst(item);
43+
return false;
44+
}
45+
else
46+
{
47+
if (_hashSet.Count >= _capacity)
48+
{
49+
// Remove the least recently used item
50+
var lruItem = _lruList.Last.Value;
51+
_hashSet.Remove(lruItem);
52+
_lruList.RemoveLast();
53+
}
54+
55+
// Add the new item to the set and list
56+
_hashSet.Add(item);
57+
_lruList.AddFirst(item);
58+
return true;
59+
}
60+
}
61+
finally
62+
{
63+
_lock.ExitWriteLock();
64+
}
65+
}
66+
67+
public bool Contains(T item)
68+
{
69+
_lock.EnterReadLock();
70+
try
71+
{
72+
return _hashSet.Contains(item);
73+
}
74+
finally
75+
{
76+
_lock.ExitReadLock();
77+
}
78+
}
79+
80+
public bool Remove(T item)
81+
{
82+
_lock.EnterWriteLock();
83+
try
84+
{
85+
if (_hashSet.Remove(item))
86+
{
87+
_lruList.Remove(item);
88+
return true;
89+
}
90+
return false;
91+
}
92+
finally
93+
{
94+
_lock.ExitWriteLock();
95+
}
96+
}
97+
98+
public int Count
99+
{
100+
get
101+
{
102+
_lock.EnterReadLock();
103+
try
104+
{
105+
return _hashSet.Count;
106+
}
107+
finally
108+
{
109+
_lock.ExitReadLock();
110+
}
111+
}
112+
}
113+
}
114+
}

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

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System;
5+
6+
namespace NewRelic.Agent.Extensions.Caching
7+
{
8+
/// <summary>
9+
/// Creates an object that can be used as a dictionary key, which holds a WeakReference&lt;T&gt;
10+
/// </summary>
11+
/// <typeparam name="T"></typeparam>
12+
public class WeakReferenceKey<T> where T : class
13+
{
14+
private WeakReference<T> WeakReference { get; }
15+
16+
public WeakReferenceKey(T cacheKey)
17+
{
18+
WeakReference = new WeakReference<T>(cacheKey);
19+
}
20+
21+
public override bool Equals(object obj)
22+
{
23+
if (obj is WeakReferenceKey<T> otherKey)
24+
{
25+
if (WeakReference.TryGetTarget(out var thisTarget) &&
26+
otherKey.WeakReference.TryGetTarget(out var otherTarget))
27+
{
28+
return ReferenceEquals(thisTarget, otherTarget);
29+
}
30+
}
31+
32+
return false;
33+
}
34+
35+
public override int GetHashCode()
36+
{
37+
if (WeakReference.TryGetTarget(out var target))
38+
{
39+
return target.GetHashCode();
40+
}
41+
42+
return 0;
43+
}
44+
45+
/// <summary>
46+
/// Gets the value from the WeakReference or null if the target has been garbage collected.
47+
/// </summary>
48+
public T Value => WeakReference.TryGetTarget(out var target) ? target : null;
49+
}
50+
}

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

+46-33
Original file line numberDiff line numberDiff line change
@@ -4,49 +4,62 @@
44
using System;
55
using NewRelic.Agent.Api;
66
using NewRelic.Agent.Extensions.AwsSdk;
7+
using NewRelic.Agent.Extensions.Caching;
78
using NewRelic.Agent.Extensions.Providers.Wrapper;
89

9-
namespace NewRelic.Providers.Wrapper.AwsSdk
10+
namespace NewRelic.Providers.Wrapper.AwsSdk;
11+
12+
public class AmazonServiceClientWrapper : IWrapper
1013
{
11-
public class AmazonServiceClientWrapper : IWrapper
14+
private const int LRUCapacity = 100;
15+
// cache the account id per instance of AmazonServiceClient.Config
16+
public static LRUCache<WeakReferenceKey<object>, string> AwsAccountIdByClientConfigCache = new(LRUCapacity);
17+
18+
// cache instances of AmazonServiceClient
19+
private static readonly LRUHashSet<WeakReferenceKey<object>> AmazonServiceClientInstanceCache = new(LRUCapacity);
20+
21+
public bool IsTransactionRequired => false;
22+
23+
public CanWrapResponse CanWrap(InstrumentedMethodInfo instrumentedMethodInfo)
1224
{
13-
/// <summary>
14-
/// The AWS account id.
15-
/// Parsed from the access key in the credentials of the client - or fall back to the configuration value if parsing fails.
16-
/// Assumes only a single account id is used in the application.
17-
/// </summary>
18-
public static string AwsAccountId { get; private set; }
25+
return new CanWrapResponse(instrumentedMethodInfo.RequestedWrapperName == nameof(AmazonServiceClientWrapper));
26+
}
27+
28+
public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
29+
{
30+
object client = instrumentedMethodCall.MethodCall.InvocationTarget;
1931

20-
public bool IsTransactionRequired => false;
32+
var weakReferenceKey = new WeakReferenceKey<object>(client);
33+
if (AmazonServiceClientInstanceCache.Contains(weakReferenceKey)) // don't do anything if we've already seen this client instance
34+
return Delegates.NoOp;
2135

22-
public CanWrapResponse CanWrap(InstrumentedMethodInfo instrumentedMethodInfo)
36+
AmazonServiceClientInstanceCache.Add(weakReferenceKey);
37+
38+
string awsAccountId;
39+
try
2340
{
24-
return new CanWrapResponse(instrumentedMethodInfo.RequestedWrapperName == nameof(AmazonServiceClientWrapper));
41+
// get the AWSCredentials parameter
42+
dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
43+
44+
dynamic immutableCredentials = awsCredentials.GetCredentials();
45+
string accessKey = immutableCredentials.AccessKey;
46+
47+
// convert the access key to an account id
48+
awsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
49+
}
50+
catch (Exception e)
51+
{
52+
agent.Logger.Info($"Unable to parse AWS Account ID from AccessKey. Using AccountId from configuration instead. Exception: {e.Message}");
53+
awsAccountId = agent.Configuration.AwsAccountId;
2554
}
2655

27-
public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
56+
return Delegates.GetDelegateFor(onComplete: () =>
2857
{
29-
if (AwsAccountId != null)
30-
return Delegates.NoOp;
31-
32-
try
33-
{
34-
// get the AWSCredentials parameter
35-
dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
36-
37-
dynamic immutableCredentials = awsCredentials.GetCredentials();
38-
string accessKey = immutableCredentials.AccessKey;
39-
40-
// convert the access key to an account id
41-
AwsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
42-
}
43-
catch (Exception e)
44-
{
45-
agent.Logger.Info($"Unable to parse AWS Account ID from AccessKey. Using AccountId from configuration instead. Exception: {e.Message}");
46-
AwsAccountId = agent.Configuration.AwsAccountId;
47-
}
58+
// get the _config field from the client
59+
object clientConfig = ((dynamic)client).Config;
4860

49-
return Delegates.NoOp;
50-
}
61+
// cache the account id using clientConfig as the key
62+
AwsAccountIdByClientConfigCache.Put(new WeakReferenceKey<object>(clientConfig), awsAccountId);
63+
});
5164
}
5265
}

0 commit comments

Comments
 (0)