diff --git a/src/Proto.Actor/Context/ActorContextDecorator.cs b/src/Proto.Actor/Context/ActorContextDecorator.cs index 39b4173ee2..83591c0405 100644 --- a/src/Proto.Actor/Context/ActorContextDecorator.cs +++ b/src/Proto.Actor/Context/ActorContextDecorator.cs @@ -65,11 +65,11 @@ public virtual void ReenterAfter(Task target, Func, Task> action) public virtual void ReenterAfter(Task target, Action action) => _context.ReenterAfter(target, action); - public void ReenterAfter(Task target, Action action) => _context.ReenterAfter(target, action); + public virtual void ReenterAfter(Task target, Action action) => _context.ReenterAfter(target, action); - public void ReenterAfter(Task target, Action> action) => _context.ReenterAfter(target, action); + public virtual void ReenterAfter(Task target, Action> action) => _context.ReenterAfter(target, action); - public void ReenterAfter(Task target, Func action) => _context.ReenterAfter(target, action); + public virtual void ReenterAfter(Task target, Func action) => _context.ReenterAfter(target, action); public CapturedContext Capture() => _context.Capture(); diff --git a/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs b/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs index 6b8c893f4e..df85ed957f 100644 --- a/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs +++ b/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs @@ -105,6 +105,48 @@ public override void ReenterAfter(Task target, Func, Task> action) base.ReenterAfter(target, a2); } + public override void ReenterAfter(Task target, Action action) + { + var current = Activity.Current?.Context ?? default; + var message = base.Message!; + Action a2 = t => + { + using var x = OpenTelemetryHelpers.BuildStartedActivity(current, Source, nameof(ReenterAfter), message, + _sendActivitySetup); + x?.SetTag(ProtoTags.ActionType, nameof(ReenterAfter)); + action(t); + }; + base.ReenterAfter(target, a2); + } + + public override void ReenterAfter(Task target, Action> action) + { + var current = Activity.Current?.Context ?? default; + var message = base.Message!; + Action> a2 = t => + { + using var x = OpenTelemetryHelpers.BuildStartedActivity(current, Source, nameof(ReenterAfter), message, + _sendActivitySetup); + x?.SetTag(ProtoTags.ActionType, nameof(ReenterAfter)); + action(t); + }; + base.ReenterAfter(target, a2); + } + + public override void ReenterAfter(Task target, Func action) + { + var current = Activity.Current?.Context ?? default; + var message = base.Message!; + Func a2 = async t => + { + using var x = OpenTelemetryHelpers.BuildStartedActivity(current, Source, nameof(ReenterAfter), message, + _sendActivitySetup); + x?.SetTag(ProtoTags.ActionType, nameof(ReenterAfter)); + await action(t).ConfigureAwait(false); + }; + base.ReenterAfter(target, a2); + } + public override PID SpawnNamed(Props props, string name, Action? callback = null) => OpenTelemetryMethodsDecorators.SpawnNamed(Source,_spawnActivitySetup, () => base.SpawnNamed(props, name, callback),name, Actor.GetType().Name); } \ No newline at end of file diff --git a/tests/Proto.OpenTelemetry.Tests/OpenTelemetryTracingTests.cs b/tests/Proto.OpenTelemetry.Tests/OpenTelemetryTracingTests.cs index 0ef69fbf37..b7cdfbaac3 100644 --- a/tests/Proto.OpenTelemetry.Tests/OpenTelemetryTracingTests.cs +++ b/tests/Proto.OpenTelemetry.Tests/OpenTelemetryTracingTests.cs @@ -145,6 +145,21 @@ await VerifyTrace(async (rootContext, target) => response.Should().BeEquivalentTo(new TraceResponse(TestBaggage)); } ); + + [Theory] + [InlineData(SendAs.ReEnterAfter1)] + [InlineData(SendAs.ReEnterAfter2)] + [InlineData(SendAs.ReEnterAfter3)] + [InlineData(SendAs.ReEnterAfter4)] + [InlineData(SendAs.ReEnterAfter5)] + public async Task TracesPropagateCorrectlyWithBaggageForReEnterAfter(SendAs reEnterType) => + await VerifyTrace(async (rootContext, target) => + { + Baggage.Current = TestBaggage; + var response = await rootContext.RequestAsync(target, new TraceMe(reEnterType)); + response.Should().BeEquivalentTo(new TraceResponse(TestBaggage)); + } + ); [Fact] public async Task TracesPropagateCorrectlyWithBaggageForRequest() => @@ -240,13 +255,18 @@ public async Task ExceptionsAreRecorded() receiveActivity.Events.Single().Tags.Where(tag => tag.Key.StartsWith("exception")).Should().NotBeEmpty(); } - private enum SendAs + public enum SendAs { RequestAsync, Request, Send, Forward, - Invalid + Invalid, + ReEnterAfter1, // void ReenterAfter(Task target, Func, Task> action); + ReEnterAfter2, // void ReenterAfter(Task target, Action action); + ReEnterAfter3, // void ReenterAfter(Task target, Action action); + ReEnterAfter4, // void ReenterAfter(Task target, Action> action); + ReEnterAfter5, // void ReenterAfter(Task target, Func action); } private record TraceMe(SendAs Method); @@ -290,6 +310,34 @@ private async Task OnTraceMe(IContext context, TraceMe msg) case SendAs.Forward: context.Forward(target); + break; + case SendAs.ReEnterAfter1: + context.ReenterAfter(Task.FromResult(1), _ => + { + context.Forward(target); + return Task.CompletedTask; + }); + + break; + case SendAs.ReEnterAfter2: + context.ReenterAfter(Task.CompletedTask, () => context.Forward(target)); + + break; + case SendAs.ReEnterAfter3: + context.ReenterAfter(Task.CompletedTask, _ => context.Forward(target)); + + break; + case SendAs.ReEnterAfter4: + context.ReenterAfter(Task.FromResult(1), _ => context.Forward(target)); + + break; + case SendAs.ReEnterAfter5: + context.ReenterAfter(Task.CompletedTask, _ => + { + context.Forward(target); + return Task.CompletedTask; + }); + break; default: throw new ArgumentOutOfRangeException(nameof(msg.Method), msg.Method.ToString()); }