Retain instrumentation after adding event to batch (Azure#35124)

JoshLove-msft · jsquire · web-flow · commit b2a2f3e281f5 · 2023-03-27T09:42:41.000-07:00
* Retain instrumentation after adding event to batch

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

Co-authored-by: Jesse Squire &lt;jesse.squire@gmail.com&gt;

* update tests

---------

Co-authored-by: Jesse Squire &lt;jesse.squire@gmail.com&gt;
diff --git a/sdk/core/Azure.Core/src/Shared/MessagingClientDiagnostics.cs b/sdk/core/Azure.Core/src/Shared/MessagingClientDiagnostics.cs
@@ -156,12 +156,10 @@ public static bool TryExtractTraceContext(IDictionary<string, object> properties
         /// </summary>
         /// <param name="properties">The dictionary of application message properties.</param>
         /// <param name="activityName">The activity name to use for the diagnostic scope.</param>
-        /// <param name="scopeCreated">Whether or not a new scope was created while instrumenting the message.</param>
         /// <param name="traceparent">The traceparent that was either added, or that already existed in the message properties.</param>
         /// <param name="tracestate">The tracestate that was either added, or that already existed in the message properties.</param>
-        public void InstrumentMessage(IDictionary<string, object> properties, string activityName, out bool scopeCreated, out string? traceparent, out string? tracestate)
+        public void InstrumentMessage(IDictionary<string, object> properties, string activityName, out string? traceparent, out string? tracestate)
         {
-            scopeCreated = false;
             traceparent = null;
             tracestate = null;
 
@@ -171,7 +169,6 @@ public void InstrumentMessage(IDictionary<string, object> properties, string act
                     activityName,
                     DiagnosticScope.ActivityKind.Producer);
                 messageScope.Start();
-                scopeCreated = true;
 
                 Activity activity = Activity.Current;
                 if (activity != null)
@@ -194,20 +191,5 @@ public void InstrumentMessage(IDictionary<string, object> properties, string act
                 TryExtractTraceContext(properties, out traceparent, out tracestate);
             }
         }
-
-        /// <summary>
-        ///   Resets the instrumentation associated with a properties bag.
-        /// </summary>
-        ///
-        /// <param name="properties">The properties to reset.</param>
-        public static void ResetEvent(IDictionary<string, object> properties)
-        {
-            properties.Remove(DiagnosticIdAttribute);
-            if (ActivityExtensions.SupportsActivitySource())
-            {
-                properties.Remove(TraceParent);
-                properties.Remove(TraceState);
-            }
-        }
     }
 }
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/CHANGELOG.md b/sdk/eventhub/Azure.Messaging.EventHubs/CHANGELOG.md
@@ -6,6 +6,8 @@
 
 ### Breaking Changes
 
+- If diagnostic tracing is enabled, diagnostic tracing information is retained on `EventData` instances when they are added to an `EventDataBatch`. This matches the existing behavior when sending events using the `SendEventsAsync` method that takes an `IEnumerable<EventData>`.
+
 ### Bugs Fixed
 
 - Changed the approach that the event processor uses to validate permissions on startup to ensure that it does not interrupt other processors already running by temporarily asserting ownership of a partition.
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventDataBatch.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventDataBatch.cs
@@ -180,34 +180,16 @@ public bool TryAdd(EventData eventData)
             {
                 AssertNotLocked();
 
-                var messageScopeCreated = false;
-                var traceparent = default(string);
+                ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out var traceparent, out var tracestate);
 
-                try
-                {
-                    ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out messageScopeCreated, out traceparent, out var tracestate);
-
-                    var added = InnerBatch.TryAdd(eventData);
-
-                    if ((added) && (traceparent != null))
-                    {
-                        EventDiagnosticIdentifiers.Add((traceparent, tracestate));
-                    }
+                var added = InnerBatch.TryAdd(eventData);
 
-                    return added;
-                }
-                finally
+                if ((added) && (traceparent != null))
                 {
-                    // If a new message scope was added when instrumenting the instance, the identifier was
-                    // added during this call.  If so, remove it so that the source event is not modified; the
-                    // instrumentation will have been captured by the batch's copy of the event, if it was accepted
-                    // into the batch.
-
-                    if ((messageScopeCreated) && (traceparent != null))
-                    {
-                        MessagingClientDiagnostics.ResetEvent(eventData.Properties);
-                    }
+                    EventDiagnosticIdentifiers.Add((traceparent, tracestate));
                 }
+
+                return added;
             }
         }
 
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubProducerClient.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubProducerClient.cs
@@ -883,7 +883,7 @@ private async Task SendInternalAsync(IReadOnlyCollection<EventData> events,
 
             foreach (var eventData in events)
             {
-                ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out _, out var traceparent, out var tracestate);
+                ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out var traceparent, out var tracestate);
 
                 if (traceparent != null)
                 {
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/tests/Diagnostics/DiagnosticsActivitySourceTests.cs b/sdk/eventhub/Azure.Messaging.EventHubs/tests/Diagnostics/DiagnosticsActivitySourceTests.cs
@@ -224,10 +224,7 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()
                 .Returns<EventData>(e =>
                 {
                     var hasSpace = writtenEventsData.Count <= 1;
-                    if (hasSpace)
-                    {
-                        writtenEventsData.Add(e.Clone());
-                    }
+                    writtenEventsData.Add(e.Clone());
                     return hasSpace;
                 });
 
@@ -254,15 +251,13 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()
             await producer.SendAsync(batch);
 
             activity.Stop();
-            Assert.That(writtenEventsData.Count, Is.EqualTo(2), "Each of the events in the batch should have been instrumented.");
+            Assert.That(writtenEventsData.Count, Is.EqualTo(3), "Each of the events should have been instrumented when attempting to add them to the batch.");
 
             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.");
                 Assert.That(value, Is.EqualTo(activity.Id), "The diagnostics identifier should match the activity in the active scope.");
             }
-
-            Assert.That(eventData3.Properties.ContainsKey(MessagingClientDiagnostics.DiagnosticIdAttribute), Is.False, "Events that were not accepted into the batch should not have been instrumented.");
         }
 
         /// <summary>
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/tests/Diagnostics/DiagnosticsTests.cs b/sdk/eventhub/Azure.Messaging.EventHubs/tests/Diagnostics/DiagnosticsTests.cs
@@ -209,10 +209,7 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()
                 .Returns<EventData>(e =>
                 {
                     var hasSpace = writtenEventsData.Count <= 1;
-                    if (hasSpace)
-                    {
-                        writtenEventsData.Add(e.Clone());
-                    }
+                    writtenEventsData.Add(e.Clone());
                     return hasSpace;
                 });
 
@@ -239,15 +236,13 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()
             await producer.SendAsync(batch);
 
             activity.Stop();
-            Assert.That(writtenEventsData.Count, Is.EqualTo(2), "Each of the events in the batch should have been instrumented.");
+            Assert.That(writtenEventsData.Count, Is.EqualTo(3), "Each of the events should have been instrumented when attempting to add them to the batch.");
 
             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.");
                 Assert.That(value, Is.EqualTo(activity.Id), "The diagnostics identifier should match the activity in the active scope.");
             }
-
-            Assert.That(eventData3.Properties.ContainsKey(MessagingClientDiagnostics.DiagnosticIdAttribute), Is.False, "Events that were not accepted into the batch should not have been instrumented.");
         }
 
         /// <summary>
diff --git a/sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusMessageBatch.cs b/sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusMessageBatch.cs
@@ -113,7 +113,7 @@ public bool TryAddMessage(ServiceBusMessage message)
             {
                 AssertNotLocked();
 
-                _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _, out var _);
+                _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _);
                 return _innerBatch.TryAddMessage(message);
             }
         }
diff --git a/sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusSender.cs b/sdk/servicebus/Azure.Messaging.ServiceBus/src/Sender/ServiceBusSender.cs
@@ -266,7 +266,7 @@ private DiagnosticScope CreateDiagnosticScope(IReadOnlyCollection<ServiceBusMess
         {
             foreach (ServiceBusMessage message in messages)
             {
-                _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _, out var _);
+                _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _);
             }
 
             // create a new scope for the specified operation

Original file line number	Diff line number	Diff line change
`@@ -883,7 +883,7 @@ private async Task SendInternalAsync(IReadOnlyCollection<EventData> events,`
`883`	`883`
`884`	`884`	`foreach (var eventData in events)`
`885`	`885`	`{`
`886`		`- ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out _, out var traceparent, out var tracestate);`
	`886`	`+ ClientDiagnostics.InstrumentMessage(eventData.Properties, DiagnosticProperty.EventActivityName, out var traceparent, out var tracestate);`
`887`	`887`
`888`	`888`	`if (traceparent != null)`
`889`	`889`	`{`
Original file line number	Diff line number	Diff line change
`@@ -224,10 +224,7 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()`
`224`	`224`	`.Returns<EventData>(e =>`
`225`	`225`	`{`
`226`	`226`	`var hasSpace = writtenEventsData.Count <= 1;`
`227`		`- if (hasSpace)`
`228`		`- {`
`229`		`- writtenEventsData.Add(e.Clone());`
`230`		`- }`
	`227`	`+ writtenEventsData.Add(e.Clone());`
`231`	`228`	`return hasSpace;`
`232`	`229`	`});`
`233`	`230`
`@@ -254,15 +251,13 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()`
`254`	`251`	`await producer.SendAsync(batch);`
`255`	`252`
`256`	`253`	`activity.Stop();`
`257`		`- Assert.That(writtenEventsData.Count, Is.EqualTo(2), "Each of the events in the batch should have been instrumented.");`
	`254`	`+ Assert.That(writtenEventsData.Count, Is.EqualTo(3), "Each of the events should have been instrumented when attempting to add them to the batch.");`
`258`	`255`
`259`	`256`	`foreach (EventData eventData in writtenEventsData)`
`260`	`257`	`{`
`261`	`258`	`Assert.That(eventData.Properties.TryGetValue(MessagingClientDiagnostics.DiagnosticIdAttribute, out object value), Is.True, "The events should have a diagnostic identifier property.");`
`262`	`259`	`Assert.That(value, Is.EqualTo(activity.Id), "The diagnostics identifier should match the activity in the active scope.");`
`263`	`260`	`}`
`264`		`-`
`265`		`- Assert.That(eventData3.Properties.ContainsKey(MessagingClientDiagnostics.DiagnosticIdAttribute), Is.False, "Events that were not accepted into the batch should not have been instrumented.");`
`266`	`261`	`}`
`267`	`262`
`268`	`263`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -209,10 +209,7 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()`
`209`	`209`	`.Returns<EventData>(e =>`
`210`	`210`	`{`
`211`	`211`	`var hasSpace = writtenEventsData.Count <= 1;`
`212`		`- if (hasSpace)`
`213`		`- {`
`214`		`- writtenEventsData.Add(e.Clone());`
`215`		`- }`
	`212`	`+ writtenEventsData.Add(e.Clone());`
`216`	`213`	`return hasSpace;`
`217`	`214`	`});`
`218`	`215`
`@@ -239,15 +236,13 @@ public async Task EventHubProducerAppliesDiagnosticIdToEventsOnBatchSend()`
`239`	`236`	`await producer.SendAsync(batch);`
`240`	`237`
`241`	`238`	`activity.Stop();`
`242`		`- Assert.That(writtenEventsData.Count, Is.EqualTo(2), "Each of the events in the batch should have been instrumented.");`
	`239`	`+ Assert.That(writtenEventsData.Count, Is.EqualTo(3), "Each of the events should have been instrumented when attempting to add them to the batch.");`
`243`	`240`
`244`	`241`	`foreach (EventData eventData in writtenEventsData)`
`245`	`242`	`{`
`246`	`243`	`Assert.That(eventData.Properties.TryGetValue(MessagingClientDiagnostics.DiagnosticIdAttribute, out object value), Is.True, "The events should have a diagnostic identifier property.");`
`247`	`244`	`Assert.That(value, Is.EqualTo(activity.Id), "The diagnostics identifier should match the activity in the active scope.");`
`248`	`245`	`}`
`249`		`-`
`250`		`- Assert.That(eventData3.Properties.ContainsKey(MessagingClientDiagnostics.DiagnosticIdAttribute), Is.False, "Events that were not accepted into the batch should not have been instrumented.");`
`251`	`246`	`}`
`252`	`247`
`253`	`248`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -113,7 +113,7 @@ public bool TryAddMessage(ServiceBusMessage message)`
`113`	`113`	`{`
`114`	`114`	`AssertNotLocked();`
`115`	`115`
`116`		`- _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _, out var _);`
	`116`	`+ _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _);`
`117`	`117`	`return _innerBatch.TryAddMessage(message);`
`118`	`118`	`}`
`119`	`119`	`}`
Original file line number	Diff line number	Diff line change
`@@ -266,7 +266,7 @@ private DiagnosticScope CreateDiagnosticScope(IReadOnlyCollection<ServiceBusMess`
`266`	`266`	`{`
`267`	`267`	`foreach (ServiceBusMessage message in messages)`
`268`	`268`	`{`
`269`		`- _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _, out var _);`
	`269`	`+ _clientDiagnostics.InstrumentMessage(message.ApplicationProperties, DiagnosticProperty.MessageActivityName, out var _, out var _);`
`270`	`270`	`}`
`271`	`271`
`272`	`272`	`// create a new scope for the specified operation`