[Serverless] [Fix PR-4535] EndInvocationAsync should be awaited. (#4752)

* EndInvocationAsync should be awaited. * add @DataDog/serverless-apm in the CODEOWNERS for the new lambda instrumentation * reduce the expected InitTime metrics by 1 since serverless removed
DataDog · Oct 25, 2023 · 7db4b8a · 7db4b8a
1 parent 9f5022c
commit 7db4b8a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -47,9 +47,13 @@ fault_tolerant_*.h                         @DataDog/debugger-dotnet
 /tracer/src/Datadog.Trace/FaultTolerant/   @DataDog/debugger-dotnet
 
 # Serverless
+/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/    @DataDog/tracing-dotnet  @DataDog/serverless-apm
+/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/    @DataDog/tracing-dotnet  @DataDog/serverless-apm
 /tracer/src/Datadog.Trace/ClrProfiler/ServerlessInstrumentation/   @DataDog/tracing-dotnet  @DataDog/serverless-apm
 /tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AWS/       @DataDog/tracing-dotnet  @DataDog/serverless-apm
 /tracer/test/test-applications/azure-functions/                    @DataDog/tracing-dotnet  @DataDog/serverless-apm
+/tracer/test/test-applications/integrations/Samples.AWS.Lambda/    @DataDog/tracing-dotnet  @DataDog/serverless-apm
+/tracer/test/test-applications/integrations/Samples.Amazon.Lambda.RuntimeSupport/        @DataDog/tracing-dotnet  @DataDog/serverless-apm
 docker-compose.serverless.yml                                      @DataDog/tracing-dotnet  @DataDog/serverless-apm
 
 # Shared code we could move to the root folder

diff --git a/...g.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs b/...g.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs
@@ -99,7 +99,7 @@ public TReturn OnDelegateEnd<TReturn>(object sender, TReturn returnValue, Except
         }
 
         /// <inheritdoc/>
-        public Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
+        public async Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
         {
             LambdaCommon.Log("DelegateWrapper Running OnDelegateEndAsync");
             try
@@ -108,22 +108,22 @@ public Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInner
                 if (proxyInstance == null)
                 {
                     LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false);
-                    LambdaCommon.EndInvocationAsync(string.Empty, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
+                    await LambdaCommon.EndInvocationAsync(string.Empty, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
                 }
                 else
                 {
                     var jsonString = ConvertPayloadStream(proxyInstance.OutputStream);
-                    LambdaCommon.EndInvocationAsync(jsonString, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
+                    await LambdaCommon.EndInvocationAsync(jsonString, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
                 }
             }
             catch (Exception ex)
             {
                 LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false);
-                LambdaCommon.EndInvocationAsync(string.Empty, ex, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
+                await LambdaCommon.EndInvocationAsync(string.Empty, ex, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
             }
 
             LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync");
-            return Task.FromResult(returnValue);
+            return returnValue;
         }
     }
 }

diff --git a/tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/TelemetryTests.cs b/tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/TelemetryTests.cs
@@ -237,7 +237,7 @@ public async Task Telemetry_SendsMetrics()
             }
 
             // The numbers here may change, but we should have _some_
-            telemetry.GetDistributions(DistributionShared.InitTime.GetName()).Sum(x => x.Points.Count).Should().BeGreaterThan(5);
+            telemetry.GetDistributions(DistributionShared.InitTime.GetName()).Sum(x => x.Points.Count).Should().BeGreaterThan(4);
 
             telemetry.GetMetricDataPoints(Count.TraceChunkEnqueued.GetName()).Sum(x => x.Value).Should().Be(ExpectedTraces);