Skip to content

Commit

Permalink
tackle issue with Task void result
Browse files Browse the repository at this point in the history
  • Loading branch information
@hjgraca
hjgraca committed Nov 6, 2024
1 parent 27bba28 commit 2752711
Show file tree
Hide file tree
Showing 2 changed files with 238 additions and 12 deletions.
35 changes: 30 additions & 5 deletions libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

using System;
using System.Linq;
using System.Runtime.ExceptionServices;
using System.Text;
using System.Threading.Tasks;
using AspectInjector.Broker;
Expand Down Expand Up @@ -102,22 +103,42 @@ public object Around(

if (result is Task task)
{
task.GetAwaiter().GetResult();
var taskResult = task.GetType().GetProperty("Result")?.GetValue(task);
HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
if (task.IsFaulted && task.Exception != null)
{
var actualException = task.Exception.InnerExceptions.Count == 1
? task.Exception.InnerExceptions[0]
: task.Exception;

// Capture and rethrow the original exception preserving the stack trace
ExceptionDispatchInfo.Capture(actualException).Throw();
}

Check warning on line 114 in libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs

View check run for this annotation

Codecov / codecov/patch

libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs#L114 

Added line #L114 was not covered by tests

// Only handle response if it's not a void Task
if (task.GetType().IsGenericType)
{
var taskType = task.GetType();
var resultProperty = taskType.GetProperty("Result");

// Handle the response only if task completed successfully
if (task.Status == TaskStatus.RanToCompletion)
{
var taskResult = resultProperty?.GetValue(task);
HandleResponse(metadataName, taskResult, trigger.CaptureMode, @namespace);
}
}

_xRayRecorder.EndSubsegment();
return task;
}

HandleResponse(metadataName, result, trigger.CaptureMode, @namespace);

_xRayRecorder.EndSubsegment();
return result;
}
catch (Exception ex)
{
HandleException(ex, metadataName, trigger.CaptureMode, @namespace);
var actualException = ex is AggregateException ae ? ae.InnerException! : ex;
HandleException(actualException, metadataName, trigger.CaptureMode, @namespace);
_xRayRecorder.EndSubsegment();
throw;
}
Expand Down Expand Up @@ -150,6 +171,10 @@ private void BeginSegment(string segmentName, string @namespace)
private void HandleResponse(string name, object result, TracingCaptureMode captureMode, string @namespace)
{
if (!CaptureResponse(captureMode)) return;
if (result == null) return; // Don't try to serialize null results

// Skip if the result is VoidTaskResult
if (result.GetType().Name == "VoidTaskResult") return;

#if NET8_0_OR_GREATER
if (!RuntimeFeatureWrapper.IsDynamicCodeSupported) // is AOT
Expand Down
215 changes: 208 additions & 7 deletions libraries/tests/AWS.Lambda.Powertools.Tracing.Tests/TracingAspectTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,7 @@ public void TracingDisabled_WhenNotInLambdaEnvironment_ReturnsTrue()
}

[Fact]
public async Task WrapVoidTask_SuccessfulExecution_HandlesResponseAndEndsSubsegment()
public async Task WrapVoidTask_SuccessfulExecution_OnlyEndsSubsegment()
{
// Arrange
var tcs = new TaskCompletionSource();
Expand All @@ -318,9 +318,9 @@ public async Task WrapVoidTask_SuccessfulExecution_HandlesResponseAndEndsSubsegm
await wrappedTask!;

// Assert
_mockXRayRecorder.Received(1).AddMetadata(
Arg.Is(nameSpace),
Arg.Is($"{methodName} response"),
_mockXRayRecorder.DidNotReceive().AddMetadata(
Arg.Any<string>(),
Arg.Any<string>(),
Arg.Any<object>()
);
_mockXRayRecorder.Received(1).EndSubsegment();
Expand Down Expand Up @@ -572,9 +572,210 @@ public async Task Around_AsyncMethodWithoutResult_HandlesNullTaskResultProperty(
// Assert with wait
await Task.Delay(100); // Give time for the continuation to complete

_mockXRayRecorder.DidNotReceive().AddMetadata(
Arg.Any<string>(),
Arg.Any<string>(),
Arg.Any<object>()
);
_mockXRayRecorder.Received(1).EndSubsegment();
}

[Fact]
public async Task Around_VoidTask_DoesNotAddResponseMetadata()
{
// Arrange
var tcs = new TaskCompletionSource();
const string methodName = "VoidTaskMethod";
const string nameSpace = "TestNamespace";

// Complete the task before passing to Around
tcs.SetResult();

// Act
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
) as Task;

await wrappedTask!;

// Assert
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
_mockXRayRecorder.Received(1).EndSubsegment();
// Verify that AddMetadata was NOT called with response
_mockXRayRecorder.DidNotReceive().AddMetadata(
Arg.Any<string>(),
Arg.Is<string>(s => s.EndsWith("response")),
Arg.Any<object>()
);
}

[Fact]
public async Task Around_VoidTask_HandlesExceptionCorrectly()
{
// Arrange
var tcs = new TaskCompletionSource();
const string methodName = "VoidTaskMethod";
const string nameSpace = "TestNamespace";
var expectedException = new Exception("Test exception");

// Fail the task before passing to Around
tcs.SetException(expectedException);

// Act & Assert
await Assert.ThrowsAsync<Exception>(async () =>
{
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.ResponseAndError } }
) as Task;

await wrappedTask!;
});

// Assert
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
_mockXRayRecorder.Received(1).AddMetadata(
"TestService",
$"{methodName} response",
null);
Arg.Is(nameSpace),
Arg.Is($"{methodName} error"),
Arg.Is<string>(s => s.Contains(expectedException.Message))
);
_mockXRayRecorder.Received(1).EndSubsegment();
}

[Fact]
public async Task Around_VoidTask_WithCancellation_EndsSegmentCorrectly()
{
// Arrange
using var cts = new CancellationTokenSource();
var tcs = new TaskCompletionSource();
const string methodName = "VoidTaskMethod";
const string nameSpace = "TestNamespace";

// Cancel before passing to Around
cts.Cancel();
tcs.SetCanceled(cts.Token);

// Act & Assert
await Assert.ThrowsAsync<TaskCanceledException>(async () =>
{
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
) as Task;

await wrappedTask!;
});

// Assert
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
_mockXRayRecorder.Received(1).EndSubsegment();
}

[Fact]
public async Task Around_TaskWithResult_AddsResponseMetadata()
{
// Arrange
var tcs = new TaskCompletionSource<string>();
const string methodName = "TaskWithResultMethod";
const string nameSpace = "TestNamespace";
const string result = "test result";

// Complete the task before passing to Around
tcs.SetResult(result);

// Act
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
) as Task<string>;

await wrappedTask!;

// Assert
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
_mockXRayRecorder.Received(1).AddMetadata(
Arg.Is(nameSpace),
Arg.Is($"{methodName} response"),
Arg.Is<string>(s => s == result)
);
_mockXRayRecorder.Received(1).EndSubsegment();
}

[Fact]
public async Task Around_NullResult_DoesNotAddResponseMetadata()
{
// Arrange
var tcs = new TaskCompletionSource<string>();
const string methodName = "NullResultMethod";
const string nameSpace = "TestNamespace";

// Complete the task with null before passing to Around
tcs.SetResult(null!);

// Act
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute { Namespace = nameSpace, CaptureMode = TracingCaptureMode.Response } }
) as Task<string>;

await wrappedTask!;

// Assert
_mockXRayRecorder.Received(1).BeginSubsegment($"## {methodName}");
_mockXRayRecorder.Received(1).EndSubsegment();
// Verify that AddMetadata was NOT called with response
_mockXRayRecorder.DidNotReceive().AddMetadata(
Arg.Any<string>(),
Arg.Is<string>(s => s.EndsWith("response")),
Arg.Any<object>()
);
}

[Fact]
public async Task Around_TracingDisabled_DoesNotAddSegments()
{
// Arrange
_mockConfigurations.TracingDisabled.Returns(true);
var tcs = new TaskCompletionSource();
const string methodName = "DisabledTracingMethod";

// Complete the task before passing to Around
tcs.SetResult();

// Act
var wrappedTask = _handler.Around(
methodName,
new object[] { tcs.Task },
args => args[0],
new Attribute[]
{ new TracingAttribute() }
) as Task;

await wrappedTask!;

// Assert
_mockXRayRecorder.DidNotReceive().BeginSubsegment(Arg.Any<string>());
_mockXRayRecorder.DidNotReceive().EndSubsegment();
_mockXRayRecorder.DidNotReceive().AddMetadata(
Arg.Any<string>(),
Arg.Any<string>(),
Arg.Any<object>()
);
}
}

0 comments on commit 2752711

Please sign in to comment.