Add scope of waitForCompletion in LRO (#29033)

pshao25 · Pan Shao · web-flow · commit 465aa616557f · 2022-06-22T11:04:07.000+08:00
* Add scope of waitForCompletion in LRO

* Change flag to string

* Resolve comments

* Resolve comments

* Add tests

* Resolve comments

* Fix tests

* Fix test

* Fix test

* Resolve comments

* Resolve comments

* Make scope name not nullable.

* Resolve comments

* Remove useless method

* Resolve comment

Co-authored-by: Pan Shao &lt;pashao@microsoft.com&gt;
diff --git a/sdk/core/Azure.Core/src/Shared/OperationInternalBase.cs b/sdk/core/Azure.Core/src/Shared/OperationInternalBase.cs
@@ -15,15 +15,20 @@ namespace Azure.Core
     internal abstract class OperationInternalBase
     {
         private readonly ClientDiagnostics _diagnostics;
-        private readonly string _updateStatusScopeName;
         private readonly IReadOnlyDictionary<string, string>? _scopeAttributes;
         private readonly DelayStrategy? _fallbackStrategy;
         private readonly AsyncLockWithValue<Response> _responseLock;
 
+        private readonly string _waitForCompletionResponseScopeName;
+        protected readonly string _updateStatusScopeName;
+        protected readonly string _waitForCompletionScopeName;
+
         protected OperationInternalBase(Response rawResponse)
         {
             _diagnostics = new ClientDiagnostics(ClientOptions.Default);
             _updateStatusScopeName = string.Empty;
+            _waitForCompletionResponseScopeName = string.Empty;
+            _waitForCompletionScopeName = string.Empty;
             _scopeAttributes = default;
             _fallbackStrategy = default;
             _responseLock = new AsyncLockWithValue<Response>(rawResponse);
@@ -32,7 +37,9 @@ protected OperationInternalBase(Response rawResponse)
         protected OperationInternalBase(ClientDiagnostics clientDiagnostics, string operationTypeName, IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null, DelayStrategy? fallbackStrategy = null)
         {
             _diagnostics = clientDiagnostics;
-            _updateStatusScopeName = $"{operationTypeName}.UpdateStatus";
+            _updateStatusScopeName = $"{operationTypeName}.{nameof(UpdateStatus)}";
+            _waitForCompletionResponseScopeName = $"{operationTypeName}.{nameof(WaitForCompletionResponse)}";
+            _waitForCompletionScopeName = $"{operationTypeName}.WaitForCompletion";
             _scopeAttributes = scopeAttributes?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
             _fallbackStrategy = fallbackStrategy;
             _responseLock = new AsyncLockWithValue<Response>();
@@ -114,7 +121,7 @@ public Response UpdateStatus(CancellationToken cancellationToken) =>
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public async ValueTask<Response> WaitForCompletionResponseAsync(CancellationToken cancellationToken)
-            => await WaitForCompletionResponseAsync(async: true, null, cancellationToken).ConfigureAwait(false);
+            => await WaitForCompletionResponseAsync(async: true, null, _waitForCompletionResponseScopeName, cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Periodically calls <see cref="UpdateStatusAsync(CancellationToken)"/> until the long-running operation completes. The interval
@@ -135,7 +142,7 @@ public async ValueTask<Response> WaitForCompletionResponseAsync(CancellationToke
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public async ValueTask<Response> WaitForCompletionResponseAsync(TimeSpan pollingInterval, CancellationToken cancellationToken)
-            => await WaitForCompletionResponseAsync(async: true, pollingInterval, cancellationToken).ConfigureAwait(false);
+            => await WaitForCompletionResponseAsync(async: true, pollingInterval, _waitForCompletionResponseScopeName, cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Periodically calls <see cref="UpdateStatus(CancellationToken)"/> until the long-running operation completes.
@@ -155,7 +162,7 @@ public async ValueTask<Response> WaitForCompletionResponseAsync(TimeSpan polling
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public Response WaitForCompletionResponse(CancellationToken cancellationToken)
-            => WaitForCompletionResponseAsync(async: false, null, cancellationToken).EnsureCompleted();
+            => WaitForCompletionResponseAsync(async: false, null, _waitForCompletionResponseScopeName, cancellationToken).EnsureCompleted();
 
         /// <summary>
         /// Periodically calls <see cref="UpdateStatus(CancellationToken)"/> until the long-running operation completes. The interval
@@ -176,9 +183,9 @@ public Response WaitForCompletionResponse(CancellationToken cancellationToken)
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public Response WaitForCompletionResponse(TimeSpan pollingInterval, CancellationToken cancellationToken)
-            => WaitForCompletionResponseAsync(async: false, pollingInterval, cancellationToken).EnsureCompleted();
+            => WaitForCompletionResponseAsync(async: false, pollingInterval, _waitForCompletionResponseScopeName, cancellationToken).EnsureCompleted();
 
-        private async ValueTask<Response> WaitForCompletionResponseAsync(bool async, TimeSpan? pollingInterval, CancellationToken cancellationToken)
+        protected async ValueTask<Response> WaitForCompletionResponseAsync(bool async, TimeSpan? pollingInterval, string scopeName, CancellationToken cancellationToken)
         {
             // If _responseLock has the value, lockOrValue will contain that value, and no lock is acquired.
             // If _responseLock doesn't have the value, GetLockOrValueAsync will acquire the lock that will be released when lockOrValue is disposed
@@ -188,20 +195,29 @@ private async ValueTask<Response> WaitForCompletionResponseAsync(bool async, Tim
                 return lockOrValue.Value;
             }
 
-            var poller = new OperationPoller(_fallbackStrategy);
-            var response = async
-                ? await poller.WaitForCompletionResponseAsync(this, pollingInterval, cancellationToken).ConfigureAwait(false)
-                : poller.WaitForCompletionResponse(this, pollingInterval, cancellationToken);
+            using var scope = CreateScope(scopeName);
+            try
+            {
+                var poller = new OperationPoller(_fallbackStrategy);
+                var response = async
+                    ? await poller.WaitForCompletionResponseAsync(this, pollingInterval, cancellationToken).ConfigureAwait(false)
+                    : poller.WaitForCompletionResponse(this, pollingInterval, cancellationToken);
 
-            lockOrValue.SetValue(response);
-            return response;
+                lockOrValue.SetValue(response);
+                return response;
+            }
+            catch (Exception e)
+            {
+                scope.Failed(e);
+                throw;
+            }
         }
 
         protected abstract ValueTask<Response> UpdateStatusAsync(bool async, CancellationToken cancellationToken);
 
-        protected DiagnosticScope CreateScope()
+        protected DiagnosticScope CreateScope(string scopeName)
         {
-            DiagnosticScope scope = _diagnostics.CreateScope(_updateStatusScopeName);
+            DiagnosticScope scope = _diagnostics.CreateScope(scopeName);
 
             if (_scopeAttributes != null)
             {
diff --git a/sdk/core/Azure.Core/src/Shared/OperationInternalOfT.cs b/sdk/core/Azure.Core/src/Shared/OperationInternalOfT.cs
@@ -174,10 +174,7 @@ public T Value
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public async ValueTask<Response<T>> WaitForCompletionAsync(CancellationToken cancellationToken)
-        {
-            var rawResponse = await WaitForCompletionResponseAsync(cancellationToken).ConfigureAwait(false);
-            return Response.FromValue(Value, rawResponse);
-        }
+            => await WaitForCompletionAsync(async: true, null, cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Periodically calls <see cref="OperationInternalBase.UpdateStatusAsync(CancellationToken)"/> until the long-running operation completes. The interval
@@ -198,10 +195,7 @@ public async ValueTask<Response<T>> WaitForCompletionAsync(CancellationToken can
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public async ValueTask<Response<T>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken)
-        {
-            var rawResponse = await WaitForCompletionResponseAsync(pollingInterval, cancellationToken).ConfigureAwait(false);
-            return Response.FromValue(Value, rawResponse);
-        }
+            => await WaitForCompletionAsync(async: true, pollingInterval, cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Periodically calls <see cref="OperationInternalBase.UpdateStatus(CancellationToken)"/> until the long-running operation completes.
@@ -220,10 +214,7 @@ public async ValueTask<Response<T>> WaitForCompletionAsync(TimeSpan pollingInter
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public Response<T> WaitForCompletion(CancellationToken cancellationToken)
-        {
-            var rawResponse = WaitForCompletionResponse(cancellationToken);
-            return Response.FromValue(Value, rawResponse);
-        }
+            => WaitForCompletionAsync(async: false, null, cancellationToken).EnsureCompleted();
 
         /// <summary>
         /// Periodically calls <see cref="OperationInternalBase.UpdateStatus(CancellationToken)"/> until the long-running operation completes. The interval
@@ -244,8 +235,11 @@ public Response<T> WaitForCompletion(CancellationToken cancellationToken)
         /// <returns>The last HTTP response received from the server, including the final result of the long-running operation.</returns>
         /// <exception cref="RequestFailedException">Thrown if there's been any issues during the connection, or if the operation has completed with failures.</exception>
         public Response<T> WaitForCompletion(TimeSpan pollingInterval, CancellationToken cancellationToken)
+            => WaitForCompletionAsync(async: false, pollingInterval, cancellationToken).EnsureCompleted();
+
+        private async ValueTask<Response<T>> WaitForCompletionAsync(bool async, TimeSpan? pollingInterval, CancellationToken cancellationToken)
         {
-            var rawResponse = WaitForCompletionResponse(pollingInterval, cancellationToken);
+            var rawResponse = await WaitForCompletionResponseAsync(async, pollingInterval, _waitForCompletionScopeName, cancellationToken).ConfigureAwait(false);
             return Response.FromValue(Value, rawResponse);
         }
 
@@ -260,7 +254,7 @@ protected override async ValueTask<Response> UpdateStatusAsync(bool async, Cance
                 return GetResponseFromState(asyncLock.Value);
             }
 
-            using var scope = CreateScope();
+            using var scope = CreateScope(_updateStatusScopeName);
             try
             {
                 var state = await _operation.UpdateStateAsync(async, cancellationToken).ConfigureAwait(false);
diff --git a/sdk/core/Azure.Core/tests/OperationInternalOfTTests.cs b/sdk/core/Azure.Core/tests/OperationInternalOfTTests.cs
@@ -148,6 +148,19 @@ public async Task UpdateStatusCreatesDiagnosticScope([Values(true, false)] bool
             testListener.AssertScope($"{expectedTypeName}.UpdateStatus", expectedAttributes);
         }
 
+        [Test]
+        public async Task UpdateStatusNotCreateDiagnosticScope([Values(true, false)] bool async)
+        {
+            using ClientDiagnosticListener testListener = new(DiagnosticNamespace);
+
+            var operationInternal = OperationInternal<int>.Succeeded(InitialResponse, 1);
+            _ = async
+                ? await operationInternal.UpdateStatusAsync(CancellationToken.None)
+                : operationInternal.UpdateStatus(CancellationToken.None);
+
+            CollectionAssert.IsEmpty(testListener.Scopes);
+        }
+
         [Test]
         public async Task UpdateStatusSetsFailedScopeWhenOperationFails([Values(true, false)] bool async)
         {
@@ -205,6 +218,78 @@ public async Task UpdateStatusPassesTheCancellationTokenToUpdateState([Values(tr
             Assert.AreEqual(originalToken, passedToken);
         }
 
+        [Test]
+        public async Task WaitForCompletionResponseCreatesDiagnosticScope([Values(true, false)] bool async, [Values(null, "CustomTypeName")] string operationTypeName, [Values(true, false)] bool suppressNestedClientActivities)
+        {
+            using ClientDiagnosticListener testListener = new(DiagnosticNamespace);
+
+            string expectedTypeName = operationTypeName ?? nameof(TestOperation);
+            KeyValuePair<string, string>[] expectedAttributes = { new("key1", "value1"), new("key2", "value2") };
+            var operationInternal = new OperationInternal<int>(new(new TestClientOptions(), suppressNestedClientActivities), TestOperation.Succeeded(1), InitialResponse, operationTypeName, expectedAttributes);
+
+            _ = async
+                ? await operationInternal.WaitForCompletionResponseAsync(CancellationToken.None)
+                : operationInternal.WaitForCompletionResponse(CancellationToken.None);
+
+            testListener.AssertScope($"{expectedTypeName}.WaitForCompletionResponse", expectedAttributes);
+#if NET5_0_OR_GREATER
+            if (suppressNestedClientActivities)
+            {
+                testListener.AssertAndRemoveScope($"{expectedTypeName}.WaitForCompletionResponse", expectedAttributes);
+                CollectionAssert.IsEmpty(testListener.Scopes);
+            }
+#endif
+        }
+
+        [Test]
+        public async Task WaitForCompletionCreatesDiagnosticScope([Values(true, false)] bool async, [Values(null, "CustomTypeName")] string operationTypeName, [Values(true, false)] bool suppressNestedClientActivities)
+        {
+            using ClientDiagnosticListener testListener = new(DiagnosticNamespace);
+
+            string expectedTypeName = operationTypeName ?? nameof(TestOperation);
+            KeyValuePair<string, string>[] expectedAttributes = { new("key1", "value1"), new("key2", "value2") };
+            var operationInternal = new OperationInternal<int>(new(new TestClientOptions(), suppressNestedClientActivities), TestOperation.Succeeded(1), InitialResponse, operationTypeName, expectedAttributes);
+
+            _ = async
+                ? await operationInternal.WaitForCompletionAsync(CancellationToken.None)
+                : operationInternal.WaitForCompletion(CancellationToken.None);
+
+            testListener.AssertScope($"{expectedTypeName}.WaitForCompletion", expectedAttributes);
+#if NET5_0_OR_GREATER
+            if (suppressNestedClientActivities)
+            {
+                testListener.AssertAndRemoveScope($"{expectedTypeName}.WaitForCompletion", expectedAttributes);
+                CollectionAssert.IsEmpty(testListener.Scopes);
+            }
+#endif
+        }
+
+        [Test]
+        public async Task WaitForCompletionResponseNotCreateDiagnosticScope([Values(true, false)] bool async)
+        {
+            using ClientDiagnosticListener testListener = new(DiagnosticNamespace);
+
+            var operationInternal = OperationInternal<int>.Succeeded(InitialResponse, 1);
+            _ = async
+                ? await operationInternal.WaitForCompletionResponseAsync(CancellationToken.None)
+                : operationInternal.WaitForCompletionResponse(CancellationToken.None);
+
+            CollectionAssert.IsEmpty(testListener.Scopes);
+        }
+
+        [Test]
+        public async Task WaitForCompletionNotCreateDiagnosticScope([Values(true, false)] bool async)
+        {
+            using ClientDiagnosticListener testListener = new(DiagnosticNamespace);
+
+            var operationInternal = OperationInternal<int>.Succeeded(InitialResponse, 1);
+            _ = async
+                ? await operationInternal.WaitForCompletionAsync(CancellationToken.None)
+                : operationInternal.WaitForCompletion(CancellationToken.None);
+
+            CollectionAssert.IsEmpty(testListener.Scopes);
+        }
+
         [Test]
         public async Task WaitForCompletionCallsUntilOperationCompletes([Values(true, false)] bool useDefaultPollingInterval)
         {
diff --git a/sdk/core/Azure.Core/tests/OperationInternalTests.cs b/sdk/core/Azure.Core/tests/OperationInternalTests.cs
