otel decorator should implement all ReenterAfter methods + add tests …

…for baggage on ReenterAfter (#2138)
asynkron · Oct 12, 2024 · 4bbf51b · 4bbf51b
1 parent bf49315
commit 4bbf51b
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 5 deletions.
diff --git a/src/Proto.Actor/Context/ActorContextDecorator.cs b/src/Proto.Actor/Context/ActorContextDecorator.cs
@@ -65,11 +65,11 @@ public virtual void ReenterAfter<T>(Task<T> target, Func<Task<T>, Task> action)
 
     public virtual void ReenterAfter(Task target, Action action) => _context.ReenterAfter(target, action);
 
-    public void ReenterAfter(Task target, Action<Task> action) => _context.ReenterAfter(target, action);
+    public virtual void ReenterAfter(Task target, Action<Task> action) => _context.ReenterAfter(target, action);
 
-    public void ReenterAfter<T>(Task<T> target, Action<Task<T>> action) => _context.ReenterAfter(target, action);
+    public virtual void ReenterAfter<T>(Task<T> target, Action<Task<T>> action) => _context.ReenterAfter(target, action);
 
-    public void ReenterAfter(Task target, Func<Task, Task> action) => _context.ReenterAfter(target, action);
+    public virtual void ReenterAfter(Task target, Func<Task, Task> action) => _context.ReenterAfter(target, action);
 
     public CapturedContext Capture() => _context.Capture();
 

diff --git a/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs b/src/Proto.OpenTelemetry/OpenTelemetryActorContextDecorator.cs
@@ -105,6 +105,48 @@ public override void ReenterAfter<T>(Task<T> target, Func<Task<T>, Task> action)
         base.ReenterAfter(target, a2);
     }
 
+    public override void ReenterAfter(Task target, Action<Task> action)
+    {
+        var current = Activity.Current?.Context ?? default;
+        var message = base.Message!;
+        Action<Task> 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<T>(Task<T> target, Action<Task<T>> action)
+    {
+        var current = Activity.Current?.Context ?? default;
+        var message = base.Message!;
+        Action<Task<T>> 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<Task, Task> action)
+    {
+        var current = Activity.Current?.Context ?? default;
+        var message = base.Message!;
+        Func<Task, Task> 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<IContext>? callback = null) => 
         OpenTelemetryMethodsDecorators.SpawnNamed(Source,_spawnActivitySetup, () => base.SpawnNamed(props, name, callback),name, Actor.GetType().Name);
 }
diff --git 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<TraceResponse>(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<T>(Task<T> target, Func<Task<T>, Task> action);
+        ReEnterAfter2, // void ReenterAfter(Task target, Action action);
+        ReEnterAfter3, // void ReenterAfter(Task target, Action<Task> action);
+        ReEnterAfter4, // void ReenterAfter<T>(Task<T> target, Action<Task<T>> action);
+        ReEnterAfter5, // void ReenterAfter(Task target, Func<Task, Task> 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());
             }