Make handler span surround handler pipeline stage (#7319) (#7326)

Newline

Pragma

Fix tests

Verify that the handler span can be accessed in the handler pipeline

Newline

Author: andreasohlund

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/Traces/When_processing_message_with_multiple_handlers.cs

@@ -1,17 +1,19 @@
 namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry.Traces;
 
+using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Threading.Tasks;
 using NServiceBus.AcceptanceTesting;
+using NServiceBus.Pipeline;
 using NUnit.Framework;
 
 public class When_processing_message_with_multiple_handlers : OpenTelemetryAcceptanceTest
 {
     [Test]
     public async Task Should_create_message_handler_spans()
     {
-        var context = await Scenario.Define<Context>()
+        await Scenario.Define<Context>()
             .WithEndpoint<ReceivingEndpoint>(b =>
                 b.When(session => session.SendLocal(new SomeMessage()))
             )
@@ -39,6 +41,8 @@ public async Task Should_create_message_handler_spans()
                 Assert.That(invokedHandlerActivity.ParentId, Is.EqualTo(receivePipelineActivities[0].Id));
                 Assert.That(invokedHandlerActivity.Status, Is.EqualTo(ActivityStatusCode.Ok));
             });
+
+            Assert.That(invokedHandlerActivity.GetTagItem("custom_handler_tag"), Is.Not.Null, "Custom tag should be set");
         }
 
         Assert.That(recordedHandlerTypes, Does.Contain(typeof(ReceivingEndpoint.HandlerOne).FullName), "invocation of handler one should be traced");
@@ -53,27 +57,28 @@ class Context : ScenarioContext
 
     class ReceivingEndpoint : EndpointConfigurationBuilder
     {
-        public ReceivingEndpoint() => EndpointSetup<OpenTelemetryEnabledEndpoint>();
+        public ReceivingEndpoint() => EndpointSetup<OpenTelemetryEnabledEndpoint>(c => c.Pipeline.Register(typeof(AddTagToHandlerSpanBehavior), "Adds a custom tag to the handler span"));
 
-        public class HandlerOne : IHandleMessages<SomeMessage>
+        class AddTagToHandlerSpanBehavior : Behavior<IInvokeHandlerContext>
         {
-            Context testContext;
-
-            public HandlerOne(Context context) => testContext = context;
+            public override async Task Invoke(IInvokeHandlerContext context, Func<Task> next)
+            {
+                Activity.Current?.AddTag("custom_handler_tag", "some value");
+                await next();
+            }
+        }
 
+        public class HandlerOne(Context testContext) : IHandleMessages<SomeMessage>
+        {
             public Task Handle(SomeMessage message, IMessageHandlerContext context)
             {
                 testContext.FirstHandlerRun = true;
                 return Task.CompletedTask;
             }
         }
 
-        public class HandlerTwo : IHandleMessages<SomeMessage>
+        public class HandlerTwo(Context testContext) : IHandleMessages<SomeMessage>
         {
-            Context testContext;
-
-            public HandlerTwo(Context context) => testContext = context;
-
             public Task Handle(SomeMessage message, IMessageHandlerContext context)
             {
                 testContext.SecondHandlerRun = true;

src/NServiceBus.Core.Tests/OpenTelemetry/ActivityFactoryTests.cs

@@ -249,7 +249,7 @@ class StartHandlerActivity : ActivityFactoryTests
         public void Should_not_start_activity_when_no_parent_activity_exists()
         {
             Type handlerType = typeof(StartHandlerActivity);
-            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, handlerType), null);
+            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, handlerType));
 
             Assert.That(activity, Is.Null, "should not start handler activity when no parent activity exists");
         }
@@ -262,7 +262,7 @@ public void Should_set_handler_type_as_tag()
             using var ambientActivity = new Activity("ambient activity");
             ambientActivity.Start();
 
-            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, handlerType), null);
+            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, handlerType));
 
             Assert.That(activity, Is.Not.Null);
             var tags = activity.Tags.ToImmutableDictionary();
@@ -277,12 +277,9 @@ public void Should_set_saga_id_when_saga()
             using var ambientActivity = new Activity("ambient activity");
             ambientActivity.Start();
 
-            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, typeof(StartHandlerActivity)), sagaInstance);
+            var activity = activityFactory.StartHandlerActivity(new MessageHandler((_, _, _) => Task.CompletedTask, typeof(StartHandlerActivity)));
 
             Assert.That(activity, Is.Not.Null);
-            var tags = activity.Tags.ToImmutableDictionary();
-
-            Assert.That(tags[ActivityTags.HandlerSagaId], Is.EqualTo(sagaInstance.SagaId));
         }
     }
 }

src/NServiceBus.Core.Tests/Pipeline/Incoming/InvokeHandlerTerminatorTest.cs

@@ -11,7 +11,7 @@
 [TestFixture]
 public class InvokeHandlerTerminatorTest
 {
-    InvokeHandlerTerminator terminator = new(new NoOpActivityFactory(), new IncomingPipelineMetrics(new TestMeterFactory(), "queue", "disc"));
+    InvokeHandlerTerminator terminator = new(new IncomingPipelineMetrics(new TestMeterFactory(), "queue", "disc"));
 
     [Test]
     public async Task When_saga_found_and_handler_is_saga_should_invoke_handler()

src/NServiceBus.Core.Tests/Unicast/LoadHandlersConnectorTests.cs

@@ -16,7 +16,7 @@ public class LoadHandlersConnectorTests
     [Test]
     public void Should_throw_when_there_are_no_registered_message_handlers()
     {
-        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry());
+        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry(), new NoOpActivityFactory());
 
         var context = new TestableIncomingLogicalMessageContext();
 
@@ -29,7 +29,7 @@ public void Should_throw_when_there_are_no_registered_message_handlers()
     [Test]
     public void Should_throw_if_ambient_transaction_is_different_from_scope_used_by_transport()
     {
-        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry());
+        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry(), new NoOpActivityFactory());
 
         var context = new TestableIncomingLogicalMessageContext();
 
@@ -49,7 +49,7 @@ public void Should_throw_if_ambient_transaction_is_different_from_scope_used_by_
     [Test]
     public void Should_throw_if_ambient_transaction_suppressed_when_transport_uses_a_scope()
     {
-        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry());
+        var behavior = new LoadHandlersConnector(new MessageHandlerRegistry(), new NoOpActivityFactory());
 
         var context = new TestableIncomingLogicalMessageContext();
 
@@ -77,7 +77,7 @@ public void Should_not_throw_if_ambient_scope_is_same_as_transport_scope()
         context.Services.AddSingleton<FakeHandler>();
         context.Extensions.Set<IOutboxTransaction>(new NoOpOutboxTransaction());
 
-        var behavior = new LoadHandlersConnector(messageHandlerRegistry);
+        var behavior = new LoadHandlersConnector(messageHandlerRegistry, new NoOpActivityFactory());
 
         using (new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
         {

src/NServiceBus.Core/OpenTelemetry/Tracing/ActivityFactory.cs

@@ -2,7 +2,6 @@
 
 using System.Diagnostics;
 using Pipeline;
-using Sagas;
 using Transport;
 
 class ActivityFactory : IActivityFactory
@@ -84,7 +83,7 @@ public Activity StartOutgoingPipelineActivity(string activityName, string displa
         return activity;
     }
 
-    public Activity StartHandlerActivity(MessageHandler messageHandler, ActiveSagaInstance saga)
+    public Activity StartHandlerActivity(MessageHandler messageHandler)
     {
         if (Activity.Current == null)
         {
@@ -98,11 +97,6 @@ public Activity StartHandlerActivity(MessageHandler messageHandler, ActiveSagaIn
         {
             activity.DisplayName = messageHandler.HandlerType.Name;
             activity.AddTag(ActivityTags.HandlerType, messageHandler.HandlerType.FullName);
-
-            if (saga != null)
-            {
-                activity.AddTag(ActivityTags.HandlerSagaId, saga.SagaId);
-            }
         }
 
         return activity;

src/NServiceBus.Core/OpenTelemetry/Tracing/IActivityFactory.cs

@@ -2,12 +2,11 @@
 
 using System.Diagnostics;
 using Pipeline;
-using Sagas;
 using Transport;
 
 interface IActivityFactory
 {
     Activity StartIncomingPipelineActivity(MessageContext context);
     Activity StartOutgoingPipelineActivity(string activityName, string displayName, IBehaviorContext outgoingContext);
-    Activity StartHandlerActivity(MessageHandler messageHandler, ActiveSagaInstance saga);
+    Activity StartHandlerActivity(MessageHandler messageHandler);
 }

src/NServiceBus.Core/OpenTelemetry/Tracing/NoOpActivityFactory.cs

@@ -2,7 +2,6 @@
 
 using System.Diagnostics;
 using Pipeline;
-using Sagas;
 using Transport;
 
 class NoOpActivityFactory : IActivityFactory
@@ -11,5 +10,5 @@ class NoOpActivityFactory : IActivityFactory
 
     public Activity StartOutgoingPipelineActivity(string activityName, string displayName, IBehaviorContext outgoingContext) => null;
 
-    public Activity StartHandlerActivity(MessageHandler messageHandler, ActiveSagaInstance saga) => null;
+    public Activity StartHandlerActivity(MessageHandler messageHandler) => null;
 }

src/NServiceBus.Core/Pipeline/Incoming/InvokeHandlerTerminator.cs

@@ -1,12 +1,11 @@
 namespace NServiceBus;
 
 using System;
-using System.Diagnostics;
 using System.Threading.Tasks;
 using Pipeline;
 using Sagas;
 
-class InvokeHandlerTerminator(IActivityFactory activityFactory, IncomingPipelineMetrics messagingMetricsMeters) : PipelineTerminator<IInvokeHandlerContext>
+class InvokeHandlerTerminator(IncomingPipelineMetrics messagingMetricsMeters) : PipelineTerminator<IInvokeHandlerContext>
 {
     protected override async Task Terminate(IInvokeHandlerContext context)
     {
@@ -15,8 +14,6 @@ protected override async Task Terminate(IInvokeHandlerContext context)
             return;
         }
 
-        using var activity = activityFactory.StartHandlerActivity(context.MessageHandler, saga);
-
         var messageHandler = context.MessageHandler;
 
         // Might as well abort before invoking the handler if we're shutting down
@@ -29,7 +26,6 @@ await messageHandler
                 .ThrowIfNull()
                 .ConfigureAwait(false);
 
-            activity?.SetStatus(ActivityStatusCode.Ok);
             messagingMetricsMeters.RecordSuccessfulMessageHandlerTime(context, DateTimeOffset.UtcNow - startTime);
         }
 #pragma warning disable PS0019 // Do not catch Exception without considering OperationCanceledException - enriching and rethrowing
@@ -42,10 +38,8 @@ await messageHandler
             ex.Data["Handler failure time"] = DateTimeOffsetHelper.ToWireFormattedString(DateTimeOffset.UtcNow);
             ex.Data["Handler canceled"] = context.CancellationToken.IsCancellationRequested;
 
-            activity?.SetErrorStatus(ex);
             messagingMetricsMeters.RecordFailedMessageHandlerTime(context, DateTimeOffset.UtcNow - startTime, ex);
             throw;
         }
     }
-
 }

src/NServiceBus.Core/Pipeline/Incoming/LoadHandlersConnector.cs

@@ -2,21 +2,20 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
 using System.Transactions;
 using Logging;
 using Microsoft.Extensions.DependencyInjection;
-using NServiceBus.Transport;
+using Transport;
 using Persistence;
 using Pipeline;
 using Unicast;
 
-class LoadHandlersConnector : StageConnector<IIncomingLogicalMessageContext, IInvokeHandlerContext>
+class LoadHandlersConnector(MessageHandlerRegistry messageHandlerRegistry, IActivityFactory activityFactory) : StageConnector<IIncomingLogicalMessageContext, IInvokeHandlerContext>
 {
-    public LoadHandlersConnector(MessageHandlerRegistry messageHandlerRegistry) => this.messageHandlerRegistry = messageHandlerRegistry;
-
     public override async Task Invoke(IIncomingLogicalMessageContext context, Func<IInvokeHandlerContext, Task> stage)
     {
         ValidateTransactionMode(context);
@@ -46,14 +45,31 @@ public override async Task Invoke(IIncomingLogicalMessageContext context, Func<I
             messageHandler.Instance = context.Builder.GetRequiredService(messageHandler.HandlerType);
 
             var handlingContext = this.CreateInvokeHandlerContext(messageHandler, storageSession, context);
-            await stage(handlingContext).ConfigureAwait(false);
+
+            using (var activity = activityFactory.StartHandlerActivity(messageHandler))
+            {
+                try
+                {
+                    await stage(handlingContext).ConfigureAwait(false);
+
+                    activity?.SetStatus(ActivityStatusCode.Ok);
+                }
+#pragma warning disable PS0019
+                catch (Exception ex)
+#pragma warning restore PS0019
+                {
+                    activity?.SetErrorStatus(ex);
+                    throw;
+                }
+            }
 
             if (handlingContext.HandlerInvocationAborted)
             {
                 //if the chain was aborted skip the other handlers
                 break;
             }
         }
+
         context.MessageHandled = true;
         await storageSession.CompleteAsync(context.CancellationToken).ConfigureAwait(false);
     }
@@ -95,8 +111,6 @@ static void LogHandlersInvocation(IIncomingLogicalMessageContext context, List<M
         logger.Debug(builder.ToString());
     }
 
-    readonly MessageHandlerRegistry messageHandlerRegistry;
-
     static readonly ILog logger = LogManager.GetLogger<LoadHandlersConnector>();
     static readonly bool isDebugIsEnabled = logger.IsDebugEnabled;
     static readonly string scopeInconsistencyMessage =

src/NServiceBus.Core/Receiving/ReceiveComponent.cs

@@ -66,12 +66,9 @@ public static ReceiveComponent Configure(
             return new TransportReceiveToPhysicalMessageConnector(storage, b.GetRequiredService<IncomingPipelineMetrics>());
         }, "Allows to abort processing the message");
 
-        pipelineSettings.Register("LoadHandlersConnector", b =>
-        {
-            return new LoadHandlersConnector(b.GetRequiredService<MessageHandlerRegistry>());
-        }, "Gets all the handlers to invoke from the MessageHandler registry based on the message type.");
+        pipelineSettings.Register("LoadHandlersConnector", b => new LoadHandlersConnector(b.GetRequiredService<MessageHandlerRegistry>(), hostingConfiguration.ActivityFactory), "Gets all the handlers to invoke from the MessageHandler registry based on the message type.");
 
-        pipelineSettings.Register("InvokeHandlers", sp => new InvokeHandlerTerminator(hostingConfiguration.ActivityFactory, sp.GetService<IncomingPipelineMetrics>()), "Calls the IHandleMessages<T>.Handle(T)");
+        pipelineSettings.Register("InvokeHandlers", sp => new InvokeHandlerTerminator(sp.GetService<IncomingPipelineMetrics>()), "Calls the IHandleMessages<T>.Handle(T)");
 
         var handlerDiagnostics = new Dictionary<string, List<string>>();
 

src/NServiceBus.Core/Sagas/SagaPersistenceBehavior.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading.Tasks;
 using Logging;
@@ -111,6 +112,8 @@ public async Task Invoke(IInvokeHandlerContext context, Func<IInvokeHandlerConte
             sagaInstanceState.AttachExistingEntity(loadedEntity);
         }
 
+        Activity.Current?.AddTag(ActivityTags.HandlerSagaId, sagaInstanceState.SagaId);
+
         await next(context).ConfigureAwait(false);
 
         if (sagaInstanceState.NotFound)
@@ -265,6 +268,7 @@ static SagaFinderDefinition GetSagaFinder(SagaMetadata metadata, IInvokeHandlerC
                 return finderDefinition;
             }
         }
+
         return null;
     }
 
@@ -314,4 +318,4 @@ IContainSagaData CreateNewSagaEntity(SagaMetadata metadata, IInvokeHandlerContex
 
     static readonly Task<IContainSagaData> DefaultSagaDataCompletedTask = Task.FromResult(default(IContainSagaData));
     static readonly ILog logger = LogManager.GetLogger<SagaPersistenceBehavior>();
-}
+}