Skip to content

Commit 2752711

Browse files
hjgracahjgraca
committed
tackle issue with Task void result
1 parent 27bba28 commit 2752711

File tree

2 files changed

+238
-12
lines changed

2 files changed

+238
-12
lines changed

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

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
using System;
1717
using System.Linq;
18+
using System.Runtime.ExceptionServices;
1819
using System.Text;
1920
using System.Threading.Tasks;
2021
using AspectInjector.Broker;
@@ -102,22 +103,42 @@ public object Around(
102103

103104
if (result is Task task)
104105
{
105-
task.GetAwaiter().GetResult();
106-
var taskResult = task.GetType().GetProperty("Result")?.GetValue(task);
107-
HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
106+
if (task.IsFaulted && task.Exception != null)
107+
{
108+
var actualException = task.Exception.InnerExceptions.Count == 1
109+
? task.Exception.InnerExceptions[0]
110+
: task.Exception;
111+
112+
// Capture and rethrow the original exception preserving the stack trace
113+
ExceptionDispatchInfo.Capture(actualException).Throw();
114+
}
115+
116+
// Only handle response if it's not a void Task
117+
if (task.GetType().IsGenericType)
118+
{
119+
var taskType = task.GetType();
120+
var resultProperty = taskType.GetProperty("Result");
121+
122+
// Handle the response only if task completed successfully
123+
if (task.Status == TaskStatus.RanToCompletion)
124+
{
125+
var taskResult = resultProperty?.GetValue(task);
126+
HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
127+
}
128+
}
108129

109130
_xRayRecorder.EndSubsegment();
110131
return task;
111132
}
112133

113134
HandleResponse(metadataName, result, trigger.CaptureMode, @namespace);
114-
115135
_xRayRecorder.EndSubsegment();
116136
return result;
117137
}
118138
catch (Exception ex)
119139
{
120-
HandleException(ex, metadataName, trigger.CaptureMode, @namespace);
140+
var actualException = ex is AggregateException ae ? ae.InnerException! : ex;
141+
HandleException(actualException, metadataName, trigger.CaptureMode, @namespace);
121142
_xRayRecorder.EndSubsegment();
122143
throw;
123144
}
@@ -150,6 +171,10 @@ private void BeginSegment(string segmentName, string @namespace)
150171
private void HandleResponse(string name, object result, TracingCaptureMode captureMode, string @namespace)
151172
{
152173
if (!CaptureResponse(captureMode)) return;
174+
if (result == null) return; // Don't try to serialize null results
175+
176+
// Skip if the result is VoidTaskResult
177+
if (result.GetType().Name == "VoidTaskResult") return;
153178

154179
#if NET8_0_OR_GREATER
155180
if (!RuntimeFeatureWrapper.IsDynamicCodeSupported) // is AOT

‎libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/TracingAspectTests.cs

Lines changed: 208 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ public void TracingDisabled_WhenNotInLambdaEnvironment_ReturnsTrue()
295295
}
296296

297297
[Fact]
298-
public async Task WrapVoidTask_SuccessfulExecution_HandlesResponseAndEndsSubsegment()
298+
public async Task WrapVoidTask_SuccessfulExecution_OnlyEndsSubsegment()
299299
{
300300
// Arrange
301301
var tcs = new TaskCompletionSource();
@@ -318,9 +318,9 @@ public async Task WrapVoidTask_SuccessfulExecution_HandlesResponseAndEndsSubsegm
318318
await wrappedTask!;
319319

320320
// Assert
321-
_mockXRayRecorder.Received(1).AddMetadata(
322-
Arg.Is(nameSpace),
323-
Arg.Is($"{methodName} response"),
321+
_mockXRayRecorder.DidNotReceive().AddMetadata(
322+
Arg.Any<string>(),
323+
Arg.Any<string>(),
324324
Arg.Any<object>()
325325
);
326326
_mockXRayRecorder.Received(1).EndSubsegment();
@@ -572,9 +572,210 @@ public async Task Around_AsyncMethodWithoutResult_HandlesNullTaskResultProperty(
572572
// Assert with wait
573573
await Task.Delay(100); // Give time for the continuation to complete
574574

575+
_mockXRayRecorder.DidNotReceive().AddMetadata(
576+
Arg.Any<string>(),
577+
Arg.Any<string>(),
578+
Arg.Any<object>()
579+
);
580+
_mockXRayRecorder.Received(1).EndSubsegment();
581+
}
582+
583+
[Fact]
584+
public async Task Around_VoidTask_DoesNotAddResponseMetadata()
585+
{
586+
// Arrange
587+
var tcs = new TaskCompletionSource();
588+
const string methodName = "VoidTaskMethod";
589+
const string nameSpace = "TestNamespace";
590+
591+
// Complete the task before passing to Around
592+
tcs.SetResult();
593+
594+
// Act
595+
var wrappedTask = _handler.Around(
596+
methodName,
597+
new object[] { tcs.Task },
598+
args => args[0],
599+
new Attribute[]
600+
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
601+
) as Task;
602+
603+
await wrappedTask!;
604+
605+
// Assert
606+
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
607+
_mockXRayRecorder.Received(1).EndSubsegment();
608+
// Verify that AddMetadata was NOT called with response
609+
_mockXRayRecorder.DidNotReceive().AddMetadata(
610+
Arg.Any<string>(),
611+
Arg.Is<string>(s => s.EndsWith("response")),
612+
Arg.Any<object>()
613+
);
614+
}
615+
616+
[Fact]
617+
public async Task Around_VoidTask_HandlesExceptionCorrectly()
618+
{
619+
// Arrange
620+
var tcs = new TaskCompletionSource();
621+
const string methodName = "VoidTaskMethod";
622+
const string nameSpace = "TestNamespace";
623+
var expectedException = new Exception("Test exception");
624+
625+
// Fail the task before passing to Around
626+
tcs.SetException(expectedException);
627+
628+
// Act & Assert
629+
await Assert.ThrowsAsync<Exception>(async () =>
630+
{
631+
var wrappedTask = _handler.Around(
632+
methodName,
633+
new object[] { tcs.Task },
634+
args => args[0],
635+
new Attribute[]
636+
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.ResponseAndError } }
637+
) as Task;
638+
639+
await wrappedTask!;
640+
});
641+
642+
// Assert
643+
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
575644
_mockXRayRecorder.Received(1).AddMetadata(
576-
"TestService",
577-
$"{methodName} response",
578-
null);
645+
Arg.Is(nameSpace),
646+
Arg.Is($"{methodName} error"),
647+
Arg.Is<string>(s => s.Contains(expectedException.Message))
648+
);
649+
_mockXRayRecorder.Received(1).EndSubsegment();
650+
}
651+
652+
[Fact]
653+
public async Task Around_VoidTask_WithCancellation_EndsSegmentCorrectly()
654+
{
655+
// Arrange
656+
using var cts = new CancellationTokenSource();
657+
var tcs = new TaskCompletionSource();
658+
const string methodName = "VoidTaskMethod";
659+
const string nameSpace = "TestNamespace";
660+
661+
// Cancel before passing to Around
662+
cts.Cancel();
663+
tcs.SetCanceled(cts.Token);
664+
665+
// Act & Assert
666+
await Assert.ThrowsAsync<TaskCanceledException>(async () =>
667+
{
668+
var wrappedTask = _handler.Around(
669+
methodName,
670+
new object[] { tcs.Task },
671+
args => args[0],
672+
new Attribute[]
673+
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
674+
) as Task;
675+
676+
await wrappedTask!;
677+
});
678+
679+
// Assert
680+
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
681+
_mockXRayRecorder.Received(1).EndSubsegment();
682+
}
683+
684+
[Fact]
685+
public async Task Around_TaskWithResult_AddsResponseMetadata()
686+
{
687+
// Arrange
688+
var tcs = new TaskCompletionSource<string>();
689+
const string methodName = "TaskWithResultMethod";
690+
const string nameSpace = "TestNamespace";
691+
const string result = "test result";
692+
693+
// Complete the task before passing to Around
694+
tcs.SetResult(result);
695+
696+
// Act
697+
var wrappedTask = _handler.Around(
698+
methodName,
699+
new object[] { tcs.Task },
700+
args => args[0],
701+
new Attribute[]
702+
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
703+
) as Task<string>;
704+
705+
await wrappedTask!;
706+
707+
// Assert
708+
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
709+
_mockXRayRecorder.Received(1).AddMetadata(
710+
Arg.Is(nameSpace),
711+
Arg.Is($"{methodName} response"),
712+
Arg.Is<string>(s => s == result)
713+
);
714+
_mockXRayRecorder.Received(1).EndSubsegment();
715+
}
716+
717+
[Fact]
718+
public async Task Around_NullResult_DoesNotAddResponseMetadata()
719+
{
720+
// Arrange
721+
var tcs = new TaskCompletionSource<string>();
722+
const string methodName = "NullResultMethod";
723+
const string nameSpace = "TestNamespace";
724+
725+
// Complete the task with null before passing to Around
726+
tcs.SetResult(null!);
727+
728+
// Act
729+
var wrappedTask = _handler.Around(
730+
methodName,
731+
new object[] { tcs.Task },
732+
args => args[0],
733+
new Attribute[]
734+
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
735+
) as Task<string>;
736+
737+
await wrappedTask!;
738+
739+
// Assert
740+
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
741+
_mockXRayRecorder.Received(1).EndSubsegment();
742+
// Verify that AddMetadata was NOT called with response
743+
_mockXRayRecorder.DidNotReceive().AddMetadata(
744+
Arg.Any<string>(),
745+
Arg.Is<string>(s => s.EndsWith("response")),
746+
Arg.Any<object>()
747+
);
748+
}
749+
750+
[Fact]
751+
public async Task Around_TracingDisabled_DoesNotAddSegments()
752+
{
753+
// Arrange
754+
_mockConfigurations.TracingDisabled.Returns(true);
755+
var tcs = new TaskCompletionSource();
756+
const string methodName = "DisabledTracingMethod";
757+
758+
// Complete the task before passing to Around
759+
tcs.SetResult();
760+
761+
// Act
762+
var wrappedTask = _handler.Around(
763+
methodName,
764+
new object[] { tcs.Task },
765+
args => args[0],
766+
new Attribute[]
767+
{ new TracingAttribute() }
768+
) as Task;
769+
770+
await wrappedTask!;
771+
772+
// Assert
773+
_mockXRayRecorder.DidNotReceive().BeginSubsegment(Arg.Any<string>());
774+
_mockXRayRecorder.DidNotReceive().EndSubsegment();
775+
_mockXRayRecorder.DidNotReceive().AddMetadata(
776+
Arg.Any<string>(),
777+
Arg.Any<string>(),
778+
Arg.Any<object>()
779+
);
579780
}
580781
}

0 commit comments

Comments
 (0)