Set the linked activity before starting the scope. (#20376)

* Set the linked activity before starting the scope. * Undo obsoletion
Azure · Apr 15, 2021 · f665705 · f665705
1 parent 8c0301c
commit f665705
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 143 deletions.
diff --git a/sdk/core/Azure.Core.TestFramework/src/ClientDiagnosticListener.cs b/sdk/core/Azure.Core.TestFramework/src/ClientDiagnosticListener.cs
@@ -16,20 +16,23 @@ public class ClientDiagnosticListener : IObserver<KeyValuePair<string, object>>,
         private readonly AsyncLocal<bool> _collectThisStack;
 
         private List<IDisposable> _subscriptions = new List<IDisposable>();
+        private readonly Action<ProducedDiagnosticScope> _scopeStartCallback;
 
         public List<ProducedDiagnosticScope> Scopes { get; } = new List<ProducedDiagnosticScope>();
 
-        public ClientDiagnosticListener(string name, bool asyncLocal = false): this(n => n == name, asyncLocal)
+        public ClientDiagnosticListener(string name, bool asyncLocal = false, Action<ProducedDiagnosticScope> scopeStartCallback = default)
+            : this(n => n == name, asyncLocal, scopeStartCallback)
         {
         }
 
-        public ClientDiagnosticListener(Func<string, bool> filter, bool asyncLocal = false)
+        public ClientDiagnosticListener(Func<string, bool> filter, bool asyncLocal = false, Action<ProducedDiagnosticScope> scopeStartCallback = default)
         {
             if (asyncLocal)
             {
                 _collectThisStack = new AsyncLocal<bool> { Value = true };
             }
             _sourceNameFilter = filter;
+            _scopeStartCallback = scopeStartCallback;
             DiagnosticListener.AllListeners.Subscribe(this);
         }
 
@@ -69,6 +72,7 @@ public void OnNext(KeyValuePair<string, object> value)
                     };
 
                     Scopes.Add(scope);
+                    _scopeStartCallback?.Invoke(scope);
                 }
                 else if (value.Key.EndsWith(stopSuffix))
                 {
@@ -144,7 +148,10 @@ public void Dispose()
             }
         }
 
-        public ProducedDiagnosticScope AssertScopeStarted(string name, params KeyValuePair<string, string>[] expectedAttributes)
+        public ProducedDiagnosticScope AssertScopeStarted(string name, params KeyValuePair<string, string>[] expectedAttributes) =>
+            AssertScopeStartedInternal(name, false, expectedAttributes);
+
+        private ProducedDiagnosticScope AssertScopeStartedInternal(string name, bool remove, params KeyValuePair<string, string>[] expectedAttributes)
         {
             lock (Scopes)
             {
@@ -160,16 +167,28 @@ public ProducedDiagnosticScope AssertScopeStarted(string name, params KeyValuePa
                             }
                         }
 
+                        if (remove)
+                        {
+                            Scopes.Remove(producedDiagnosticScope);
+                        }
+
                         return producedDiagnosticScope;
                     }
                 }
                 throw new InvalidOperationException($"Event '{name}' was not started");
             }
         }
 
-        public ProducedDiagnosticScope AssertScope(string name, params KeyValuePair<string, string>[] expectedAttributes)
+        public ProducedDiagnosticScope AssertScope(string name, params KeyValuePair<string, string>[] expectedAttributes) =>
+            AssertScopeInternal(name, false, expectedAttributes);
+
+        public ProducedDiagnosticScope AssertAndRemoveScope(string name, params KeyValuePair<string, string>[] expectedAttributes) =>
+            AssertScopeInternal(name, true, expectedAttributes);
+
+        private ProducedDiagnosticScope AssertScopeInternal(string name, bool remove,
+            params KeyValuePair<string, string>[] expectedAttributes)
         {
-            ProducedDiagnosticScope scope = AssertScopeStarted(name, expectedAttributes);
+            ProducedDiagnosticScope scope = AssertScopeStartedInternal(name, remove, expectedAttributes);
             if (!scope.IsCompleted)
             {
                 throw new InvalidOperationException($"'{name}' is not completed");

diff --git a/sdk/core/Azure.Core.TestFramework/src/TestDiagnosticListener.cs b/sdk/core/Azure.Core.TestFramework/src/TestDiagnosticListener.cs
@@ -7,6 +7,7 @@
 
 namespace Azure.Core.Tests
 {
+    // DO NOT USE - use ClientDiagnosticListener instead
     public class TestDiagnosticListener : IObserver<DiagnosticListener>, IDisposable
     {
         private readonly Func<DiagnosticListener, bool> _selector;

diff --git a/sdk/servicebus/Azure.Messaging.ServiceBus/src/Processor/ReceiverManager.cs b/sdk/servicebus/Azure.Messaging.ServiceBus/src/Processor/ReceiverManager.cs
@@ -116,8 +116,8 @@ await RaiseExceptionReceived(
         protected async Task ProcessOneMessageWithinScopeAsync(ServiceBusReceivedMessage message, string activityName, CancellationToken cancellationToken)
         {
             using DiagnosticScope scope = _scopeFactory.CreateScope(activityName, DiagnosticProperty.ConsumerKind);
-            scope.Start();
             scope.SetMessageData(new ServiceBusReceivedMessage[] { message });
+            scope.Start();
             try
             {
                 await ProcessOneMessage(