diff --git a/tracer/src/Datadog.Trace/Util/Delegates/DelegateInstrumentation.cs b/tracer/src/Datadog.Trace/Util/Delegates/DelegateInstrumentation.cs index 7f8b43c764d6..6013917d5b3b 100644 --- a/tracer/src/Datadog.Trace/Util/Delegates/DelegateInstrumentation.cs +++ b/tracer/src/Datadog.Trace/Util/Delegates/DelegateInstrumentation.cs @@ -824,6 +824,11 @@ public Func0Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -832,11 +837,6 @@ public Func0Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } @@ -894,6 +894,11 @@ public Func1Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -902,11 +907,6 @@ public Func1Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } @@ -965,6 +965,11 @@ public Func2Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -973,11 +978,6 @@ public Func2Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } @@ -1037,6 +1037,11 @@ public Func3Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -1045,11 +1050,6 @@ public Func3Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } @@ -1110,6 +1110,11 @@ public Func4Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -1118,11 +1123,6 @@ public Func4Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } @@ -1184,6 +1184,11 @@ public Func5Wrapper(Delegate? target, TCallbacks callbacks) { try { + if (SetContinuation is { } setContinuation) + { + returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); + } + returnValue = Callbacks.OnDelegateEnd(sender, returnValue, exception, state); } catch (Exception innerException) @@ -1192,11 +1197,6 @@ public Func5Wrapper(Delegate? target, TCallbacks callbacks) } } - if (SetContinuation is { } setContinuation) - { - returnValue = setContinuation(Callbacks, sender, exception, state, returnValue); - } - return returnValue; } } diff --git a/tracer/test/Datadog.Trace.Tests/Util/Delegates/FuncInstrumentationTests.cs b/tracer/test/Datadog.Trace.Tests/Util/Delegates/FuncInstrumentationTests.cs index 67dda521b60c..ee1bb492f5d0 100644 --- a/tracer/test/Datadog.Trace.Tests/Util/Delegates/FuncInstrumentationTests.cs +++ b/tracer/test/Datadog.Trace.Tests/Util/Delegates/FuncInstrumentationTests.cs @@ -479,6 +479,126 @@ public async Task Async1Test() value.Should().Be(4); } + [Fact] + public async Task Async1WithAsyncExceptionTest() + { + var result = new StrongBox(0); + var callbacks = new Async1Callbacks(); + Func> func = async (arg1) => + { + callbacks.Count.Value++; + // force an exception + int x = 0, y = 0, z = 0; + z = x / y; + await Task.Yield(); + return 42; + }; + func = func.Instrument(callbacks); + await Assert.ThrowsAsync( + async () => + { + result = new StrongBox(await func("Arg01").ConfigureAwait(false)); + }).ConfigureAwait(false); + + callbacks.Count.Value.Should().Be(4); + + var value = 0; + CustomFunc> func2 = async (arg1) => + { + Interlocked.Increment(ref value).Should().Be(2); + // force an exception + int x = 0, y = 0, z = 0; + z = x / y; + await Task.Yield(); + return 42; + }; + func2 = DelegateInstrumentation.Wrap(func2, new DelegateFunc1Callbacks( + (target, arg1) => + { + arg1.Should().Be("Arg01"); + Interlocked.Increment(ref value).Should().Be(1); + return null; + }, + (target, returnValue, exception, state) => + { + Interlocked.Increment(ref value).Should().Be(4); + return returnValue; + }, + onDelegateAsyncEnd: async (sender, returnValue, exception, state) => + { + Interlocked.Increment(ref value).Should().Be(3); + await Task.Delay(100).ConfigureAwait(false); + return ((int)returnValue) + 1; + })); + + await Assert.ThrowsAsync( + async () => + { + result = new StrongBox(await func2("Arg01").ConfigureAwait(false)); + }).ConfigureAwait(false); + + value.Should().Be(4); + } + + [Fact] + public async Task Async1WithExceptionTest() + { + var result = new StrongBox(0); + var callbacks = new Async1Callbacks(); + Func> func = (arg1) => + { + callbacks.Count.Value++; + // force an exception + int x = 0, y = 0, z = 0; + z = x / y; + return Task.FromResult(42); + }; + func = func.Instrument(callbacks); + await Assert.ThrowsAsync( + async () => + { + result = new StrongBox(await func("Arg01").ConfigureAwait(false)); + }).ConfigureAwait(false); + + callbacks.Count.Value.Should().Be(4); + + var value = 0; + CustomFunc> func2 = (arg1) => + { + Interlocked.Increment(ref value).Should().Be(2); + // force an exception + int x = 0, y = 0, z = 0; + z = x / y; + return Task.FromResult(42); + }; + func2 = DelegateInstrumentation.Wrap(func2, new DelegateFunc1Callbacks( + (target, arg1) => + { + arg1.Should().Be("Arg01"); + Interlocked.Increment(ref value).Should().Be(1); + return null; + }, + (target, returnValue, exception, state) => + { + Interlocked.Increment(ref value).Should().Be(4); + return returnValue; + }, + onDelegateAsyncEnd: async (sender, returnValue, exception, state) => + { + Interlocked.Increment(ref value).Should().Be(3); + await Task.Delay(100).ConfigureAwait(false); + return ((int)returnValue) + 1; + })); + + await Assert.ThrowsAsync( + async () => + { + result = new StrongBox(await func2("Arg01").ConfigureAwait(false)); + }).ConfigureAwait(false); + + value.Should().Be(4); + } + public readonly struct Async1Callbacks : IBegin1Callbacks, IReturnCallback, IReturnAsyncCallback { public Async1Callbacks()