feat: Instrument Lambda invocations in AWS SDK

.github/workflows/all_solutions.yml

@@ -270,6 +270,7 @@ jobs:
           AwsLambda.Sns,
           AwsLambda.Sqs,
           AwsLambda.WebRequest,
+          AwsSdk,
           AzureFunction,
           BasicInstrumentation, 
           CatInbound, 

src/Agent/NewRelic/Agent/Core/AgentHealth/AgentHealthReporter.cs

@@ -685,6 +685,7 @@ private void CollectOneTimeMetrics()
             ReportIfLoggingDisabled();
             ReportIfInstrumentationIsDisabled();
             ReportIfGCSamplerV2IsEnabled();
+            ReportIfAwsAccountIdProvided();
         }
 
         public void CollectMetrics()
@@ -847,8 +848,14 @@ private void ReportIfGCSamplerV2IsEnabled()
             {
                 ReportSupportabilityCountMetric(MetricNames.SupportabilityGCSamplerV2Enabled);
             }
-
         }
 
+        private void ReportIfAwsAccountIdProvided()
+        {
+            if (!string.IsNullOrEmpty(_configuration.AwsAccountId))
+            {
+                ReportSupportabilityCountMetric(MetricNames.SupportabilityAwsAccountIdProvided);
+            }
+        }
     }
 }

src/Agent/NewRelic/Agent/Core/Attributes/AttributeDefinitionService.cs

@@ -126,6 +126,7 @@ public interface IAttributeDefinitions
 
         AttributeDefinition<object, object> GetLambdaAttribute(string name);
         AttributeDefinition<object, object> GetFaasAttribute(string name);
+        AttributeDefinition<object, object> GetCloudSdkAttribute(string name);
 
         AttributeDefinition<string, string> GetRequestParameterAttribute(string paramName);
 
@@ -190,6 +191,7 @@ public AttributeDefinitions(IAttributeFilter attribFilter)
         private readonly ConcurrentDictionary<string, AttributeDefinition<string, string>> _requestHeadersAttributes = new ConcurrentDictionary<string, AttributeDefinition<string, string>>();
         private readonly ConcurrentDictionary<string, AttributeDefinition<object, object>> _lambdaAttributes = new ConcurrentDictionary<string, AttributeDefinition<object, object>>();
         private readonly ConcurrentDictionary<string, AttributeDefinition<object, object>> _faasAttributes = new();
+        private readonly ConcurrentDictionary<string, AttributeDefinition<object, object>> _cloudSdkAttributes = new();
 
         private readonly ConcurrentDictionary<TypeAttributeValue, AttributeDefinition<TypeAttributeValue, string>> _typeAttributes = new ConcurrentDictionary<TypeAttributeValue, AttributeDefinition<TypeAttributeValue, string>>();
 
@@ -281,6 +283,19 @@ public AttributeDefinition<object, object> GetFaasAttribute(string name)
         }
 
 
+        private AttributeDefinition<object, object> CreateCloudSdkAttribute(string attribName)
+        {
+            return AttributeDefinitionBuilder
+                .Create<object, object>(attribName, AttributeClassification.AgentAttributes)
+                .AppliesTo(AttributeDestinations.TransactionTrace)
+                .AppliesTo(AttributeDestinations.SpanEvent)
+                .Build(_attribFilter);
+        }
+
+        public AttributeDefinition<object, object> GetCloudSdkAttribute(string name)
+        {
+            return _cloudSdkAttributes.GetOrAdd(name, CreateCloudSdkAttribute);
+        }
         public AttributeDefinition<object, object> GetCustomAttributeForTransaction(string name)
         {
             return _trxCustomAttributes.GetOrAdd(name, CreateCustomAttributeForTransaction);

src/Agent/NewRelic/Agent/Core/Metrics/MetricNames.cs

@@ -838,6 +838,7 @@ public static string GetSupportabilityInstallType(string installType)
 
         public const string SupportabilityIgnoredInstrumentation = SupportabilityDotnetPs + "IgnoredInstrumentation";
         public const string SupportabilityGCSamplerV2Enabled = SupportabilityDotnetPs + "GCSamplerV2/Enabled";
+        public const string SupportabilityAwsAccountIdProvided = SupportabilityDotnetPs + "AwsAccountId/Config";
 
         #endregion Supportability
 

src/Agent/NewRelic/Agent/Core/Transactions/NoOpTransaction.cs

@@ -331,5 +331,10 @@ public void AddFaasAttribute(string name, object value)
         {
             return;
         }
+
+        public void AddCloudSdkAttribute(string name, object value)
+        {
+            return;
+        }
     }
 }

src/Agent/NewRelic/Agent/Core/Transactions/Transaction.cs

@@ -1374,7 +1374,7 @@ public void AddLambdaAttribute(string name, object value)
         {
             if (string.IsNullOrWhiteSpace(name))
             {
-                Log.Debug($"AddLambdaAttribute - Unable to set Lambda value on transaction because the key is null/empty");
+                Log.Debug($"AddLambdaAttribute - Name cannot be null/empty");
                 return;
             }
 
@@ -1386,12 +1386,24 @@ public void AddFaasAttribute(string name, object value)
         {
             if (string.IsNullOrWhiteSpace(name))
             {
-                Log.Debug($"AddFaasAttribute - Unable to set FaaS value on transaction because the key is null/empty");
+                Log.Debug($"AddFaasAttribute - Name cannot be null/empty");
                 return;
             }
 
             var faasAttrib = _attribDefs.GetFaasAttribute(name);
             TransactionMetadata.UserAndRequestAttributes.TrySetValue(faasAttrib, value);
         }
+
+        public void AddCloudSdkAttribute(string name, object value)
+        {
+            if (string.IsNullOrWhiteSpace(name))
+            {
+                Log.Debug($"AddCloudSdkAttribute - Name cannot be null/empty");
+                return;
+            }
+
+            var cloudAttrib = _attribDefs.GetCloudSdkAttribute(name);
+            TransactionMetadata.UserAndRequestAttributes.TrySetValue(cloudAttrib, value);
+        }
     }
 }

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Api/ITransaction.cs

@@ -315,5 +315,7 @@ ISegment StartMessageBrokerSegment(MethodCall methodCall, MessageBrokerDestinati
         void AddLambdaAttribute(string name, object value);
 
         void AddFaasAttribute(string name, object value);
+
+        void AddCloudSdkAttribute(string name, object value);
     }
 }

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/AwsSdk/ArnBuilder.cs

@@ -0,0 +1,160 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Linq;
+using System.Text.RegularExpressions;
+using NewRelic.Agent.Extensions.Collections;
+
+namespace NewRelic.Agent.Extensions.AwsSdk
+{
+    public class ArnBuilder
+    {
+        public readonly string Partition;
+        public readonly string Region;
+        public readonly string AccountId;
+
+        public ArnBuilder(string partition, string region, string accountId)
+        {
+            Partition = partition ?? "";
+            Region = region ?? "";
+            AccountId = accountId ?? "";
+        }
+
+        public string Build(string service, string resource) => ConstructArn(Partition, service, Region, AccountId, resource);
+
+        // This is the full regex pattern for a Lambda ARN:
+        // (arn:(aws[a-zA-Z-]*)?:lambda:)?([a-z]{2}((-gov)|(-iso([a-z]?)))?-[a-z]+-\d{1}:)?(\d{12}:)?(function:)?([a-zA-Z0-9-_\.]+)(:(\$LATEST|[a-zA-Z0-9-_]+))?
+
+        // If it's a full ARN, it has to start with 'arn:'
+        // A partial ARN can contain up to 5 segments separated by ':'
+        // 1. Region
+        // 2. Account ID
+        // 3. 'function' (fixed string)
+        // 4. Function name
+        // 5. Alias or version
+        // Only the function name is required, the rest are all optional. e.g. you could have region and function name and nothing else
+        public string BuildFromPartialLambdaArn(string invocationName)
+        {
+            if (invocationName.StartsWith("arn:"))
+            {
+                return invocationName;
+            }
+            var segments = invocationName.Split(':');
+            string functionName = null;
+            string alias = null;
+            string fallback = null;
+            string region = null;
+            string accountId = null;
+
+            // If there's only one segment, assume it's the function name
+            if (segments.Length == 1)
+            {
+                functionName = segments[0];
+            }
+            else
+            {
+                // All we should need is the function name, but if we find a region or account ID, we'll use it
+                // since it should be more accurate
+                foreach (var segment in segments)
+                {
+                    // A string that looks like a region or account ID could also be the function name
+                    // Assume it's the former, unless we never find a function name
+                    if (LooksLikeARegion(segment))
+                    {
+                        if (string.IsNullOrEmpty(region))
+                        {
+                            region = segment;
+                        }
+                        else
+                        {
+                            fallback = segment;
+                        }
+                        continue;
+                    }
+                    else if (LooksLikeAnAccountId(segment))
+                    {
+                        if (string.IsNullOrEmpty(accountId))
+                        {
+                            accountId = segment;
+                        }
+                        else
+                        {
+                            fallback = segment;
+                        }
+                        continue;
+                    }
+                    else if (segment == "function")
+                    {
+                        continue;
+                    }
+                    else if (functionName == null)
+                    {
+                        functionName = segment;
+                    }
+                    else if (alias == null)
+                    {
+                        alias = segment;
+                    }
+                    else
+                    {
+                        return null;
+                    }
+                }
+            }
+
+            if (string.IsNullOrEmpty(functionName))
+            {
+                if (!string.IsNullOrEmpty(fallback))
+                {
+                    functionName = fallback;
+                }
+                else
+                {
+                    return null;
+                }
+            }
+
+            accountId = !string.IsNullOrEmpty(accountId) ? accountId : AccountId;
+            if (string.IsNullOrEmpty(accountId))
+            {
+                return null;
+            }
+
+            region = !string.IsNullOrEmpty(region) ? region : Region;
+            if (string.IsNullOrEmpty(region))
+            {
+                return null;
+            }
+
+            if (!string.IsNullOrEmpty(alias))
+            {
+                functionName += $":{alias}";
+            }
+            return ConstructArn(Partition, "lambda", region, accountId, $"function:{functionName}");
+        }
+
+        public override string ToString()
+        {
+            string partition = string.IsNullOrEmpty(Partition) ? "[Missing]" : Partition;
+            string region = string.IsNullOrEmpty(Region) ? "[Missing]" : Region;
+            string idPresent = string.IsNullOrEmpty(AccountId) ? "[Missing]" : "[Present]";
+
+            return $"Partition: {partition}, Region: {region}, AccountId: {idPresent}";
+        }
+
+        private static Regex RegionRegex = new Regex(@"^[a-z]{2}((-gov)|(-iso([a-z]?)))?-[a-z]+-\d{1}$", RegexOptions.Compiled);
+        private static bool LooksLikeARegion(string text) => RegionRegex.IsMatch(text);
+        private static bool LooksLikeAnAccountId(string text) => (text.Length == 12) && text.All(c => c >= '0' && c <= '9');
+
+        private string ConstructArn(string partition, string service, string region, string accountId, string resource)
+        {
+            if (string.IsNullOrEmpty(partition) || string.IsNullOrEmpty(region) || string.IsNullOrEmpty(accountId)
+                || string.IsNullOrEmpty(service) || string.IsNullOrEmpty(resource))
+            {
+                return null;
+            }
+            return "arn:" + partition + ":" + service + ":" + region + ":" + accountId + ":" + resource;
+        }
+
+    }
+}

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Collections/ConcurrentHashSet.cs

@@ -47,7 +47,13 @@ public void Add(T item)
                 _hashSet.Add(item);
             }
         }
-
+        public bool TryAdd(T item)
+        {
+            using (_writeLock())
+            {
+                return _hashSet.Add(item);
+            }
+        }
         public void Clear()
         {
             using (_writeLock())

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/HandlerMethodWrapper.cs

@@ -193,10 +193,9 @@ private void InitLambdaData(InstrumentedMethodCall instrumentedMethodCall, IAgen
                 {
                     agent.Logger.Log(Agent.Extensions.Logging.Level.Debug, $"Supported Event Type found: {_functionDetails.EventType}");
                 }
-                else if (!_unsupportedInputTypes.Contains(name))
+                else if (_unsupportedInputTypes.TryAdd(name))
                 {
                     agent.Logger.Log(Agent.Extensions.Logging.Level.Warn, $"Unsupported input object type: {name}. Unable to provide additional instrumentation.");
-                    _unsupportedInputTypes.Add(name);
                 }
             }
 
@@ -344,10 +343,9 @@ private void CaptureResponseData(ITransaction transaction, object response, IAge
                 ||
                 (_functionDetails.EventType == AwsLambdaEventType.ApplicationLoadBalancerRequest && responseType != "Amazon.Lambda.ApplicationLoadBalancerEvents.ApplicationLoadBalancerResponse"))
             {
-                if (!_unexpectedResponseTypes.Contains(responseType))
+                if (_unexpectedResponseTypes.TryAdd(responseType))
                 {
                     agent.Logger.Log(Agent.Extensions.Logging.Level.Warn, $"Unexpected response type {responseType} for request event type {_functionDetails.EventType}. Not capturing any response data.");
-                    _unexpectedResponseTypes.Add(responseType);
                 }
 
                 return;