Initial support for OpenTelemetry (#6410)

Co-Authored-By: Hadi Eskandari <[email protected]>
Co-authored-by: Mike Minutillo <[email protected]>
Co-authored-by: Laila Bougria <[email protected]>
Co-authored-by: Tim Bussmann <[email protected]>

Author: 4 people

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/ActivityTestingExtensions.cs

@@ -0,0 +1,38 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry;
+
+using System;
+using System.Collections.Immutable;
+using System.Diagnostics;
+using NUnit.Framework;
+
+public static class ActivityTestingExtensions
+{
+    public static void VerifyTag(this ImmutableDictionary<string, string> tags, string tagName, string expectedValue)
+    {
+        Assert.IsTrue(tags.TryGetValue(tagName, out var tagValue), $"Tags should contain key '{tagName}'");
+        Assert.AreEqual(expectedValue, tagValue, $"Tag value with key '{tagName}' is incorrect");
+    }
+
+    /// <summary>
+    /// Checks tags for duplicate tag keys.
+    /// </summary>
+    public static void VerifyUniqueTags(this Activity activity)
+    {
+        var tagsList = activity.Tags.ToImmutableList();
+
+        if (tagsList.Count < 2)
+        {
+            return;
+        }
+
+        var sortedTags = tagsList.Sort((a, b) => StringComparer.CurrentCultureIgnoreCase.Compare(a.Key, b.Key));
+
+        for (int i = 0; i < sortedTags.Count - 1; i++)
+        {
+            if (StringComparer.InvariantCultureIgnoreCase.Equals(sortedTags[i].Key, sortedTags[i + 1].Key))
+            {
+                Assert.Fail($"duplicate tag found: {sortedTags[i].Key}. Tags should be unique.");
+            }
+        }
+    }
+}

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/OpenTelemetryAcceptanceTest.cs

@@ -0,0 +1,19 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry;
+
+using NUnit.Framework;
+
+[NonParallelizable] // Ensure only activities for the current test are captured
+public class OpenTelemetryAcceptanceTest : NServiceBusAcceptanceTest
+{
+    protected TestingActivityListener NServicebusActivityListener { get; private set; }
+
+    [SetUp]
+    public void Setup() => NServicebusActivityListener = TestingActivityListener.SetupDiagnosticListener("NServiceBus.Core");
+
+    [TearDown]
+    public void Cleanup()
+    {
+        NServicebusActivityListener?.VerifyAllActivitiesCompleted();
+        NServicebusActivityListener?.Dispose();
+    }
+}

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/TestingActivityListener.cs

@@ -0,0 +1,62 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry
+{
+    using System;
+    using System.Collections.Concurrent;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Linq;
+    using NUnit.Framework;
+
+    public class TestingActivityListener : IDisposable
+    {
+        readonly ActivityListener activityListener;
+
+        public static TestingActivityListener SetupDiagnosticListener(string sourceName)
+        {
+            var testingListener = new TestingActivityListener(sourceName);
+
+            ActivitySource.AddActivityListener(testingListener.activityListener);
+            return testingListener;
+        }
+
+        TestingActivityListener(string sourceName = null)
+        {
+            activityListener = new ActivityListener
+            {
+                ShouldListenTo = source => string.IsNullOrEmpty(sourceName) || source.Name == sourceName,
+                Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData,
+                SampleUsingParentId = (ref ActivityCreationOptions<string> options) => ActivitySamplingResult.AllData
+            };
+            activityListener.ActivityStarted += activity =>
+            {
+                StartedActivities.Enqueue(activity);
+            };
+            activityListener.ActivityStopped += activity =>
+            {
+                CompletedActivities.Enqueue(activity);
+            };
+        }
+
+        public void Dispose() => activityListener?.Dispose();
+
+        public ConcurrentQueue<Activity> StartedActivities { get; } = new ConcurrentQueue<Activity>();
+        public ConcurrentQueue<Activity> CompletedActivities { get; } = new ConcurrentQueue<Activity>();
+
+        public void VerifyAllActivitiesCompleted()
+        {
+            Assert.AreEqual(StartedActivities.Count, CompletedActivities.Count, "all started activities should be completed");
+        }
+    }
+
+    static class ActivityExtensions
+    {
+        public static List<Activity> GetReceiveMessageActivities(this ConcurrentQueue<Activity> activities, bool includeControlMessages = false)
+            => activities.Where(a => a.OperationName == "NServiceBus.Diagnostics.ReceiveMessage")
+                         .Where(a => includeControlMessages || !Convert.ToBoolean(a.GetTagItem("nservicebus.control_message")))
+                         .ToList();
+
+        public static List<Activity> GetSendMessageActivities(this ConcurrentQueue<Activity> activities) => activities.Where(a => a.OperationName == "NServiceBus.Diagnostics.SendMessage").ToList();
+        public static List<Activity> GetPublishEventActivities(this ConcurrentQueue<Activity> activities) => activities.Where(a => a.OperationName == "NServiceBus.Diagnostics.PublishMessage").ToList();
+        public static List<Activity> GetInvokedHandlerActivities(this ConcurrentQueue<Activity> activities) => activities.Where(a => a.OperationName == "NServiceBus.Diagnostics.InvokeHandler").ToList();
+    }
+}

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/TestingMetricListener.cs

@@ -0,0 +1,86 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry
+{
+    using System;
+    using System.Collections.Concurrent;
+    using System.Collections.Generic;
+    using System.Diagnostics.Metrics;
+    using System.Linq;
+    using NUnit.Framework;
+
+    class TestingMetricListener : IDisposable
+    {
+        readonly MeterListener meterListener;
+
+        public TestingMetricListener(string sourceName)
+        {
+            meterListener = new()
+            {
+                InstrumentPublished = (instrument, listener) =>
+                {
+                    if (instrument.Meter.Name == sourceName)
+                    {
+                        TestContext.WriteLine($"Subscribing to {instrument.Meter.Name}\\{instrument.Name}");
+                        listener.EnableMeasurementEvents(instrument);
+                    }
+                }
+            };
+
+            meterListener.SetMeasurementEventCallback((Instrument instrument,
+                long measurement,
+                ReadOnlySpan<KeyValuePair<string, object>> t,
+                object _) =>
+            {
+                TestContext.WriteLine($"{instrument.Meter.Name}\\{instrument.Name}:{measurement}");
+
+                var tags = t.ToArray();
+                ReportedMeters.AddOrUpdate(instrument.Name, measurement, (_, val) => val + measurement);
+                Tags.AddOrUpdate(instrument.Name, _ => tags, (_, _) => tags);
+            });
+            meterListener.Start();
+        }
+
+        public static TestingMetricListener SetupNServiceBusMetricsListener() =>
+            SetupMetricsListener("NServiceBus.Core");
+
+        public static TestingMetricListener SetupMetricsListener(string sourceName)
+        {
+            var testingMetricListener = new TestingMetricListener(sourceName);
+            return testingMetricListener;
+        }
+
+        public void Dispose() => meterListener?.Dispose();
+
+        public ConcurrentDictionary<string, long> ReportedMeters { get; } = new();
+        public ConcurrentDictionary<string, KeyValuePair<string, object>[]> Tags { get; } = new();
+
+        public void AssertMetric(string metricName, long expected)
+        {
+            if (expected == 0)
+            {
+                Assert.False(ReportedMeters.ContainsKey(metricName), $"Should not have '{metricName}' metric reported.");
+            }
+            else
+            {
+                Assert.True(ReportedMeters.ContainsKey(metricName), $"'{metricName}' metric was not reported.");
+                Assert.AreEqual(expected, ReportedMeters[metricName]);
+            }
+        }
+
+        public object AssertTagKeyExists(string metricName, string tagKey)
+        {
+            if (!Tags.ContainsKey(metricName))
+            {
+                Assert.Fail($"'{metricName}' metric was not reported");
+            }
+
+            var emptyTag = default(KeyValuePair<string, object>);
+            var meterTag = Tags[metricName].FirstOrDefault(t => t.Key == tagKey);
+            if (meterTag.Equals(emptyTag))
+            {
+                Assert.Fail($"'{tagKey}' tag was not found.");
+            }
+
+            return meterTag.Value;
+        }
+    }
+}

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/When_ambient_trace_in_message_session.cs

@@ -0,0 +1,78 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry
+{
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Threading.Tasks;
+    using NServiceBus.AcceptanceTesting;
+    using NServiceBus.AcceptanceTests.EndpointTemplates;
+    using NUnit.Framework;
+
+    public class When_ambient_trace_in_message_session : OpenTelemetryAcceptanceTest
+    {
+        [Test]
+        public async Task Should_attach_to_ambient_trace()
+        {
+            using var externalActivitySource = new ActivitySource("external trace source");
+            using var _ = TestingActivityListener.SetupDiagnosticListener(externalActivitySource.Name); // need to have a registered listener for activities to be created
+
+            const string wrapperActivityTraceState = "test trace state";
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<EndpointWithAmbientActivity>(b => b
+                    .When(async (s, ctx) =>
+                    {
+                        // Otherwise the activity is created with a hierarchical ID format on .NET Framework which resets the RootId once it's converted to W3C format in the send pipeline.
+                        var activityTraceContext = new ActivityContext(ActivityTraceId.CreateRandom(),
+                            ActivitySpanId.CreateRandom(), ActivityTraceFlags.Recorded);
+                        using (var wrapperActivity = externalActivitySource.StartActivity("ambient span", ActivityKind.Server, activityTraceContext))
+                        {
+                            wrapperActivity.TraceStateString = wrapperActivityTraceState;
+                            ctx.WrapperActivityId = wrapperActivity.Id;
+                            ctx.WrapperActivityRootId = wrapperActivity.RootId;
+                            await s.SendLocal(new LocalMessage());
+                        }
+                    }))
+                .Done(c => c.OutgoingMessageReceived)
+                .Run();
+
+            var outgoingMessageActivity = NServicebusActivityListener.CompletedActivities.GetSendMessageActivities().Single();
+            var incomingMessageActivity = NServicebusActivityListener.CompletedActivities.GetReceiveMessageActivities().Single();
+
+            Assert.AreEqual(context.WrapperActivityId, outgoingMessageActivity.ParentId, "outgoing message should be connected to the ambient span");
+            Assert.AreEqual(context.WrapperActivityRootId, outgoingMessageActivity.RootId, "outgoing message should be connected to the ambient trace");
+            Assert.AreEqual(wrapperActivityTraceState, outgoingMessageActivity.TraceStateString, "ambient trace state should be floated to outgoing message span");
+            Assert.AreEqual(outgoingMessageActivity.Id, incomingMessageActivity.ParentId, "received message should be connected to send operation span");
+            Assert.AreEqual(context.WrapperActivityRootId, incomingMessageActivity.RootId, "received message should be connected to the ambient trace");
+            Assert.AreEqual(wrapperActivityTraceState, incomingMessageActivity.TraceStateString, "ambient trace state should be floated to incoming message span");
+        }
+
+        class Context : ScenarioContext
+        {
+            public bool OutgoingMessageReceived { get; set; }
+            public string WrapperActivityId { get; set; }
+            public string WrapperActivityRootId { get; set; }
+        }
+
+        class EndpointWithAmbientActivity : EndpointConfigurationBuilder
+        {
+            public EndpointWithAmbientActivity() => EndpointSetup<DefaultServer>();
+
+            public class MessageHandler : IHandleMessages<LocalMessage>
+            {
+                Context testContext;
+
+                public MessageHandler(Context testContext) => this.testContext = testContext;
+
+                public Task Handle(LocalMessage message, IMessageHandlerContext context)
+                {
+                    testContext.OutgoingMessageReceived = true;
+                    return Task.CompletedTask;
+                }
+            }
+        }
+
+        public class LocalMessage : IMessage
+        {
+        }
+    }
+}

src/NServiceBus.AcceptanceTests/Core/OpenTelemetry/When_ambient_trace_in_pipeline.cs

@@ -0,0 +1,87 @@
+namespace NServiceBus.AcceptanceTests.Core.OpenTelemetry
+{
+    using System;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Threading.Tasks;
+    using NServiceBus.AcceptanceTesting;
+    using NServiceBus.AcceptanceTests.EndpointTemplates;
+    using NUnit.Framework;
+
+    public class When_ambient_trace_in_pipeline : OpenTelemetryAcceptanceTest
+    {
+        static ActivitySource externalActivitySource = new(Guid.NewGuid().ToString());
+
+        [Test]
+        public async Task Should_attach_to_ambient_trace()
+        {
+            using var _ = TestingActivityListener.SetupDiagnosticListener(externalActivitySource.Name); // need to have a registered listener for activities to be created
+
+            var context = await Scenario.Define<Context>()
+                .WithEndpoint<EndpointWithAmbientActivity>(e => e
+                    .When(s => s.SendLocal(new TriggerMessage())))
+                .Done(c => c.MessageReceived)
+                .Run();
+
+            var handlerActivity = NServicebusActivityListener.CompletedActivities.GetInvokedHandlerActivities().First();
+            var sendFromHandlerActivity = NServicebusActivityListener.CompletedActivities.GetSendMessageActivities().Last();
+            Assert.AreEqual(context.AmbientActivityId, sendFromHandlerActivity.ParentId, "the outgoing message should be connected to the ambient span");
+            Assert.AreEqual(context.AmbientActivityRootId, sendFromHandlerActivity.RootId, "outgoing and ambient activity should belong to same trace");
+            Assert.AreEqual(ExpectedTraceState, sendFromHandlerActivity.TraceStateString, "outgoing activity should capture ambient trace state");
+            Assert.AreEqual(handlerActivity.Id, context.AmbientActivityParentId, "the ambient activity should be connected to the handler span");
+            Assert.AreEqual(handlerActivity.RootId, context.AmbientActivityRootId, "handler and ambient activity should belong to same trace");
+        }
+
+        class Context : ScenarioContext
+        {
+            public bool MessageReceived { get; set; }
+            public string AmbientActivityId { get; set; }
+            public string AmbientActivityParentId { get; set; }
+            public string AmbientActivityRootId { get; set; }
+            public string AmbientActivityState { get; set; }
+        }
+
+        class EndpointWithAmbientActivity : EndpointConfigurationBuilder
+        {
+            public EndpointWithAmbientActivity() => EndpointSetup<DefaultServer>();
+
+            class MessageHandler : IHandleMessages<TriggerMessage>, IHandleMessages<MessageFromAmbientTrace>
+            {
+                Context testContext;
+
+                public MessageHandler(Context testContext) => this.testContext = testContext;
+
+                public async Task Handle(TriggerMessage message, IMessageHandlerContext context)
+                {
+                    using (var ambientActivity = externalActivitySource.StartActivity())
+                    {
+                        // set/modify trace state:
+                        ambientActivity.TraceStateString = ExpectedTraceState;
+
+                        testContext.AmbientActivityId = ambientActivity.Id;
+                        testContext.AmbientActivityParentId = ambientActivity.ParentId;
+                        testContext.AmbientActivityRootId = ambientActivity.RootId;
+                        testContext.AmbientActivityState = ambientActivity.TraceStateString;
+                        await context.SendLocal(new MessageFromAmbientTrace());
+                    }
+                }
+
+                public Task Handle(MessageFromAmbientTrace message, IMessageHandlerContext context)
+                {
+                    testContext.MessageReceived = true;
+                    return Task.CompletedTask;
+                }
+            }
+        }
+
+        public class TriggerMessage : IMessage
+        {
+        }
+
+        public class MessageFromAmbientTrace : IMessage
+        {
+        }
+
+        const string ExpectedTraceState = "trace state from ambient activity";
+    }
+}