Added new Lambda Multiple Span Flag (#178)

Author: AsakerMohd

src/AWS.Distro.OpenTelemetry.AutoInstrumentation/AwsAttributeKeys.cs

@@ -28,6 +28,7 @@ internal sealed class AwsAttributeKeys
     internal static readonly string AttributeAWSRegion = "aws.region";
     internal static readonly string AttributeAWSRequestId = "aws.requestId";
     internal static readonly string AttributeAWSTraceFlagSampled = "aws.trace.flag.sampled";
+    internal static readonly string AttributeAWSTraceLambdaFlagMultipleServer = "aws.trace.lambda.multiple-server";
 
     // The below semantic names were copied over from various sources.
     // https://github.com/open-telemetry/opentelemetry-dotnet-contrib/blob/4c6474259ccb08a41eb45ea6424243d4d2c707db/src/OpenTelemetry.Instrumentation.AWS/Implementation/AWSSemanticConventions.cs

src/AWS.Distro.OpenTelemetry.AutoInstrumentation/AwsLambdaSpanProcessor.cs

@@ -0,0 +1,44 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Diagnostics;
+using OpenTelemetry;
+using static AWS.Distro.OpenTelemetry.AutoInstrumentation.AwsAttributeKeys;
+
+namespace AWS.Distro.OpenTelemetry.AutoInstrumentation;
+
+/// <summary>
+/// Simple Span Processor that adds a new Lambda specific flag to the Lambda Handler Span
+/// This is used to show that the lambda span belongs to a trace with multiple server spans
+/// </summary>
+public class AwsLambdaSpanProcessor : BaseProcessor<Activity>
+{
+    private Activity? lambdaActivity;
+
+    /// <summary>
+    /// OnStart caches a reference to the lambda activity if it exists and adds
+    /// a flag to signify whether there are multiple server spans or not.
+    /// </summary>
+    /// <param name="activity"><see cref="Activity"/> to configure</param>
+    public override void OnStart(Activity activity)
+    {
+        if (activity.Source.Name.Equals("OpenTelemetry.Instrumentation.AWSLambda"))
+        {
+            this.lambdaActivity = activity;
+        }
+
+        if (activity.Source.Name.Equals("Microsoft.AspNetCore") && activity.ParentId != null &&
+            (activity.ParentId.Equals(this.lambdaActivity?.SpanId) || activity.ParentId.Equals(this.lambdaActivity?.Id)))
+        {
+            this.lambdaActivity.SetTag(AttributeAWSTraceLambdaFlagMultipleServer, "true");
+        }
+    }
+
+    /// <summary>
+    /// OnEnd Function
+    /// </summary>
+    /// <param name="activity"><see cref="Activity"/> to configure</param>
+    public override void OnEnd(Activity activity)
+    {
+    }
+}

src/AWS.Distro.OpenTelemetry.AutoInstrumentation/AwsSpanProcessingUtil.cs

@@ -115,17 +115,17 @@ internal static string GetIngressOperation(Activity span)
             return InternalOperation;
         }
 
-        // There is the case where we can run an ASP.NET application in a lambda environment.
-        // In this case, we need to decide which takes precedence for the aws.local.operation
-        // between functionName/FunctionHandler or the route template.
-        else if (IsLambdaEnvironment())
-        {
-            return AwsLambdaFunctionName + "/FunctionHandler";
-        }
+        // this takes precedence over the FunctionHandler path. Basically, if HttpContextWeakRef exists,
+        // this means the the span is coming from ASP.NET Core instrumentation and we want the
+        // operation name to be the API Route
         else if (span.GetCustomProperty("HttpContextWeakRef") != null)
         {
             return GetRouteTemplate(span);
         }
+        else if (IsLambdaEnvironment())
+        {
+            return AwsLambdaFunctionName + "/FunctionHandler";
+        }
 
         return RouteFallback(span);
     }

src/AWS.Distro.OpenTelemetry.AutoInstrumentation/Plugin.cs

@@ -148,20 +148,24 @@ public void TracerProviderInitialized(TracerProvider tracerProvider)
 
         // We want to be adding the exporter as the last processor in the traceProvider since processors
         // are executed in the order they were added to the provider.
-        if (AwsSpanProcessingUtil.IsLambdaEnvironment() && !this.HasCustomTracesEndpoint())
+        if (AwsSpanProcessingUtil.IsLambdaEnvironment())
         {
-            Resource processResource = tracerProvider.GetResource();
+            tracerProvider.AddProcessor(new AwsLambdaSpanProcessor());
 
-            // UDP exporter for sampled spans
-            var sampledSpanExporter = new OtlpUdpExporter(processResource, AwsXrayDaemonAddress, FormatOtelSampledTracesBinaryPrefix);
-            tracerProvider.AddProcessor(new BatchActivityExportProcessor(exporter: sampledSpanExporter, maxExportBatchSize: LambdaSpanExportBatchSize));
-
-            if (this.IsApplicationSignalsEnabled())
+            if (!this.HasCustomTracesEndpoint())
             {
-                // Register UDP Exporter to export unsampled traces in Lambda
-                // only when Application Signals enabled
-                var unsampledSpanExporter = new OtlpUdpExporter(processResource, AwsXrayDaemonAddress, FormatOtelUnSampledTracesBinaryPrefix);
-                tracerProvider.AddProcessor(new AwsBatchUnsampledSpanExportProcessor(exporter: unsampledSpanExporter, maxExportBatchSize: LambdaSpanExportBatchSize));
+                Resource processResource = tracerProvider.GetResource();
+
+                // UDP exporter for sampled spans
+                var sampledSpanExporter = new OtlpUdpExporter(processResource, AwsXrayDaemonAddress, FormatOtelSampledTracesBinaryPrefix);
+                tracerProvider.AddProcessor(new BatchActivityExportProcessor(exporter: sampledSpanExporter, maxExportBatchSize: LambdaSpanExportBatchSize));
+                if (this.IsApplicationSignalsEnabled())
+                {
+                    // Register UDP Exporter to export unsampled traces in Lambda
+                    // only when Application Signals enabled
+                    var unsampledSpanExporter = new OtlpUdpExporter(processResource, AwsXrayDaemonAddress, FormatOtelUnSampledTracesBinaryPrefix);
+                    tracerProvider.AddProcessor(new AwsBatchUnsampledSpanExportProcessor(exporter: unsampledSpanExporter, maxExportBatchSize: LambdaSpanExportBatchSize));
+                }
             }
         }
 

test/AWS.Distro.OpenTelemetry.AutoInstrumentation.Tests/AwsLambdaSpanProcessorTest.cs

@@ -0,0 +1,89 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Diagnostics;
+using System.Reflection;
+using Moq;
+using OpenTelemetry;
+using static AWS.Distro.OpenTelemetry.AutoInstrumentation.AwsAttributeKeys;
+
+namespace AWS.Distro.OpenTelemetry.AutoInstrumentation.Tests;
+
+/// <summary>
+/// AwsBatchUnsampledSpanExportProcessor test class
+/// </summary>
+[System.Diagnostics.CodeAnalysis.SuppressMessage("StyleCop.CSharp.DocumentationRules", "SA1600:Elements should be documented", Justification = "Tests")]
+public class AwsLambdaSpanProcessorTest
+{
+    private ActivitySource lambdaActivitySource = new ActivitySource("OpenTelemetry.Instrumentation.AWSLambda");
+    private ActivitySource aspNetSource = new ActivitySource("Microsoft.AspNetCore");
+    private ActivitySource httpSource = new ActivitySource("Microsoft.Http");
+
+    public AwsLambdaSpanProcessorTest()
+    {
+        var listener = new ActivityListener
+        {
+            ShouldListenTo = (activitySource) => true,
+            Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllData,
+        };
+        ActivitySource.AddActivityListener(listener);
+    }
+
+    /// <summary>
+    /// Unit test to check if multiple server spans exists, one from lambda, and
+    /// other from aspnet core and that AttributeAWSTraceLambdaFlagMultipleServer is set to true.
+    /// </summary>
+    [Fact]
+    public void CheckThatMultipleServerSpanFlagAdded()
+    {
+        using var processor = new AwsLambdaSpanProcessor();
+
+        var lambdaActivity = this.lambdaActivitySource.StartActivity("lambdaSpan", ActivityKind.Server);
+        if (lambdaActivity == null)
+        {
+            return;
+        }
+
+        Activity.Current = lambdaActivity;
+
+        var aspNetCoreActivity = this.aspNetSource.StartActivity("aspSpan", ActivityKind.Server);
+        if (aspNetCoreActivity == null)
+        {
+            return;
+        }
+
+        processor.OnStart(lambdaActivity);
+        processor.OnStart(aspNetCoreActivity);
+        processor.Shutdown();
+
+        Assert.Equal("true", lambdaActivity.GetTagItem(AttributeAWSTraceLambdaFlagMultipleServer));
+    }
+
+    /// <summary>
+    /// Unit test to check if only lambda span exists, the flag is not added
+    /// </summary>
+    [Fact]
+    public void CheckThatMultipleServerSpanFlagNotAdded()
+    {
+        using var processor = new AwsLambdaSpanProcessor();
+
+        var lambdaActivity = this.lambdaActivitySource.StartActivity("lambdaSpan", ActivityKind.Server);
+        if (lambdaActivity == null)
+        {
+            return;
+        }
+
+        Activity.Current = lambdaActivity;
+
+        var httpActivity = this.aspNetSource.StartActivity("httpSpan", ActivityKind.Client);
+        if (httpActivity == null)
+        {
+            return;
+        }
+
+        processor.OnStart(lambdaActivity);
+        processor.Shutdown();
+
+        Assert.Null(lambdaActivity.GetTagItem(AttributeAWSTraceLambdaFlagMultipleServer));
+    }
+}