[Instrumentation.AWS]: always add context propagation data to requests

open-telemetry · Dec 25, 2024 · ec7a0d8 · ec7a0d8
1 parent bfdf029
commit ec7a0d8
Showing 1 changed file with 27 additions and 14 deletions.
diff --git a/src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSTracingPipelineHandler.cs b/src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSTracingPipelineHandler.cs
@@ -46,6 +46,22 @@ public override async Task<T> InvokeAsync<T>(IExecutionContext executionContext)
         return ret;
     }
 
+    private static void AddPropagationDataToRequest(Activity activity, IRequestContext requestContext)
+    {
+        var service = requestContext.ServiceMetaData.ServiceId;
+
+        if (AWSServiceType.IsSqsService(service))
+        {
+            SqsRequestContextHelper.AddAttributes(
+                requestContext, AWSMessagingUtils.InjectIntoDictionary(new PropagationContext(activity.Context, Baggage.Current)));
+        }
+        else if (AWSServiceType.IsSnsService(service))
+        {
+            SnsRequestContextHelper.AddAttributes(
+                requestContext, AWSMessagingUtils.InjectIntoDictionary(new PropagationContext(activity.Context, Baggage.Current)));
+        }
+    }
+
 #if NET
     [System.Diagnostics.CodeAnalysis.UnconditionalSuppressMessage(
         "Trimming",
@@ -167,16 +183,6 @@ private void AddRequestSpecificInformation(Activity activity, IRequestContext re
         {
             this.awsSemanticConventions.TagBuilder.SetTagAttributeDbSystemToDynamoDb(activity);
         }
-        else if (AWSServiceType.IsSqsService(service))
-        {
-            SqsRequestContextHelper.AddAttributes(
-                requestContext, AWSMessagingUtils.InjectIntoDictionary(new PropagationContext(activity.Context, Baggage.Current)));
-        }
-        else if (AWSServiceType.IsSnsService(service))
-        {
-            SnsRequestContextHelper.AddAttributes(
-                requestContext, AWSMessagingUtils.InjectIntoDictionary(new PropagationContext(activity.Context, Baggage.Current)));
-        }
         else if (AWSServiceType.IsBedrockRuntimeService(service))
         {
             this.awsSemanticConventions.TagBuilder.SetTagAttributeGenAiSystemToBedrock(activity);
@@ -202,14 +208,21 @@ private void ProcessEndRequest(Activity? activity, IExecutionContext executionCo
 
         var currentActivity = Activity.Current;
 
-        if (currentActivity == null
-            || !currentActivity.Source.Name.StartsWith(TelemetryConstants.TelemetryScopePrefix, StringComparison.Ordinal)
-            || !currentActivity.IsAllDataRequested)
+        if (currentActivity == null)
         {
             return null;
         }
 
-        this.AddRequestSpecificInformation(currentActivity, executionContext.RequestContext);
+        if (currentActivity.IsAllDataRequested
+            && currentActivity.Source.Name.StartsWith(TelemetryConstants.TelemetryScopePrefix, StringComparison.Ordinal))
+        {
+            this.AddRequestSpecificInformation(currentActivity, executionContext.RequestContext);
+        }
+
+        // Context propagation should always happen regardless of sampling decision (which affects Activity.IsAllDataRequested and Activity.Source).
+        // Otherwise, downstream services can make inconsistent sampling decisions and create incomplete traces.
+        AddPropagationDataToRequest(currentActivity, executionContext.RequestContext);
+
         return currentActivity;
     }
 }