#1034: added sqs messages contexts as activity links

open-telemetry · Mar 27, 2023 · e88edbe · e88edbe
1 parent bdb0077
commit e88edbe
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 14 deletions.
diff --git a/src/OpenTelemetry.Instrumentation.AWSLambda/AWSLambdaWrapper.cs b/src/OpenTelemetry.Instrumentation.AWSLambda/AWSLambdaWrapper.cs
@@ -15,6 +15,7 @@
 // </copyright>
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
@@ -170,9 +171,10 @@ public static Task<TResult> TraceAsync<TInput, TResult>(
 
     internal static Activity OnFunctionStart<TInput>(TInput input, ILambdaContext context, ActivityContext parentContext = default)
     {
+        IEnumerable<ActivityLink> links = null;
         if (parentContext == default)
         {
-            parentContext = AWSLambdaUtils.ExtractParentContext(input);
+            (parentContext, links) = AWSLambdaUtils.ExtractParentContext(input);
             if (parentContext == default && !DisableAwsXRayContextExtraction)
             {
                 parentContext = AWSLambdaUtils.GetXRayParentContext();
@@ -184,7 +186,7 @@ internal static Activity OnFunctionStart<TInput>(TInput input, ILambdaContext co
 
         // We assume that functionTags and httpTags have no intersection.
         var activityName = AWSLambdaUtils.GetFunctionName(context) ?? "AWS Lambda Invoke";
-        var activity = AWSLambdaActivitySource.StartActivity(activityName, ActivityKind.Server, parentContext, functionTags.Concat(httpTags));
+        var activity = AWSLambdaActivitySource.StartActivity(activityName, ActivityKind.Server, parentContext, functionTags.Concat(httpTags), links);
 
         return activity;
     }

diff --git a/src/OpenTelemetry.Instrumentation.AWSLambda/Implementation/AWSLambdaUtils.cs b/src/OpenTelemetry.Instrumentation.AWSLambda/Implementation/AWSLambdaUtils.cs
@@ -65,32 +65,33 @@ internal static ActivityContext GetXRayParentContext()
         return activityContext;
     }
 
-    internal static ActivityContext ExtractParentContext<TInput>(TInput input)
+    internal static (ActivityContext ParentContext, IEnumerable<ActivityLink> Links) ExtractParentContext<TInput>(TInput input)
     {
-        PropagationContext propagationContext = default;
+        PropagationContext parentContext = default;
+        IEnumerable<ActivityLink> links = null;
         switch (input)
         {
             case APIGatewayProxyRequest apiGatewayProxyRequest:
-                propagationContext = Propagators.DefaultTextMapPropagator.Extract(default, apiGatewayProxyRequest, GetHeaderValues);
+                parentContext = Propagators.DefaultTextMapPropagator.Extract(default, apiGatewayProxyRequest, GetHeaderValues);
                 break;
             case APIGatewayHttpApiV2ProxyRequest apiGatewayHttpApiV2ProxyRequest:
-                propagationContext = Propagators.DefaultTextMapPropagator.Extract(default, apiGatewayHttpApiV2ProxyRequest, GetHeaderValues);
+                parentContext = Propagators.DefaultTextMapPropagator.Extract(default, apiGatewayHttpApiV2ProxyRequest, GetHeaderValues);
                 break;
             case SQSEvent sqsEvent:
-                propagationContext = AWSMessagingUtils.ExtractParentContext(sqsEvent);
+                (parentContext, links) = AWSMessagingUtils.ExtractParentContext(sqsEvent);
                 break;
             case SQSEvent.SQSMessage sqsMessage:
-                propagationContext = AWSMessagingUtils.ExtractParentContext(sqsMessage);
+                parentContext = AWSMessagingUtils.ExtractParentContext(sqsMessage);
                 break;
             case SNSEvent snsEvent:
-                propagationContext = AWSMessagingUtils.ExtractParentContext(snsEvent);
+                parentContext = AWSMessagingUtils.ExtractParentContext(snsEvent);
                 break;
             case SNSEvent.SNSRecord snsRecord:
-                propagationContext = AWSMessagingUtils.ExtractParentContext(snsRecord);
+                parentContext = AWSMessagingUtils.ExtractParentContext(snsRecord);
                 break;
         }
 
-        return propagationContext.ActivityContext;
+        return (parentContext.ActivityContext, links);
     }
 
     internal static string GetCloudProvider()

diff --git a/src/OpenTelemetry.Instrumentation.AWSLambda/Implementation/AWSMessagingUtils.cs b/src/OpenTelemetry.Instrumentation.AWSLambda/Implementation/AWSMessagingUtils.cs
@@ -16,6 +16,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using Amazon.Lambda.SNSEvents;
 using Amazon.Lambda.SQSEvents;
@@ -31,11 +32,28 @@ internal class AWSMessagingUtils
     private const string SnsAttributeTypeStringArray = "String.Array";
     private const string SnsMessageAttributes = "MessageAttributes";
 
-    internal static PropagationContext ExtractParentContext(SQSEvent sqsEvent)
+    internal static (PropagationContext ParentContext, IEnumerable<ActivityLink> Links) ExtractParentContext(SQSEvent sqsEvent)
     {
+        if (sqsEvent?.Records == null)
+        {
+            return (default, null);
+        }
+
         // We assume there can be only one parent that's why we consider only a single (the last) record as the carrier.
-        var message = sqsEvent?.Records?.LastOrDefault();
-        return ExtractParentContext(message);
+        var parentRecord = sqsEvent.Records.LastOrDefault();
+        var parentContext = ExtractParentContext(parentRecord);
+
+        var links = new List<ActivityLink>();
+        foreach (var record in sqsEvent.Records)
+        {
+            var context = ExtractParentContext(record);
+            if (context != default)
+            {
+                links.Add(new ActivityLink(context.ActivityContext));
+            }
+        }
+
+        return (parentContext, links);
     }
 
     internal static PropagationContext ExtractParentContext(SQSEvent.SQSMessage sqsMessage)