Instrument events when they are enqueued in the buffered producer (Azure#38273)

* Instrument events when they are enqueued in the buffered producer

* Fix tests

* Fix tests

* Set W3C format on net462

* updated wrong test

* PR fb

* Fix spacing

Author: JoshLove-msft

sdk/eventhub/Azure.Messaging.EventHubs/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- When using the `EventHubBufferedProducerClient`, events are now instrumented when `EnqueueEventAsync` or `EnqueueEventsAsync` is called, rather than when the event is published. This ensures that the instrumentation is accurate when the event is published, regardless of whether the event is published immediately or buffered for a period of time.
+
 ### Other Changes
 
 ## 5.9.2 (2023-06-06)

sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs

@@ -14,6 +14,7 @@
 using Azure.Core;
 using Azure.Core.Diagnostics;
 using Azure.Core.Pipeline;
+using Azure.Core.Shared;
 using Azure.Messaging.EventHubs.Amqp;
 using Azure.Messaging.EventHubs.Core;
 using Azure.Messaging.EventHubs.Diagnostics;
@@ -128,6 +129,9 @@ public class EventHubBufferedProducerClient : IAsyncDisposable
         /// <summary>Indicates whether or not this instance has been closed.</summary>
         private volatile bool _isClosed;
 
+        /// <summary>The client diagnostics instance used to instrument events when enqueueing.</summary>
+        private readonly MessagingClientDiagnostics _clientDiagnostics;
+
         /// <summary>
         ///   The fully qualified Event Hubs namespace that this producer is currently associated with, which will likely be similar
         ///   to <c>{yournamespace}.servicebus.windows.net</c>.
@@ -459,6 +463,13 @@ public EventHubBufferedProducerClient(string connectionString,
         {
             Argument.AssertNotNullOrEmpty(connectionString, nameof(connectionString));
             _producer = new EventHubProducerClient(connectionString, eventHubName, (clientOptions ?? DefaultOptions).ToEventHubProducerClientOptions());
+
+            _clientDiagnostics = new MessagingClientDiagnostics(
+                DiagnosticProperty.DiagnosticNamespace,
+                DiagnosticProperty.ResourceProviderNamespace,
+                DiagnosticProperty.EventHubsServiceContext,
+                _producer.FullyQualifiedNamespace,
+                _producer.EventHubName);
         }
 
         /// <summary>
@@ -520,6 +531,13 @@ public EventHubBufferedProducerClient(EventHubConnection connection,
                                               EventHubBufferedProducerClientOptions clientOptions = default) : this(clientOptions)
         {
             _producer = new EventHubProducerClient(connection, (clientOptions ?? DefaultOptions).ToEventHubProducerClientOptions());
+
+            _clientDiagnostics = new MessagingClientDiagnostics(
+                DiagnosticProperty.DiagnosticNamespace,
+                DiagnosticProperty.ResourceProviderNamespace,
+                DiagnosticProperty.EventHubsServiceContext,
+                _producer.FullyQualifiedNamespace,
+                _producer.EventHubName);
         }
 
         /// <summary>
@@ -538,6 +556,12 @@ internal EventHubBufferedProducerClient(EventHubProducerClient producer,
                                                 EventHubBufferedProducerClientOptions clientOptions = default) : this(clientOptions)
         {
             _producer = producer;
+            _clientDiagnostics = new MessagingClientDiagnostics(
+                DiagnosticProperty.DiagnosticNamespace,
+                DiagnosticProperty.ResourceProviderNamespace,
+                DiagnosticProperty.EventHubsServiceContext,
+                _producer.FullyQualifiedNamespace,
+                _producer.EventHubName);
         }
 
         /// <summary>
@@ -575,6 +599,13 @@ private EventHubBufferedProducerClient(string fullyQualifiedNamespace,
                 AzureNamedKeyCredential keyCred =>  new EventHubProducerClient(fullyQualifiedNamespace, eventHubName, keyCred, options),
                 _ => throw new ArgumentException(Resources.UnsupportedCredential, nameof(credential))
             };
+
+            _clientDiagnostics = new MessagingClientDiagnostics(
+                DiagnosticProperty.DiagnosticNamespace,
+                DiagnosticProperty.ResourceProviderNamespace,
+                DiagnosticProperty.EventHubsServiceContext,
+                _producer.FullyQualifiedNamespace,
+                _producer.EventHubName);
         }
 
         /// <summary>
@@ -789,6 +820,7 @@ public virtual async Task<int> EnqueueEventAsync(EventData eventData,
                 var partitionState = _activePartitionStateMap.GetOrAdd(partitionId, partitionId => new PartitionPublishingState(partitionId, _options));
                 var writer = partitionState.PendingEventsWriter;
 
+                _clientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out _, out _);
                 await writer.WriteAsync(eventData, cancellationToken).ConfigureAwait(false);
 
                 var count = Interlocked.Increment(ref _totalBufferedEventCount);
@@ -968,6 +1000,7 @@ public virtual async Task<int> EnqueueEventsAsync(IEnumerable<EventData> events,
                     var activePartitionState = partitionState ?? _activePartitionStateMap.GetOrAdd(eventPartitionId, partitionId => new PartitionPublishingState(eventPartitionId, _options));
                     var writer = activePartitionState.PendingEventsWriter;
 
+                    _clientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out _, out _);
                     await writer.WriteAsync(eventData, cancellationToken).ConfigureAwait(false);
 
                     var count = Interlocked.Increment(ref _totalBufferedEventCount);

sdk/eventhub/Azure.Messaging.EventHubs/tests/Diagnostics/DiagnosticsTests.cs

@@ -187,6 +187,136 @@ await producer.SendAsync(new[]
             }
         }
 
+        /// <summary>
+        ///   Verifies diagnostics functionality of the <see cref="EventHubBufferedProducerClient" />
+        ///   class.
+        /// </summary>
+        ///
+        [Test]
+        public async Task EventHubBufferedProducerAppliesDiagnosticIdToEventsOnEnqueue()
+        {
+            using var testListener = new ClientDiagnosticListener(DiagnosticProperty.DiagnosticNamespace);
+            Activity activity = new Activity("SomeActivity");
+#if NET462
+            activity.SetIdFormat(ActivityIdFormat.W3C);
+#endif
+            activity.Start();
+
+            var eventHubName = "SomeName";
+            var endpoint = "some.endpoint.com";
+            var fakeConnection = new MockConnection(endpoint, eventHubName);
+
+            var writtenEventsData = new List<EventData>();
+
+            var mockProducer = new Mock<EventHubProducerClient>(fakeConnection, new EventHubProducerClientOptions());
+            mockProducer
+                .Setup(producer => producer.GetPartitionIdsAsync(It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new[] { "0"});
+            int times = 0;
+            mockProducer
+                .Setup(producer => producer.CreateBatchAsync(It.IsAny<CreateBatchOptions>(), It.IsAny<CancellationToken>()))
+                .Returns<CreateBatchOptions, CancellationToken>((options, token) =>
+                {
+                    List<EventData> events = new List<EventData>();
+                    if (times++ == 0)
+                    {
+                        events = writtenEventsData;
+                    }
+
+                    return new ValueTask<EventDataBatch>(
+                        EventHubsModelFactory.EventDataBatch(long.MaxValue, events, options));
+                });
+
+            mockProducer
+                .Setup(m => m.SendAsync(It.IsAny<EventDataBatch>(), It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask);
+
+            var bufferedProducer = new EventHubBufferedProducerClient(mockProducer.Object);
+
+            bufferedProducer.SendEventBatchFailedAsync += args => Task.CompletedTask;
+            await bufferedProducer.EnqueueEventsAsync(new[]
+            {
+                new EventData(ReadOnlyMemory<byte>.Empty),
+                new EventData(ReadOnlyMemory<byte>.Empty)
+            });
+
+            // Stop the activity before flushing so that we verify that the events are instrumented at time of enqueueing
+            activity.Stop();
+
+            await bufferedProducer.FlushAsync();
+            Assert.That(writtenEventsData.Count, Is.EqualTo(2), "All events should have been instrumented.");
+
+            foreach (EventData eventData in writtenEventsData)
+            {
+                Assert.That(eventData.Properties.TryGetValue(MessagingClientDiagnostics.DiagnosticIdAttribute, out object value), Is.True, "The events should have a diagnostic identifier property.");
+                ActivityContext.TryParse((string)value, null, out var context);
+                Assert.That(context.TraceId, Is.EqualTo(activity.TraceId), "The trace identifier should match the activity in the active scope when the events were enqueued.");
+            }
+        }
+
+        /// <summary>
+        ///   Verifies diagnostics functionality of the <see cref="EventHubBufferedProducerClient" />
+        ///   class when enqueueing a single event.
+        /// </summary>
+        ///
+        [Test]
+        public async Task EventHubBufferedProducerAppliesDiagnosticIdToSingleEventOnEnqueue()
+        {
+            using var testListener = new ClientDiagnosticListener(DiagnosticProperty.DiagnosticNamespace);
+            Activity activity = new Activity("SomeActivity");
+#if NET462
+            activity.SetIdFormat(ActivityIdFormat.W3C);
+#endif
+            activity.Start();
+
+            var eventHubName = "SomeName";
+            var endpoint = "some.endpoint.com";
+            var fakeConnection = new MockConnection(endpoint, eventHubName);
+
+            var writtenEventsData = new List<EventData>();
+
+            var mockProducer = new Mock<EventHubProducerClient>(fakeConnection, new EventHubProducerClientOptions());
+            mockProducer
+                .Setup(producer => producer.GetPartitionIdsAsync(It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new[] { "0"});
+            int times = 0;
+            mockProducer
+                .Setup(producer => producer.CreateBatchAsync(It.IsAny<CreateBatchOptions>(), It.IsAny<CancellationToken>()))
+                .Returns<CreateBatchOptions, CancellationToken>((options, token) =>
+                {
+                    List<EventData> events = new List<EventData>();
+                    if (times++ == 0)
+                    {
+                        events = writtenEventsData;
+                    }
+
+                    return new ValueTask<EventDataBatch>(
+                        EventHubsModelFactory.EventDataBatch(long.MaxValue, events, options));
+                });
+
+            mockProducer
+                .Setup(m => m.SendAsync(It.IsAny<EventDataBatch>(), It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask);
+
+            var bufferedProducer = new EventHubBufferedProducerClient(mockProducer.Object);
+
+            bufferedProducer.SendEventBatchFailedAsync += args => Task.CompletedTask;
+            await bufferedProducer.EnqueueEventAsync(new EventData(ReadOnlyMemory<byte>.Empty));
+
+            // Stop the activity before flushing so that we verify that the events are instrumented at time of enqueueing
+            activity.Stop();
+
+            await bufferedProducer.FlushAsync();
+            Assert.That(writtenEventsData.Count, Is.EqualTo(1), "All events should have been instrumented.");
+
+            foreach (EventData eventData in writtenEventsData)
+            {
+                Assert.That(eventData.Properties.TryGetValue(MessagingClientDiagnostics.DiagnosticIdAttribute, out object value), Is.True, "The events should have a diagnostic identifier property.");
+                ActivityContext.TryParse((string)value, null, out var context);
+                Assert.That(context.TraceId, Is.EqualTo(activity.TraceId), "The trace identifier should match the activity in the active scope when the events were enqueued.");
+            }
+        }
+
         /// <summary>
         ///   Verifies diagnostics functionality of the <see cref="EventHubProducerClient" />
         ///   class.