Synchronize OperationInternal calls

sdk/core/Azure.Core/src/Shared/OperationInternal.cs

@@ -42,6 +42,9 @@ namespace Azure.Core
     /// </summary>
     internal class OperationInternal : OperationInternalBase
     {
+        // To minimize code duplication and avoid introduction of another type,
+        // OperationInternal delegates implementation to the OperationInternal<VoidValue>.
+        // VoidValue is a private empty struct which only purpose is to be used as generic parameter.
         private readonly OperationInternal<VoidValue> _internalOperation;
 
         /// <summary>
@@ -89,7 +92,7 @@ public OperationInternal(
             DelayStrategy? fallbackStrategy = null)
             :base(clientDiagnostics, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy)
         {
-            _internalOperation = new OperationInternal<VoidValue>(clientDiagnostics, new OperationWrapper(operation), rawResponse, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy);
+            _internalOperation = new OperationInternal<VoidValue>(clientDiagnostics, new OperationToOperationOfTProxy(operation), rawResponse, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy);
         }
 
         private OperationInternal(OperationState finalState)
@@ -109,23 +112,30 @@ protected override async ValueTask<Response> UpdateStatusAsync(bool async, Cance
 
         private readonly struct VoidValue { }
 
-        private readonly struct OperationWrapper : IOperation<VoidValue>
+        // Wrapper type that converts OperationState to OperationState<T> and can be passed to `OperationInternal<T>` constructor.
+        private class OperationToOperationOfTProxy : IOperation<VoidValue>
         {
             private readonly IOperation _operation;
 
-            public OperationWrapper(IOperation operation)
+            public OperationToOperationOfTProxy(IOperation operation)
             {
                 _operation = operation;
             }
 
             public async ValueTask<OperationState<VoidValue>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
             {
                 var state = await _operation.UpdateStateAsync(async, cancellationToken).ConfigureAwait(false);
-                return state.HasCompleted
-                    ? state.HasSucceeded
-                        ? OperationState<VoidValue>.Success(state.RawResponse, new VoidValue())
-                        : OperationState<VoidValue>.Failure(state.RawResponse, state.OperationFailedException)
-                    : OperationState<VoidValue>.Pending(state.RawResponse);
+                if (!state.HasCompleted)
+                {
+                    return OperationState<VoidValue>.Pending(state.RawResponse);
+                }
+
+                if (state.HasSucceeded)
+                {
+                    return OperationState<VoidValue>.Success(state.RawResponse, new VoidValue());
+                }
+
+                return OperationState<VoidValue>.Failure(state.RawResponse, state.OperationFailedException);
             }
         }
     }

sdk/core/Azure.Core/src/Shared/OperationInternalBase.cs

@@ -180,18 +180,20 @@ public Response WaitForCompletionResponse(TimeSpan pollingInterval, Cancellation
 
         private async ValueTask<Response> WaitForCompletionResponseAsync(bool async, TimeSpan? pollingInterval, CancellationToken cancellationToken)
         {
-            using var asyncLock = await _responseLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
-            if (asyncLock.HasValue)
+            // If _responseLock has the value, lockOrValue will return 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
+            using var lockOrValue = await _responseLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
+            if (lockOrValue.HasValue)
             {
-                return asyncLock.Value;
+                return lockOrValue.Value;
             }
 
             var poller = new OperationPoller(_fallbackStrategy);
             var response = async
                 ? await poller.WaitForCompletionResponseAsync(this, pollingInterval, cancellationToken).ConfigureAwait(false)
                 : poller.WaitForCompletionResponse(this, pollingInterval, cancellationToken);
 
-            asyncLock.SetValue(response);
+            lockOrValue.SetValue(response);
             return response;
         }
 

sdk/core/Azure.Core/src/Shared/OperationInternalOfT.cs

@@ -109,6 +109,8 @@ public OperationInternal(
         private OperationInternal(OperationState<T> finalState)
             : base(finalState.RawResponse)
         {
+            // FinalOperation represents operation that is in final state and can't be updated.
+            // It implements IOperation<T> and throws exception when UpdateStateAsync is called.
             _operation = new FinalOperation();
             _rawResponse = finalState.RawResponse;
             _stateLock = new AsyncLockWithValue<OperationState<T>>(finalState);
@@ -249,6 +251,9 @@ public Response<T> WaitForCompletion(TimeSpan pollingInterval, CancellationToken
 
         protected override async ValueTask<Response> UpdateStatusAsync(bool async, CancellationToken cancellationToken)
         {
+            // If _stateLock has the final state, lockOrValue will contain that state, and no lock is acquired.
+            // If _stateLock doesn't have the state, GetLockOrValueAsync will acquire the lock that will be released when lockOrValue is disposed
+            // While _responseLock is used for the whole WaitForCompletionResponseAsync, _stateLock is used for individual calls of UpdateStatusAsync
             using var asyncLock = await _stateLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
             if (asyncLock.HasValue)
             {