Skip to content

Commit 27bba28

Browse files
hjgracahjgraca
committed
fix async segment not output correctly, add tests
1 parent 764727e commit 27bba28

File tree

4 files changed

+327
-160
lines changed

4 files changed

+327
-160
lines changed

‎libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs

Lines changed: 43 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -86,12 +86,14 @@ public object Around(
8686
if (TracingDisabled())
8787
return target(args);
8888

89-
var @namespace = !string.IsNullOrWhiteSpace(trigger.Namespace) ? trigger.Namespace : _powertoolsConfigurations.Service;
90-
91-
var (segmentName, metadataName) = string.IsNullOrWhiteSpace(trigger.SegmentName)
92-
? ($"## {name}", name)
89+
var @namespace = !string.IsNullOrWhiteSpace(trigger.Namespace)
90+
? trigger.Namespace
91+
: _powertoolsConfigurations.Service;
92+
93+
var (segmentName, metadataName) = string.IsNullOrWhiteSpace(trigger.SegmentName)
94+
? ($"## {name}", name)
9395
: (trigger.SegmentName, trigger.SegmentName);
94-
96+
9597
BeginSegment(segmentName, @namespace);
9698

9799
try
@@ -100,68 +102,29 @@ public object Around(
100102

101103
if (result is Task task)
102104
{
103-
return WrapTask(task, metadataName, trigger.CaptureMode, @namespace);
104-
}
105-
106-
HandleResponse(metadataName, result, trigger.CaptureMode, @namespace);
107-
_xRayRecorder.EndSubsegment(); // End segment here for sync methods
108-
return result;
109-
}
110-
catch (Exception ex)
111-
{
112-
HandleException(ex, metadataName, trigger.CaptureMode, @namespace);
113-
_xRayRecorder.EndSubsegment(); // End segment here for sync methods with exceptions
114-
throw;
115-
}
116-
}
117-
118-
private object WrapTask(Task task, string name, TracingCaptureMode captureMode, string @namespace)
119-
{
120-
if (task.GetType() == typeof(Task))
121-
{
122-
return WrapVoidTask(task, name, captureMode, @namespace);
123-
}
105+
task.GetAwaiter().GetResult();
106+
var taskResult = task.GetType().GetProperty("Result")?.GetValue(task);
107+
HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
124108

125-
// Create an async wrapper task that returns the original task type
126-
async Task AsyncWrapper()
127-
{
128-
try
129-
{
130-
await task;
131-
var result = task.GetType().GetProperty("Result")?.GetValue(task);
132-
HandleResponse(name, result, captureMode, @namespace);
133-
}
134-
catch (Exception ex)
135-
{
136-
HandleException(ex, name, captureMode, @namespace);
137-
throw;
138-
}
139-
finally
140-
{
141109
_xRayRecorder.EndSubsegment();
110+
return task;
142111
}
143-
}
144112

145-
// Start the wrapper and return original task
146-
_ = AsyncWrapper();
147-
return task;
148-
}
113+
HandleResponse(metadataName, result, trigger.CaptureMode, @namespace);
149114

150-
private async Task WrapVoidTask(Task task, string name, TracingCaptureMode captureMode, string @namespace)
151-
{
152-
try
153-
{
154-
await task;
155-
HandleResponse(name, null, captureMode, @namespace);
115+
_xRayRecorder.EndSubsegment();
116+
return result;
156117
}
157118
catch (Exception ex)
158119
{
159-
HandleException(ex, name, captureMode, @namespace);
120+
HandleException(ex, metadataName, trigger.CaptureMode, @namespace);
121+
_xRayRecorder.EndSubsegment();
160122
throw;
161123
}
162124
finally
163125
{
164-
_xRayRecorder.EndSubsegment();
126+
if (_isAnnotationsCaptured)
127+
_captureAnnotations = true;
165128
}
166129
}
167130

@@ -184,15 +147,16 @@ private void BeginSegment(string segmentName, string @namespace)
184147
_isColdStart = false;
185148
}
186149

187-
private void HandleResponse(string segmentName, object result, TracingCaptureMode captureMode, string @namespace)
150+
private void HandleResponse(string name, object result, TracingCaptureMode captureMode, string @namespace)
188151
{
189152
if (!CaptureResponse(captureMode)) return;
153+
190154
#if NET8_0_OR_GREATER
191155
if (!RuntimeFeatureWrapper.IsDynamicCodeSupported) // is AOT
192156
{
193157
_xRayRecorder.AddMetadata(
194158
@namespace,
195-
$"{segmentName} response",
159+
$"{name} response",
196160
Serializers.PowertoolsTracingSerializer.Serialize(result)
197161
);
198162
return;
@@ -201,22 +165,34 @@ private void HandleResponse(string segmentName, object result, TracingCaptureMod
201165

202166
_xRayRecorder.AddMetadata(
203167
@namespace,
204-
$"{segmentName} response",
168+
$"{name} response",
205169
result
206170
);
207171
}
208172

209-
/// <summary>
210-
/// When Aspect Handler exits runs this code to end subsegment
211-
/// </summary>
212-
[Advice(Kind.After)]
213-
public void OnExit()
173+
private void HandleException(Exception exception, string name, TracingCaptureMode captureMode, string @namespace)
214174
{
215-
if (TracingDisabled())
216-
return;
175+
if (!CaptureError(captureMode)) return;
176+
177+
var sb = new StringBuilder();
178+
sb.AppendLine($"Exception type: {exception.GetType()}");
179+
sb.AppendLine($"Exception message: {exception.Message}");
180+
sb.AppendLine($"Stack trace: {exception.StackTrace}");
181+
182+
if (exception.InnerException != null)
183+
{
184+
sb.AppendLine("---BEGIN InnerException--- ");
185+
sb.AppendLine($"Exception type {exception.InnerException.GetType()}");
186+
sb.AppendLine($"Exception message: {exception.InnerException.Message}");
187+
sb.AppendLine($"Stack trace: {exception.InnerException.StackTrace}");
188+
sb.AppendLine("---END Inner Exception");
189+
}
217190

218-
if (_isAnnotationsCaptured)
219-
_captureAnnotations = true;
191+
_xRayRecorder.AddMetadata(
192+
@namespace,
193+
$"{name} error",
194+
sb.ToString()
195+
);
220196
}
221197

222198
private bool TracingDisabled()
@@ -258,35 +234,6 @@ private bool CaptureError(TracingCaptureMode captureMode)
258234
};
259235
}
260236

261-
private void HandleException(Exception exception, string name, TracingCaptureMode captureMode, string @namespace)
262-
{
263-
if (!CaptureError(captureMode)) return;
264-
265-
var nameSpace = @namespace;
266-
var sb = new StringBuilder();
267-
sb.AppendLine($"Exception type: {exception.GetType()}");
268-
sb.AppendLine($"Exception message: {exception.Message}");
269-
sb.AppendLine($"Stack trace: {exception.StackTrace}");
270-
271-
if (exception.InnerException != null)
272-
{
273-
sb.AppendLine("---BEGIN InnerException--- ");
274-
sb.AppendLine($"Exception type {exception.InnerException.GetType()}");
275-
sb.AppendLine($"Exception message: {exception.InnerException.Message}");
276-
sb.AppendLine($"Stack trace: {exception.InnerException.StackTrace}");
277-
sb.AppendLine("---END Inner Exception");
278-
}
279-
280-
_xRayRecorder.AddMetadata(
281-
nameSpace,
282-
$"{name} error",
283-
sb.ToString()
284-
);
285-
}
286-
287-
/// <summary>
288-
/// Resets static variables for test.
289-
/// </summary>
290237
internal static void ResetForTest()
291238
{
292239
_isColdStart = true;

‎libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/Handlers/FullExampleHandler.cs

Lines changed: 96 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
* permissions and limitations under the License.
1414
*/
1515

16+
using System.Collections.Generic;
17+
using System.Net.Http;
1618
using System.Threading.Tasks;
1719
using Amazon.Lambda.Core;
1820

@@ -25,10 +27,24 @@ public async Task<string> Handle(string text, ILambdaContext context)
2527
{
2628
Tracing.AddAnnotation("annotation", "value");
2729
await BusinessLogic1();
30+
31+
var response = DoAction(text);
2832

29-
return await Task.FromResult(text.ToUpper());
33+
return await Task.FromResult(response);
34+
}
35+
36+
[Tracing(SegmentName = "Do Action")]
37+
private string DoAction(string text)
38+
{
39+
return ToUpper(text);
40+
}
41+
42+
[Tracing(SegmentName = "To Upper")]
43+
private string ToUpper(string text)
44+
{
45+
return text.ToUpper();
3046
}
31-
47+
3248
[Tracing(SegmentName = "First Call")]
3349
private async Task BusinessLogic1()
3450
{
@@ -51,4 +67,82 @@ private void GetSomething()
5167
{
5268
Tracing.AddAnnotation("getsomething", "value");
5369
}
70+
}
71+
72+
public class FullExampleHandler2
73+
{
74+
[Tracing]
75+
public static async Task<string> FunctionHandler(string request, ILambdaContext context)
76+
{
77+
string myIp = await GetIp();
78+
79+
var response = CallDynamo(myIp);
80+
81+
return "hello";
82+
}
83+
84+
[Tracing(SegmentName = "Call DynamoDB")]
85+
private static List<string> CallDynamo(string myIp)
86+
{
87+
var newList = ToUpper(myIp, new List<string>() { "Hello", "World" });
88+
return newList;
89+
}
90+
91+
[Tracing(SegmentName = "To Upper")]
92+
private static List<string> ToUpper(string myIp, List<string> response)
93+
{
94+
var newList = new List<string>();
95+
foreach (var item in response)
96+
{
97+
newList.Add(item.ToUpper());
98+
}
99+
100+
newList.Add(myIp);
101+
return newList;
102+
}
103+
104+
[Tracing(SegmentName = "Get Ip Address")]
105+
private static async Task<string> GetIp()
106+
{
107+
return await Task.FromResult("127.0.0.1");
108+
}
109+
}
110+
111+
public class FullExampleHandler3
112+
{
113+
[Tracing]
114+
public static async Task<string> FunctionHandler(string request, ILambdaContext context)
115+
{
116+
string myIp = GetIp();
117+
118+
var response = CallDynamo(myIp);
119+
120+
return "hello";
121+
}
122+
123+
[Tracing(SegmentName = "Call DynamoDB")]
124+
private static List<string> CallDynamo(string myIp)
125+
{
126+
var newList = ToUpper(myIp, new List<string>() { "Hello", "World" });
127+
return newList;
128+
}
129+
130+
[Tracing(SegmentName = "To Upper")]
131+
private static List<string> ToUpper(string myIp, List<string> response)
132+
{
133+
var newList = new List<string>();
134+
foreach (var item in response)
135+
{
136+
newList.Add(item.ToUpper());
137+
}
138+
139+
newList.Add(myIp);
140+
return newList;
141+
}
142+
143+
[Tracing(SegmentName = "Get Ip Address")]
144+
private static string GetIp()
145+
{
146+
return "127.0.0.1";
147+
}
54148
}

0 commit comments

Comments
 (0)