fix: Resolve Lambda auto-instrumentation race condition

build/Dotty/packageInfo.json

@@ -19,6 +19,9 @@
     {
         "packageName": "amazon.lambda.kinesisfirehoseevents"
     },
+    {
+        "packageName": "amazon.lambda.runtimesupport"
+    },
     {
         "packageName": "amazon.lambda.s3events"
     },

build/Dotty/projectInfo.json

@@ -2,6 +2,9 @@
     {
         "projectFile": "tests/Agent/IntegrationTests/SharedApplications/Common/MFALatestPackages/MFALatestPackages.csproj"
     },
+    {
+        "projectFile": "tests/Agent/IntegrationTests/Applications/AspNetCoreWebApiLambdaApplication/AspNetCoreWebApiLambdaApplication.csproj"
+    },
     {
         "projectFile": "tests/Agent/IntegrationTests/Applications/AzureFunctionApplication/AzureFunctionApplication.csproj"
     },

src/Agent/NewRelic/Agent/Core/Utilities/ExtensionsLoader.cs

@@ -86,6 +86,13 @@ public static void Initialize(string installPathExtensionsDirectory)
             { "KafkaProducerWrapper",                                                                          Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
             { "KafkaSerializerWrapper",                                                                        Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
             { "KafkaConsumerWrapper",                                                                          Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
+
+            // NewRelic.Providers.Wrapper.AwsLambda.dll depends on Amazon.Lambda.RuntimeSupport; therefore, it
+            // should only be loaded by the agent when running in a Lambda environment, otherwise an assembly
+            // load exception will occur.
+            { "NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper",                          Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
+            { "NewRelic.Providers.Wrapper.AwsLambda.HandlerMethod",                                            Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
+            { "OpenTracingWrapper",                                                                            Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
         };
 
         var nonAutoReflectedAssemblies = _dynamicLoadWrapperAssemblies.Values.Distinct().ToList();

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/AwsLambda.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <TargetFramework>netstandard2.0</TargetFramework>
     <AssemblyName>NewRelic.Providers.Wrapper.AwsLambda</AssemblyName>
@@ -11,6 +11,7 @@
 		</Content>
 	</ItemGroup>
   <ItemGroup>
+    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.7.0" />
     <PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
   </ItemGroup>
   <ItemGroup>

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/AwsLambdaWrapperExtensions.cs

@@ -0,0 +1,11 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using NewRelic.Agent.Configuration;
+
+namespace NewRelic.Providers.Wrapper.AwsLambda;
+
+internal static class AwsLambdaWrapperExtensions
+{
+    public static string GetTransactionCategory(IConfiguration configuration) => configuration.AwsLambdaApmModeEnabled ? "Function" : "Lambda";
+}

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/EnsureTransformCompletesWrapper.cs

@@ -0,0 +1,64 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Threading.Tasks;
+using Amazon.Lambda.RuntimeSupport;
+using NewRelic.Agent.Api;
+using NewRelic.Agent.Extensions.Providers.Wrapper;
+using NewRelic.Reflection;
+
+namespace NewRelic.Providers.Wrapper.AwsLambda;
+
+public class EnsureTransformCompletesWrapper : IWrapper
+{
+    public bool IsTransactionRequired => false;
+
+    public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
+    {
+        return new CanWrapResponse("NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper".Equals(methodInfo.RequestedWrapperName));
+    }
+
+    public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
+    {
+        // The instrumented method is only called once in the lifetime of a lambda instance so we can just create a handler wrapper each time.
+        CreateAndSetWrappedHandlerMethod(agent, transaction, instrumentedMethodCall.MethodCall.InvocationTarget);
+
+        return Delegates.NoOp;
+    }
+
+    private static void CreateAndSetWrappedHandlerMethod(IAgent agent, ITransaction transaction, object bootstrapper)
+    {
+        var bootstrapperType = bootstrapper.GetType();
+
+        var handlerReadAccessor = VisibilityBypasser.Instance.GenerateFieldReadAccessor<object>(bootstrapperType, "_handler");
+        var originalHandler = handlerReadAccessor(bootstrapper);
+
+        // replace the original handler with NewHandler, which calls the original handler and waits for the transaction to finish before returning a response to the lambda runtime.
+        var handlerWriteAccessor = VisibilityBypasser.Instance.GenerateFieldWriteAccessor<LambdaBootstrapHandler>(bootstrapperType, "_handler");
+        handlerWriteAccessor(bootstrapper, NewHandler);
+        return;
+
+        // Creates a _handler that will wait for the transaction to finish being transformed and harvested before allowing
+        // the lambda runtime to send a response back. This prevents a race condition where async continuations can run in different orders
+        // which sometimes causes the lambda runtime to return a response before the agent can generate a lambda payload.
+        // ref: https://github.com/aws/aws-lambda-dotnet/blob/c28fcfaba68607f785662ff1d232eb9b26d0fa09/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs#L369
+        // and https://github.com/aws/aws-lambda-dotnet/blob/c28fcfaba68607f785662ff1d232eb9b26d0fa09/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs#L387
+        async Task<InvocationResponse> NewHandler(InvocationRequest request)
+        {
+            transaction = agent.CreateTransaction(isWeb: false, category: AwsLambdaWrapperExtensions.GetTransactionCategory(agent.Configuration), transactionDisplayName: "TempLambdaName", doNotTrackAsUnitOfWork: true);
+            agent.Logger.Finest("EnsureTransformCompletesWrapper: Started transaction.");
+
+            try
+            {
+                agent.Logger.Finest("EnsureTransformCompletesWrapper: Calling original handler.");
+                var typedHandler = (LambdaBootstrapHandler)originalHandler;
+                return await typedHandler(request);
+            }
+            finally
+            {
+                agent.Logger.Finest("EnsureTransformCompletesWrapper: Ending transaction.");
+                transaction.End();
+            }
+        }
+    }
+}

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/HandlerMethodWrapper.cs

@@ -123,23 +123,23 @@ public bool ValidateWebRequestParameters(InstrumentedMethodCall instrumentedMeth
                 switch (EventType)
                 {
                     case AwsLambdaEventType.APIGatewayHttpApiV2ProxyRequest:
-                    {
-                        if (input.RequestContext != null)
                         {
-                            dynamic requestContext = input.RequestContext;
+                            if (input.RequestContext != null)
+                            {
+                                dynamic requestContext = input.RequestContext;
 
-                            if (requestContext.Http != null)
-                                return !string.IsNullOrEmpty(requestContext.Http.Method) && !string.IsNullOrEmpty(requestContext.Http.Path);
-                        }
+                                if (requestContext.Http != null)
+                                    return !string.IsNullOrEmpty(requestContext.Http.Method) && !string.IsNullOrEmpty(requestContext.Http.Path);
+                            }
 
-                        return false;
-                    }
+                            return false;
+                        }
                     case AwsLambdaEventType.APIGatewayProxyRequest:
                     case AwsLambdaEventType.ApplicationLoadBalancerRequest:
-                    {
-                        dynamic webReq = input;
-                        return !string.IsNullOrEmpty(webReq.HttpMethod) && !string.IsNullOrEmpty(webReq.Path);
-                    }
+                        {
+                            dynamic webReq = input;
+                            return !string.IsNullOrEmpty(webReq.HttpMethod) && !string.IsNullOrEmpty(webReq.Path);
+                        }
                     default:
                         return true;
                 }
@@ -156,7 +156,7 @@ public bool ValidateWebRequestParameters(InstrumentedMethodCall instrumentedMeth
     private static object _initLock = new object();
     private static FunctionDetails _functionDetails = null;
 
-    public bool IsTransactionRequired => false;
+    public bool IsTransactionRequired => true;
 
     private static bool _coldStart = true;
     private ConcurrentHashSet<string> _unexpectedResponseTypes = new();
@@ -246,13 +246,18 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
 
         // create a transaction for the function invocation if AwsLambdaApmModeEnabled is enabled then based on spec WebTransaction/Function/APIGATEWAY myFunction
         // else WebTransaction/Function/myFunction
-        transaction = agent.CreateTransaction(
-            isWeb: _functionDetails.EventType.IsWebEvent(),
-            category: agent.Configuration.AwsLambdaApmModeEnabled ? "Function" : "Lambda",
-            transactionDisplayName: agent.Configuration.AwsLambdaApmModeEnabled
-                ? _functionDetails.EventType.ToEventTypeString().ToUpper() + " " + _functionDetails.FunctionName
-                : _functionDetails.FunctionName,
-            doNotTrackAsUnitOfWork: true);
+        var transactionName  = agent.Configuration.AwsLambdaApmModeEnabled
+            ? _functionDetails.EventType.ToEventTypeString().ToUpper() + " " + _functionDetails.FunctionName
+            : _functionDetails.FunctionName;
+
+        if (_functionDetails.IsWebRequest)
+        {
+            transaction.SetWebTransactionName("Lambda", transactionName, TransactionNamePriority.Handler); // low priority to allow for override 
+        }
+        else
+        {
+            transaction.SetOtherTransactionName(AwsLambdaWrapperExtensions.GetTransactionCategory(agent.Configuration), transactionName, TransactionNamePriority.Handler); // low priority to allow for override 
+        }
 
         if (isAsync)
         {
@@ -311,7 +316,6 @@ void InvokeTryProcessResponse(Task responseTask)
                 finally
                 {
                     segment?.End();
-                    transaction.End();
                 }
             }
         }
@@ -329,7 +333,6 @@ void InvokeTryProcessResponse(Task responseTask)
                 onFailure: exception =>
                 {
                     segment.End(exception);
-                    transaction.End();
                 });
         }
     }
@@ -381,4 +384,4 @@ private static bool ValidTaskResponse(Task response)
         return response?.Status == TaskStatus.RanToCompletion;
     }
 
-}
+}

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsLambda/Instrumentation.xml

@@ -6,7 +6,7 @@ SPDX-License-Identifier: Apache-2.0
 <extension xmlns="urn:newrelic-extension">
   <instrumentation>
     <!-- Start the agent init process during the init stage of the lambda for better performance and reduced cost. -->
-    <tracerFactory name="NewRelic.Agent.Core.Wrapper.NoOpWrapper">
+    <tracerFactory name="NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper">
       <match assemblyName="Amazon.Lambda.RuntimeSupport" className="Amazon.Lambda.RuntimeSupport.LambdaBootstrap">
         <exactMethodMatcher methodName="RunAsync" />
       </match>

tests/Agent/IntegrationTests/Applications/AspNetCoreWebApiLambdaApplication/AspNetCoreWebApiLambdaApplication.csproj

@@ -1,19 +1,27 @@
-<Project Sdk="Microsoft.NET.Sdk.Web">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <OutputType>Exe</OutputType>
-    <TargetFramework>net8.0</TargetFramework>
+    <TargetFrameworks>net8.0;net10.0</TargetFrameworks>
     <ImplicitUsings>enable</ImplicitUsings>
     <GenerateRuntimeConfigurationFiles>true</GenerateRuntimeConfigurationFiles>
     <AWSProjectType>Lambda</AWSProjectType>
     <!-- This property makes the build directory similar to a publish directory and helps the AWS .NET Lambda Mock Test Tool find project dependencies. -->
     <CopyLocalLockFileAssemblies>true</CopyLocalLockFileAssemblies>
   </PropertyGroup>
+  <!-- Use Microsoft.NET.Sdk (not .Web) to avoid DeployOnBuild/Publish conflicts with
+       multi-TFM builds (NETSDK1129). FrameworkReference provides ASP.NET Core APIs. -->
   <ItemGroup>
+    <FrameworkReference Include="Microsoft.AspNetCore.App" />
+  </ItemGroup>
+  <ItemGroup>
+    <!-- 1.7.0 is from Feb 2022 -->
+    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.7.0" Condition="'$(TargetFramework)' == 'net8.0'" />
+    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.14.2" Condition="'$(TargetFramework)' == 'net10.0'" />
+
     <PackageReference Include="Amazon.Lambda.AspNetCoreServer" Version="9.0.1" />
-    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.10.0" />
     <PackageReference Include="Amazon.Lambda.Serialization.SystemTextJson" Version="2.4.3" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="..\..\ApplicationHelperLibraries\ApplicationLifecycle\ApplicationLifecycle.csproj" />
   </ItemGroup>
-</Project>
+</Project>

tests/Agent/IntegrationTests/Applications/AspNetCoreWebApiLambdaApplication/GlobalUsings.cs

@@ -0,0 +1,8 @@
+// Copyright 2020 New Relic, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+global using Microsoft.AspNetCore.Builder;
+global using Microsoft.AspNetCore.Hosting;
+global using Microsoft.AspNetCore.Http;
+global using Microsoft.Extensions.Configuration;
+global using Microsoft.Extensions.DependencyInjection;
+global using Microsoft.Extensions.Hosting;

tests/Agent/IntegrationTests/Applications/LambdaSelfExecutingAssembly/LambdaSelfExecutingAssembly.csproj

@@ -38,8 +38,12 @@
     <!-- 1.2.0 is from October 2020 -->
     <PackageReference Include="Amazon.Lambda.SQSEvents" Version="2.2.1" Condition="'$(TargetFramework)' == 'net8.0'" />
     <PackageReference Include="Amazon.Lambda.SQSEvents" Version="2.2.1" Condition="'$(TargetFramework)' == 'net10.0'" />
+
+    <!-- 1.7.0 is from Feb 2022 -->
+    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.7.0" Condition="'$(TargetFramework)' == 'net8.0'" />
+    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.14.2" Condition="'$(TargetFramework)' == 'net10.0'" />
+
     <!-- Non-event libraries -->
-    <PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.11.0" />
     <PackageReference Include="Amazon.Lambda.Serialization.SystemTextJson" Version="2.4.3" />
     <PackageReference Include="Microsoft.VisualStudio.Threading.Analyzers" Version="17.14.15">
       <PrivateAssets>all</PrivateAssets>

tests/Agent/IntegrationTests/IntegrationTests/AwsLambda/AwsLambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTest.cs

@@ -82,3 +82,11 @@ public AwsLambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTestCoreOldest
     {
     }
 }
+
+public class AwsLambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTestCoreLatest : AwsLambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTest<LambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTriggerFixtureCoreLatest>
+{
+    public AwsLambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTestCoreLatest(LambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTriggerFixtureCoreLatest fixture, ITestOutputHelper output)
+        : base(fixture, output, "WebTransaction/MVC/Values/Get")
+    {
+    }
+}

tests/Agent/IntegrationTests/IntegrationTests/AwsLambda/AwsLambdaAPIGatewayRequestAutoInstrumentationTest.cs

@@ -89,9 +89,17 @@ private void ValidateServerlessPayload(ServerlessPayload serverlessPayload)
     }
 }
 
-public class AwsLambdaAPIGatewayRequestAutoInstrumentationTestTestCoreOldest : AwsLambdaAPIGatewayRequestAutoInstrumentationTest<LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreOldest>
+public class AwsLambdaAPIGatewayRequestAutoInstrumentationTestCoreOldest : AwsLambdaAPIGatewayRequestAutoInstrumentationTest<LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreOldest>
 {
-    public AwsLambdaAPIGatewayRequestAutoInstrumentationTestTestCoreOldest(LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreOldest fixture, ITestOutputHelper output)
+    public AwsLambdaAPIGatewayRequestAutoInstrumentationTestCoreOldest(LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreOldest fixture, ITestOutputHelper output)
+        : base(fixture, output, "WebTransaction/MVC/Values/Get")
+    {
+    }
+}
+
+public class AwsLambdaAPIGatewayRequestAutoInstrumentationTestCoreLatest : AwsLambdaAPIGatewayRequestAutoInstrumentationTest<LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreLatest>
+{
+    public AwsLambdaAPIGatewayRequestAutoInstrumentationTestCoreLatest(LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreLatest fixture, ITestOutputHelper output)
         : base(fixture, output, "WebTransaction/MVC/Values/Get")
     {
     }

tests/Agent/IntegrationTests/IntegrationTests/AwsLambda/AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTest.cs

@@ -78,9 +78,17 @@ private void ValidateServerlessPayload(ServerlessPayload serverlessPayload)
     }
 }
 
-public class AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestTestCoreOldest : AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTest<LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest>
+public class AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestCoreOldest : AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTest<LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest>
 {
-    public AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestTestCoreOldest(LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest fixture, ITestOutputHelper output)
+    public AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestCoreOldest(LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest fixture, ITestOutputHelper output)
+        : base(fixture, output, "WebTransaction/MVC/Values/Get")
+    {
+    }
+}
+
+public class AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestCoreLatest : AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTest<LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest>
+{
+    public AwsLambdaApplicationLoadBalancerRequestAutoInstrumentationTestCoreLatest(LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest fixture, ITestOutputHelper output)
         : base(fixture, output, "WebTransaction/MVC/Values/Get")
     {
     }

tests/Agent/IntegrationTests/IntegrationTests/RemoteServiceFixtures/AwsLambda/LambdaAutoInstrumentationTriggerFixtures.cs

@@ -159,13 +159,22 @@ public LambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTriggerFixtureCor
         CommandLineArguments = "--handler APIGatewayHttpApiV2ProxyFunctionEntryPoint";
     }
 }
+
 public class LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest : AspNetCoreWebApiLambdaFixtureBase
 {
     public LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreOldest() : base(CoreOldestTFM)
     {
         CommandLineArguments = "--handler ApplicationLoadBalancerFunctionEntryPoint";
     }
 }
+public class LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest : AspNetCoreWebApiLambdaFixtureBase
+{
+    public LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest() : base(CoreLatestTFM)
+    {
+        CommandLineArguments = "--handler ApplicationLoadBalancerFunctionEntryPoint";
+    }
+}
+
 public class LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreLatest : AspNetCoreWebApiLambdaFixtureBase
 {
     public LambdaAPIGatewayProxyRequestAutoInstrumentationTriggerFixtureCoreLatest() : base(CoreLatestTFM)
@@ -181,10 +190,3 @@ public LambdaAPIGatewayHttpApiV2ProxyRequestAutoInstrumentationTriggerFixtureCor
         CommandLineArguments = "--handler APIGatewayHttpApiV2ProxyFunctionEntryPoint";
     }
 }
-public class LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest : AspNetCoreWebApiLambdaFixtureBase
-{
-    public LambdaApplicationLoadBalancerRequestAutoInstrumentationTriggerFixtureCoreLatest() : base(CoreLatestTFM)
-    {
-        CommandLineArguments = "--handler ApplicationLoadBalancerFunctionEntryPoint";
-    }
-}