[Serverless] Fix not able to disable tracing when running in AWS Lambda

.azure-pipelines/ultimate-pipeline.yml

@@ -2711,7 +2711,9 @@ stages:
           serverless-lambda-generic-abstract-async \
           serverless-lambda-generic-complex \
           serverless-lambda-generic-complex-nested \
-          serverless-lambda-toplevel-statement
+          serverless-lambda-toplevel-statement \
+          serverless-lambda-disabled-no-param-sync \
+          serverless-lambda-disabled-toplevel-statement
       env:
         baseImage: $(baseImage)
         framework: $(publishTargetFramework)

docker-compose.serverless.yml

@@ -41,6 +41,8 @@ services:
     - serverless-lambda-generic-complex
     - serverless-lambda-generic-complex-nested
     - serverless-lambda-toplevel-statement
+    - serverless-lambda-disabled-no-param-sync
+    - serverless-lambda-disabled-toplevel-statement
     - serverless-dummy-api
     environment:
     - AWS_LAMBDA_ENDPOINT_NO_PARAM_SYNC=http://serverless-lambda-no-param-sync:8080
@@ -81,6 +83,8 @@ services:
     - AWS_LAMBDA_ENDPOINT_GENERICBASE_COMPLEX=http://serverless-lambda-generic-complex:8080
     - AWS_LAMBDA_ENDPOINT_GENERICBASE_COMPLEX_NESTED=http://serverless-lambda-generic-complex-nested:8080
     - AWS_LAMBDA_ENDPOINT_TOPLEVEL_STATEMENT=http://serverless-lambda-toplevel-statement:8080
+    - AWS_LAMBDA_ENDPOINT_DISABLED_NO_PARAM_SYNC=http://serverless-lambda-disabled-no-param-sync:8080
+    - AWS_LAMBDA_ENDPOINT_DISABLED_TOPLEVEL_STATEMENT=http://serverless-lambda-disabled-toplevel-statement:8080
     - DUMMY_API_HOST=http://serverless-dummy-api:9005
 
   StartDependencies.Serverless:
@@ -124,6 +128,8 @@ services:
     - serverless-lambda-generic-complex
     - serverless-lambda-generic-complex-nested
     - serverless-lambda-toplevel-statement
+    - serverless-lambda-disabled-no-param-sync
+    - serverless-lambda-disabled-toplevel-statement
     - serverless-dummy-api
     environment:
       - TIMEOUT_LENGTH=120
@@ -166,6 +172,8 @@ services:
       serverless-lambda-generic-complex:8080
       serverless-lambda-generic-complex-nested:8080
       serverless-lambda-toplevel-statement:8080
+      serverless-lambda-disabled-no-param-sync:8080
+      serverless-lambda-disabled-toplevel-statement:8080
       serverless-dummy-api:9005
 
   serverless-lambda-no-param-sync:
@@ -1042,6 +1050,55 @@ services:
       - ./:/project
       - ./artifacts/build_data/logs/serverless-lambda-toplevel-statement:/var/log/datadog/dotnet
 
+  serverless-lambda-disabled-no-param-sync:
+    build:
+      context: ./artifacts/bin/Samples.AWS.Lambda
+      dockerfile: serverless.lambda.dockerfile
+      args:
+        - lambdaBaseImage=${lambdaBaseImage:-public.ecr.aws/lambda/dotnet:6}
+        - framework=${framework:-net6.0}
+    depends_on:
+      - serverless-dummy-api
+    command: "Samples.AWS.Lambda::Samples.AWS.Lambda.Function::HandlerNoParamSync"
+    image: dd-trace-dotnet/serverless-lambda-disabled-no-param-sync
+    environment:
+      - DD_TRACE_AGENT_URL=http://integrationtests:5002
+      - _DD_EXTENSION_ENDPOINT=http://integrationtests:9003
+      - DUMMY_API_HOST=http://serverless-dummy-api:9005
+      - DD_INSTRUMENTATION_TELEMETRY_ENABLED=0
+      - DD_TRACE_DEBUG=1
+      - DD_TRACE_AwsLambda_ENABLED=0
+    ports:
+      - "8080"
+    volumes:
+      - ./:/project
+      - ./artifacts/build_data/logs/serverless-lambda-disabled-no-param-sync:/var/log/datadog/dotnet
+
+  serverless-lambda-disabled-toplevel-statement:
+    build:
+      context: ./artifacts/bin/Samples.Amazon.Lambda.RuntimeSupport
+      dockerfile: serverless.lambda.dockerfile
+      args:
+        - lambdaBaseImage=${lambdaBaseImage:-public.ecr.aws/lambda/dotnet:6}
+        - framework=${framework:-net6.0}
+    depends_on:
+      - serverless-dummy-api
+    command: "Samples.Amazon.Lambda.RuntimeSupport"
+    image: dd-trace-dotnet/serverless-lambda-disabled-toplevel-statement
+    environment:
+      - DD_TRACE_AGENT_URL=http://integrationtests:5002
+      - _DD_EXTENSION_ENDPOINT=http://integrationtests:9003
+      - DUMMY_API_HOST=http://serverless-dummy-api:9005
+      - DD_INSTRUMENTATION_TELEMETRY_ENABLED=0
+      - DD_TRACE_ENABLED=0
+      - DD_TRACE_DEBUG=1
+      - DD_TRACE_AwsLambda_ENABLED=0
+    ports:
+      - "8080"
+    volumes:
+      - ./:/project
+      - ./artifacts/build_data/logs/serverless-lambda-disabled-toplevel-statement:/var/log/datadog/dotnet
+
   # The serverless function calls this API, which always returns 200 OK
   serverless-dummy-api:
     image: andrewlock/ok-api:latest

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs

@@ -68,6 +68,12 @@ private static string ConvertPayloadStream(Stream payloadStream)
 
         public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
         {
+            if (!Tracer.Instance.Settings.IsIntegrationEnabled(IntegrationId.AwsLambda))
+            {
+                // integration disabled, don't create a scope, skip this trace
+                return CallTargetState.GetDefault();
+            }
+
             LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin");
 
             Scope scope;
@@ -101,6 +107,14 @@ public TReturn OnDelegateEnd<TReturn>(object sender, TReturn returnValue, Except
         /// <inheritdoc/>
         public async Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
         {
+            if (!Tracer.Instance.Settings.IsIntegrationEnabled(IntegrationId.AwsLambda))
+            {
+                // integration disabled, don't create a scope, skip this trace,
+                // but we still need to make sure we flush any traces
+                await LambdaCommon.FlushTracesAsync().ConfigureAwait(false);
+                return returnValue;
+            }
+
             LambdaCommon.Log("DelegateWrapper Running OnDelegateEndAsync");
             try
             {

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs

@@ -67,7 +67,7 @@ internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, S
         }
     }
 
-    internal static async Task EndInvocationAsync(string returnValue, Exception exception, Scope scope, ILambdaExtensionRequest requestBuilder)
+    internal static async Task FlushTracesAsync()
     {
         try
         {
@@ -79,6 +79,11 @@ await Tracer.Instance.TracerManager.AgentWriter.FlushTracesAsync()
         {
             Log("Could not flush to the extension", ex, false);
         }
+    }
+
+    internal static async Task EndInvocationAsync(string returnValue, Exception exception, Scope scope, ILambdaExtensionRequest requestBuilder)
+    {
+        await FlushTracesAsync().ConfigureAwait(false);
 
         try
         {

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AWS/AwsLambdaTests.cs

@@ -56,9 +56,12 @@ public async Task SubmitsTraces()
                                   + 3 // throwing (manual only)
                                   + 3 // throwing with context
                                   + 8 // Generic types
-                                  + 1; // Toplevel Statement
+                                  + 1 // Toplevel Statement
+                                  + 2; // Disabled
 
-                var expectedSpans = requests * 2; // we manually instrument each request too
+                // We have HTTP Request + Manual span + Fake Lambda span for each request
+                // For disabled tests, we still get all the direct spans, we just lose the fake Lambda ones
+                var expectedSpans = requests * 2;
 
                 var spans = agent.WaitForSpans(expectedSpans, 15_000).ToArray();
 

tracer/test/snapshots/AwsLambdaTests.verified.txt

@@ -1978,5 +1978,45 @@
     Metrics: {
       _dd.top_level: 1.0
     }
+  },
+  {
+    TraceId: Id_153,
+    SpanId: Id_154,
+    Name: manual.HandlerNoParamSync,
+    Resource: manual.HandlerNoParamSync,
+    Service: Samples.AWS.Lambda,
+    Tags: {
+      language: dotnet,
+      runtime-id: Guid_39
+    },
+    Metrics: {
+      process_id: 0,
+      _dd.top_level: 1.0,
+      _dd.tracer_kr: 1.0,
+      _sampling_priority_v1: 1.0
+    }
+  },
+  {
+    TraceId: Id_153,
+    SpanId: Id_155,
+    Name: http.request,
+    Resource: GET serverless-dummy-api:9005/function/HandlerNoParamSync,
+    Service: Samples.AWS.Lambda-http-client,
+    Type: http,
+    ParentId: Id_154,
+    Tags: {
+      component: WebRequest,
+      http.method: GET,
+      http.status_code: 200,
+      http.url: http://serverless-dummy-api:9005/function/HandlerNoParamSync,
+      language: dotnet,
+      out.host: serverless-dummy-api,
+      runtime-id: Guid_39,
+      span.kind: client,
+      _dd.base_service: Samples.AWS.Lambda
+    },
+    Metrics: {
+      _dd.top_level: 1.0
+    }
   }
 ]

tracer/test/test-applications/integrations/Samples.AWS.Lambda/Program.cs

@@ -34,6 +34,13 @@ private static async Task Main(string[] args)
             // in the corresponding container. Each of the lambda instances are configured to use
             // ONE of the handler methods described in the #handler region below.
 
+            // Disabled
+            Thread.Sleep(1000);
+            await Post(Environment.GetEnvironmentVariable("AWS_LAMBDA_ENDPOINT_DISABLED_NO_PARAM_SYNC"));
+            Thread.Sleep(1000);
+            await Post(Environment.GetEnvironmentVariable("AWS_LAMBDA_ENDPOINT_DISABLED_TOPLEVEL_STATEMENT"));
+
+            // Param tests
             await Post(Environment.GetEnvironmentVariable("AWS_LAMBDA_ENDPOINT_NO_PARAM_SYNC"));
             Thread.Sleep(1000);
             await Post(Environment.GetEnvironmentVariable("AWS_LAMBDA_ENDPOINT_ONE_PARAM_SYNC"));
@@ -121,9 +128,10 @@ private static async Task Main(string[] args)
             // Toplevel Statements
             Thread.Sleep(1000);
             await Post(Environment.GetEnvironmentVariable("AWS_LAMBDA_ENDPOINT_TOPLEVEL_STATEMENT"));
-            
+
             static async Task Post(string url)
             {
+                Console.WriteLine("Sending request to " + url);
                 HttpClient client = new HttpClient();
                 client.BaseAddress = new Uri(url);
                 client.DefaultRequestHeaders