Skip to content

Commit 25c59e7

Browse files
tippmar-nrtippmar-nr
authored
fix: Resolve Lambda auto-instrumentation race condition (#3508)
1 parent d543f42 commit 25c59e7

15 files changed

Lines changed: 179 additions & 41 deletions

‎build/Dotty/packageInfo.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
{
2020
"packageName": "amazon.lambda.kinesisfirehoseevents"
2121
},
22+
{
23+
"packageName": "amazon.lambda.runtimesupport"
24+
},
2225
{
2326
"packageName": "amazon.lambda.s3events"
2427
},

‎build/Dotty/projectInfo.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
{
33
"projectFile": "tests/Agent/IntegrationTests/SharedApplications/Common/MFALatestPackages/MFALatestPackages.csproj"
44
},
5+
{
6+
"projectFile": "tests/Agent/IntegrationTests/Applications/AspNetCoreWebApiLambdaApplication/AspNetCoreWebApiLambdaApplication.csproj"
7+
},
58
{
69
"projectFile": "tests/Agent/IntegrationTests/Applications/AzureFunctionApplication/AzureFunctionApplication.csproj"
710
},

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,13 @@ public static void Initialize(string installPathExtensionsDirectory)
8686
{ "KafkaProducerWrapper", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
8787
{ "KafkaSerializerWrapper", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
8888
{ "KafkaConsumerWrapper", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.Kafka.dll") },
89+
90+
// NewRelic.Providers.Wrapper.AwsLambda.dll depends on Amazon.Lambda.RuntimeSupport; therefore, it
91+
// should only be loaded by the agent when running in a Lambda environment, otherwise an assembly
92+
// load exception will occur.
93+
{ "NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
94+
{ "NewRelic.Providers.Wrapper.AwsLambda.HandlerMethod", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
95+
{ "OpenTracingWrapper", Path.Combine(_installPathExtensionsDirectory, "NewRelic.Providers.Wrapper.AwsLambda.dll") },
8996
};
9097

9198
var nonAutoReflectedAssemblies = _dynamicLoadWrapperAssemblies.Values.Distinct().ToList();

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
<Project Sdk="Microsoft.NET.Sdk">
1+
<Project Sdk="Microsoft.NET.Sdk">
22
<PropertyGroup>
33
<TargetFramework>netstandard2.0</TargetFramework>
44
<AssemblyName>NewRelic.Providers.Wrapper.AwsLambda</AssemblyName>
@@ -11,6 +11,7 @@
1111
</Content>
1212
</ItemGroup>
1313
<ItemGroup>
14+
<PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.7.0" />
1415
<PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
1516
</ItemGroup>
1617
<ItemGroup>

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using NewRelic.Agent.Configuration;
5+
6+
namespace NewRelic.Providers.Wrapper.AwsLambda;
7+
8+
internal static class AwsLambdaWrapperExtensions
9+
{
10+
public static string GetTransactionCategory(IConfiguration configuration) => configuration.AwsLambdaApmModeEnabled ? "Function" : "Lambda";
11+
}

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
using System.Threading.Tasks;
5+
using Amazon.Lambda.RuntimeSupport;
6+
using NewRelic.Agent.Api;
7+
using NewRelic.Agent.Extensions.Providers.Wrapper;
8+
using NewRelic.Reflection;
9+
10+
namespace NewRelic.Providers.Wrapper.AwsLambda;
11+
12+
public class EnsureTransformCompletesWrapper : IWrapper
13+
{
14+
public bool IsTransactionRequired => false;
15+
16+
public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
17+
{
18+
return new CanWrapResponse("NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper".Equals(methodInfo.RequestedWrapperName));
19+
}
20+
21+
public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
22+
{
23+
// The instrumented method is only called once in the lifetime of a lambda instance so we can just create a handler wrapper each time.
24+
CreateAndSetWrappedHandlerMethod(agent, transaction, instrumentedMethodCall.MethodCall.InvocationTarget);
25+
26+
return Delegates.NoOp;
27+
}
28+
29+
private static void CreateAndSetWrappedHandlerMethod(IAgent agent, ITransaction transaction, object bootstrapper)
30+
{
31+
var bootstrapperType = bootstrapper.GetType();
32+
33+
var handlerReadAccessor = VisibilityBypasser.Instance.GenerateFieldReadAccessor<object>(bootstrapperType, "_handler");
34+
var originalHandler = handlerReadAccessor(bootstrapper);
35+
36+
// 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.
37+
var handlerWriteAccessor = VisibilityBypasser.Instance.GenerateFieldWriteAccessor<LambdaBootstrapHandler>(bootstrapperType, "_handler");
38+
handlerWriteAccessor(bootstrapper, NewHandler);
39+
return;
40+
41+
// Creates a _handler that will wait for the transaction to finish being transformed and harvested before allowing
42+
// the lambda runtime to send a response back. This prevents a race condition where async continuations can run in different orders
43+
// which sometimes causes the lambda runtime to return a response before the agent can generate a lambda payload.
44+
// ref: https://github.com/aws/aws-lambda-dotnet/blob/c28fcfaba68607f785662ff1d232eb9b26d0fa09/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs#L369
45+
// and https://github.com/aws/aws-lambda-dotnet/blob/c28fcfaba68607f785662ff1d232eb9b26d0fa09/Libraries/src/Amazon.Lambda.RuntimeSupport/Bootstrap/LambdaBootstrap.cs#L387
46+
async Task<InvocationResponse> NewHandler(InvocationRequest request)
47+
{
48+
transaction = agent.CreateTransaction(isWeb: false, category: AwsLambdaWrapperExtensions.GetTransactionCategory(agent.Configuration), transactionDisplayName: "TempLambdaName", doNotTrackAsUnitOfWork: true);
49+
agent.Logger.Finest("EnsureTransformCompletesWrapper: Started transaction.");
50+
51+
try
52+
{
53+
agent.Logger.Finest("EnsureTransformCompletesWrapper: Calling original handler.");
54+
var typedHandler = (LambdaBootstrapHandler)originalHandler;
55+
return await typedHandler(request);
56+
}
57+
finally
58+
{
59+
agent.Logger.Finest("EnsureTransformCompletesWrapper: Ending transaction.");
60+
transaction.End();
61+
}
62+
}
63+
}
64+
}

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

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -123,23 +123,23 @@ public bool ValidateWebRequestParameters(InstrumentedMethodCall instrumentedMeth
123123
switch (EventType)
124124
{
125125
case AwsLambdaEventType.APIGatewayHttpApiV2ProxyRequest:
126-
{
127-
if (input.RequestContext != null)
128126
{
129-
dynamic requestContext = input.RequestContext;
127+
if (input.RequestContext != null)
128+
{
129+
dynamic requestContext = input.RequestContext;
130130

131-
if (requestContext.Http != null)
132-
return !string.IsNullOrEmpty(requestContext.Http.Method) && !string.IsNullOrEmpty(requestContext.Http.Path);
133-
}
131+
if (requestContext.Http != null)
132+
return !string.IsNullOrEmpty(requestContext.Http.Method) && !string.IsNullOrEmpty(requestContext.Http.Path);
133+
}
134134

135-
return false;
136-
}
135+
return false;
136+
}
137137
case AwsLambdaEventType.APIGatewayProxyRequest:
138138
case AwsLambdaEventType.ApplicationLoadBalancerRequest:
139-
{
140-
dynamic webReq = input;
141-
return !string.IsNullOrEmpty(webReq.HttpMethod) && !string.IsNullOrEmpty(webReq.Path);
142-
}
139+
{
140+
dynamic webReq = input;
141+
return !string.IsNullOrEmpty(webReq.HttpMethod) && !string.IsNullOrEmpty(webReq.Path);
142+
}
143143
default:
144144
return true;
145145
}
@@ -156,7 +156,7 @@ public bool ValidateWebRequestParameters(InstrumentedMethodCall instrumentedMeth
156156
private static object _initLock = new object();
157157
private static FunctionDetails _functionDetails = null;
158158

159-
public bool IsTransactionRequired => false;
159+
public bool IsTransactionRequired => true;
160160

161161
private static bool _coldStart = true;
162162
private ConcurrentHashSet<string> _unexpectedResponseTypes = new();
@@ -246,13 +246,18 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
246246

247247
// create a transaction for the function invocation if AwsLambdaApmModeEnabled is enabled then based on spec WebTransaction/Function/APIGATEWAY myFunction
248248
// else WebTransaction/Function/myFunction
249-
transaction = agent.CreateTransaction(
250-
isWeb: _functionDetails.EventType.IsWebEvent(),
251-
category: agent.Configuration.AwsLambdaApmModeEnabled ? "Function" : "Lambda",
252-
transactionDisplayName: agent.Configuration.AwsLambdaApmModeEnabled
253-
? _functionDetails.EventType.ToEventTypeString().ToUpper() + " " + _functionDetails.FunctionName
254-
: _functionDetails.FunctionName,
255-
doNotTrackAsUnitOfWork: true);
249+
var transactionName = agent.Configuration.AwsLambdaApmModeEnabled
250+
? _functionDetails.EventType.ToEventTypeString().ToUpper() + " " + _functionDetails.FunctionName
251+
: _functionDetails.FunctionName;
252+
253+
if (_functionDetails.IsWebRequest)
254+
{
255+
transaction.SetWebTransactionName("Lambda", transactionName, TransactionNamePriority.Handler); // low priority to allow for override
256+
}
257+
else
258+
{
259+
transaction.SetOtherTransactionName(AwsLambdaWrapperExtensions.GetTransactionCategory(agent.Configuration), transactionName, TransactionNamePriority.Handler); // low priority to allow for override
260+
}
256261

257262
if (isAsync)
258263
{
@@ -311,7 +316,6 @@ void InvokeTryProcessResponse(Task responseTask)
311316
finally
312317
{
313318
segment?.End();
314-
transaction.End();
315319
}
316320
}
317321
}
@@ -329,7 +333,6 @@ void InvokeTryProcessResponse(Task responseTask)
329333
onFailure: exception =>
330334
{
331335
segment.End(exception);
332-
transaction.End();
333336
});
334337
}
335338
}
@@ -381,4 +384,4 @@ private static bool ValidTaskResponse(Task response)
381384
return response?.Status == TaskStatus.RanToCompletion;
382385
}
383386

384-
}
387+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ SPDX-License-Identifier: Apache-2.0
66
<extension xmlns="urn:newrelic-extension">
77
<instrumentation>
88
<!-- Start the agent init process during the init stage of the lambda for better performance and reduced cost. -->
9-
<tracerFactory name="NewRelic.Agent.Core.Wrapper.NoOpWrapper">
9+
<tracerFactory name="NewRelic.Providers.Wrapper.AwsLambda.EnsureTransformCompletesWrapper">
1010
<match assemblyName="Amazon.Lambda.RuntimeSupport" className="Amazon.Lambda.RuntimeSupport.LambdaBootstrap">
1111
<exactMethodMatcher methodName="RunAsync" />
1212
</match>

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

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,27 @@
1-
<Project Sdk="Microsoft.NET.Sdk.Web">
1+
<Project Sdk="Microsoft.NET.Sdk">
22
<PropertyGroup>
33
<OutputType>Exe</OutputType>
4-
<TargetFramework>net8.0</TargetFramework>
4+
<TargetFrameworks>net8.0;net10.0</TargetFrameworks>
55
<ImplicitUsings>enable</ImplicitUsings>
66
<GenerateRuntimeConfigurationFiles>true</GenerateRuntimeConfigurationFiles>
77
<AWSProjectType>Lambda</AWSProjectType>
88
<!-- This property makes the build directory similar to a publish directory and helps the AWS .NET Lambda Mock Test Tool find project dependencies. -->
99
<CopyLocalLockFileAssemblies>true</CopyLocalLockFileAssemblies>
1010
</PropertyGroup>
11+
<!-- Use Microsoft.NET.Sdk (not .Web) to avoid DeployOnBuild/Publish conflicts with
12+
multi-TFM builds (NETSDK1129). FrameworkReference provides ASP.NET Core APIs. -->
1113
<ItemGroup>
14+
<FrameworkReference Include="Microsoft.AspNetCore.App" />
15+
</ItemGroup>
16+
<ItemGroup>
17+
<!-- 1.7.0 is from Feb 2022 -->
18+
<PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.7.0" Condition="'$(TargetFramework)' == 'net8.0'" />
19+
<PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.14.2" Condition="'$(TargetFramework)' == 'net10.0'" />
20+
1221
<PackageReference Include="Amazon.Lambda.AspNetCoreServer" Version="9.0.1" />
13-
<PackageReference Include="Amazon.Lambda.RuntimeSupport" Version="1.10.0" />
1422
<PackageReference Include="Amazon.Lambda.Serialization.SystemTextJson" Version="2.4.3" />
1523
</ItemGroup>
1624
<ItemGroup>
1725
<ProjectReference Include="..\..\ApplicationHelperLibraries\ApplicationLifecycle\ApplicationLifecycle.csproj" />
1826
</ItemGroup>
19-
</Project>
27+
</Project>

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
// Copyright 2020 New Relic, Inc. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
global using Microsoft.AspNetCore.Builder;
4+
global using Microsoft.AspNetCore.Hosting;
5+
global using Microsoft.AspNetCore.Http;
6+
global using Microsoft.Extensions.Configuration;
7+
global using Microsoft.Extensions.DependencyInjection;
8+
global using Microsoft.Extensions.Hosting;

0 commit comments

Comments
 (0)