Skip to content

Commit b478ec9

Browse files
joeyzhao2018joeyzhao2018
committed
1st working version using the DelegateInstrumentation util
1 parent 767ac88 commit b478ec9

File tree

3 files changed

+85
-73
lines changed

3 files changed

+85
-73
lines changed

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

Lines changed: 78 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -33,85 +33,95 @@ public class HandlerWrapperSetHandlerIntegration
3333
{
3434
private const string IntegrationName = nameof(IntegrationId.AwsLambda);
3535
private static readonly ILambdaExtensionRequest RequestBuilder = new LambdaRequestBuilder();
36+
private static readonly Async1Callbacks _callbacks = new Async1Callbacks();
3637

3738
/// <summary>
38-
/// OnMethodBegin callback. The input Delegate handler is the customer's handler wrapped by the HandlerWrapper.
39-
/// Here we will try further wrap it with our DelegateWrapper class in order to make it notify the Datadog-Extension
40-
/// before and after the handler's each invocation.
39+
/// OnMethodBegin callback. The input Delegate handler is the customer's handler.
40+
/// And it's already wrapped by the Amazon.Lambda.RuntimeSupport.HandlerWrapper.
41+
/// Here we will try further wrap it with our DelegateWrapper class
42+
/// in order to make it notify the Datadog-Extension before and after the handler's each invocation.
4143
/// </summary>
4244
/// <typeparam name="TTarget">Type of the target</typeparam>
4345
/// <param name="instance">Instance value, aka `this` of the instrumented method</param>
4446
/// <param name="handler">Instance of Amazon.Lambda.RuntimeSupport.LambdaBootstrapHandler</param>
4547
/// <returns>Calltarget state value</returns>
4648
internal static CallTargetState OnMethodBegin<TTarget>(TTarget instance, ref Delegate handler)
4749
{
48-
Serverless.Debug("DelegateWrapper Wrapping the Handler");
50+
handler = handler.Instrument(_callbacks);
51+
return CallTargetState.GetDefault();
52+
}
4953

50-
var state = CallTargetState.GetDefault();
51-
handler = DelegateInstrumentation.Wrap(handler, new DelegateFunc1Callbacks(
52-
(sender, arg) =>
53-
{
54-
Serverless.Debug("DelegateWrapper Running onDelegateBegin");
55-
try
56-
{
57-
var proxyInstance = arg.DuckCast<IInvocationRequest>();
58-
if (proxyInstance != null)
59-
{
60-
var reader = new StreamReader(proxyInstance.InputStream, encoding: Encoding.UTF8, detectEncodingFromByteOrderMarks: true, bufferSize: -1, leaveOpen: true);
61-
string json = reader.ReadToEnd();
62-
proxyInstance.InputStream.Seek(0, SeekOrigin.Begin);
63-
var scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), json, proxyInstance.LambdaContext?.ClientContext?.Custom);
64-
state = new CallTargetState(scope);
65-
}
66-
else
67-
{
68-
var scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), string.Empty, null);
69-
state = new CallTargetState(scope);
70-
}
71-
}
72-
catch (Exception ex)
73-
{
74-
Serverless.Error("Could not send payload to the extension", ex);
75-
Console.WriteLine(ex.StackTrace);
76-
}
54+
private readonly struct Async1Callbacks : IBegin1Callbacks, IReturnCallback, IReturnAsyncCallback
55+
{
56+
public bool PreserveAsyncContext => false;
7757

78-
return null;
79-
},
80-
(sender, arg, returnValue, exception) =>
81-
{
82-
Serverless.Debug("DelegateWrapper Running onDelegateEnd");
83-
return returnValue;
84-
},
85-
onDelegateAsyncEnd: async (sender, returnValue, exception, arg) =>
86-
{
87-
return await Task.Run(() =>
88-
{
89-
Serverless.Debug("DelegateWrapper Running onDelegateAsyncEnd");
90-
try
91-
{
92-
var proxyInstance = returnValue.DuckCast<IInvocationResponse>();
93-
if (proxyInstance != null)
94-
{
95-
var reader = new StreamReader(proxyInstance.OutputStream, encoding: Encoding.UTF8, detectEncodingFromByteOrderMarks: true, bufferSize: -1, leaveOpen: true);
96-
string json = reader.ReadToEnd();
97-
proxyInstance.OutputStream.Seek(0, SeekOrigin.Begin);
98-
LambdaCommon.EndInvocationSync(json, exception, state.Scope, RequestBuilder);
99-
}
100-
else
101-
{
102-
LambdaCommon.EndInvocationSync(string.Empty, exception, state.Scope, RequestBuilder);
103-
}
104-
}
105-
catch (Exception ex)
106-
{
107-
Serverless.Error("Could not call EndInvocationSync", ex);
108-
Console.WriteLine(ex.StackTrace);
109-
}
58+
public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
59+
{
60+
Serverless.Debug("DelegateWrapper Running onDelegateBegin");
61+
try
62+
{
63+
var proxyInstance = arg.DuckCast<IInvocationRequest>();
64+
if (proxyInstance != null)
65+
{
66+
var reader = new StreamReader(proxyInstance.InputStream, Encoding.UTF8, leaveOpen: true);
67+
var json = reader.ReadToEnd();
68+
proxyInstance.InputStream.Seek(0, SeekOrigin.Begin);
69+
var scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), json, proxyInstance.LambdaContext?.ClientContext?.Custom);
70+
return new CallTargetState(scope);
71+
}
72+
else
73+
{
74+
var scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), string.Empty, null);
75+
return new CallTargetState(scope);
76+
}
77+
}
78+
catch (Exception ex)
79+
{
80+
Serverless.Error("Could not send payload to the extension", ex);
81+
Console.WriteLine(ex.StackTrace);
82+
}
11083

111-
Serverless.Debug("DelegateWrapper FINISHED Running onDelegateAsyncEnd");
112-
return returnValue;
113-
}).ConfigureAwait(false);
114-
}));
115-
return CallTargetState.GetDefault();
84+
return CallTargetState.GetDefault();
85+
}
86+
87+
/// <inheritdoc/>
88+
public TReturn OnDelegateEnd<TReturn>(object sender, TReturn returnValue, Exception exception, object state)
89+
{
90+
return returnValue;
91+
}
92+
93+
/// <inheritdoc/>
94+
public Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
95+
{
96+
Serverless.Debug("DelegateWrapper Running onDelegateAsyncEnd");
97+
try
98+
{
99+
var proxyInstance = returnValue.DuckCast<IInvocationResponse>();
100+
if (proxyInstance != null)
101+
{
102+
var reader = new StreamReader(proxyInstance.OutputStream, Encoding.UTF8, leaveOpen: true);
103+
var json = reader.ReadToEnd();
104+
proxyInstance.OutputStream.Seek(0, SeekOrigin.Begin);
105+
LambdaCommon.EndInvocationSync(json, exception, ((CallTargetState)state!).Scope, RequestBuilder);
106+
}
107+
else
108+
{
109+
LambdaCommon.EndInvocationSync(string.Empty, exception, ((CallTargetState)state!).Scope, RequestBuilder);
110+
}
111+
}
112+
catch (Exception ex)
113+
{
114+
Serverless.Debug(ex.StackTrace);
115+
LambdaCommon.EndInvocationSync(string.Empty, ex, ((CallTargetState)state!).Scope, RequestBuilder);
116+
}
117+
118+
Serverless.Debug("DelegateWrapper FINISHED Running onDelegateAsyncEnd");
119+
return Task.FromResult(returnValue);
120+
}
121+
122+
/// <inheritdoc/>
123+
public void OnException(object sender, Exception ex)
124+
{
125+
}
116126
}
117127
}

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/IInvocationRequest.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@
99
namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK;
1010

1111
/// <summary>
12-
/// Class that contains all the information necessary to handle an invocation of an AWS Lambda function.
12+
/// Interface that contains all the information necessary to handle an invocation of an AWS Lambda function.
13+
/// This is the DuckType for Amazon.Lambda.RuntimeSupport.InvocationRequest
1314
/// </summary>
1415
public interface IInvocationRequest : IDisposable
1516
{

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/IInvocationResponse.cs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,18 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK;
99

1010
/// <summary>
1111
/// Interface that contains the response for an invocation of an AWS Lambda function.
12+
/// This is the DuckType for Amazon.Lambda.RuntimeSupport.InvocationResponse
1213
/// </summary>
13-
public class IInvocationResponse
14+
public interface IInvocationResponse
1415
{
1516
/// <summary>
16-
/// Gets or sets output from the function invocation.
17+
/// Gets output from the function invocation.
1718
/// </summary>
18-
public Stream OutputStream { get; set; }
19+
public Stream OutputStream { get; internal set; }
1920

2021
/// <summary>
2122
/// Gets a value indicating whether true if the LambdaBootstrap should dispose the stream after it's read, false otherwise.
2223
/// Set this to false if you plan to reuse the same output stream for multiple invocations of the function.
2324
/// </summary>
24-
public bool DisposeOutputStream { get; private set; } = true;
25+
public bool DisposeOutputStream { get; internal set; }
2526
}

0 commit comments

Comments
 (0)