[FormRecognizer] Updated LROs to use new OperationInternal class (#22629)

Author: kinelski

sdk/formrecognizer/Azure.AI.FormRecognizer/src/Azure.AI.FormRecognizer.csproj

@@ -25,7 +25,6 @@
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScope.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScopeFactory.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureCoreSharedSources)HttpMessageSanitizer.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
-    <Compile Include="$(AzureCoreSharedSources)OperationHelpers.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureCoreSharedSources)OperationInternal.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
     <Compile Include="$(AzureCoreSharedSources)TaskExtensions.cs" Link="Shared\%(RecursiveDir)\%(Filename)%(Extension)" />
   </ItemGroup>

sdk/formrecognizer/Azure.AI.FormRecognizer/src/CopyModelOperation.cs

@@ -14,23 +14,16 @@ namespace Azure.AI.FormRecognizer.Training
     /// <summary>
     /// Tracks the status of a long-running operation for copying a custom model into a target Form Recognizer resource.
     /// </summary>
-    public class CopyModelOperation : Operation<CustomFormModelInfo>
+    public class CopyModelOperation : Operation<CustomFormModelInfo>, IOperation<CustomFormModelInfo>
     {
+        private readonly OperationInternal<CustomFormModelInfo> _operationInternal;
+
         /// <summary>Provides communication with the Form Recognizer Azure Cognitive Service through its REST API.</summary>
         private readonly FormRecognizerRestClient _serviceClient;
 
         /// <summary>Provides tools for exception creation in case of failure.</summary>
         private readonly ClientDiagnostics _diagnostics;
 
-        /// <summary>The last HTTP response received from the server. <c>null</c> until the first response is received.</summary>
-        private Response _response;
-
-        /// <summary>The result of the long-running operation. <c>null</c> until result is received on status update.</summary>
-        private CustomFormModelInfo _value;
-
-        /// <summary><c>true</c> if the long-running operation has completed. Otherwise, <c>false</c>.</summary>
-        private bool _hasCompleted;
-
         /// <summary>The ID of the model to use for copy.</summary>
         private readonly string _modelId;
 
@@ -40,8 +33,6 @@ public class CopyModelOperation : Operation<CustomFormModelInfo>
         /// <summary>An ID representing the operation that can be used along with <see cref="_modelId"/> to poll for the status of the long-running operation.</summary>
         private readonly string _resultId;
 
-        private RequestFailedException _requestFailedException;
-
         /// <summary>
         /// Gets an ID representing the operation that can be used to poll for the status
         /// of the long-running operation.
@@ -54,28 +45,17 @@ public class CopyModelOperation : Operation<CustomFormModelInfo>
         /// <remarks>
         /// This property can be accessed only after the operation completes successfully (HasValue is true).
         /// </remarks>
-        public override CustomFormModelInfo Value
-        {
-            get
-            {
-                if (HasCompleted && !HasValue)
-#pragma warning disable CA1065 // Do not raise exceptions in unexpected locations
-                    throw _requestFailedException;
-#pragma warning restore CA1065 // Do not raise exceptions in unexpected locations
-                else
-                    return OperationHelpers.GetValue(ref _value);
-            }
-        }
+        public override CustomFormModelInfo Value => _operationInternal.Value;
 
         /// <summary>
         /// Returns true if the long-running operation completed.
         /// </summary>
-        public override bool HasCompleted => _hasCompleted;
+        public override bool HasCompleted => _operationInternal.HasCompleted;
 
         /// <summary>
         /// Returns true if the long-running operation completed successfully and has produced final result (accessible by Value property).
         /// </summary>
-        public override bool HasValue => _value != null;
+        public override bool HasValue => _operationInternal.HasValue;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="CopyModelOperation"/> class which
@@ -92,6 +72,7 @@ public CopyModelOperation(string operationId, string targetModelId, FormTraining
             _serviceClient = client.ServiceClient;
             _diagnostics = client.Diagnostics;
             _targetModelId = targetModelId;
+            _operationInternal = new(_diagnostics, this, rawResponse: null);
 
             // TODO: Use regex to parse ids.
             // https://github.com/Azure/azure-sdk-for-net/issues/11505
@@ -124,6 +105,7 @@ internal CopyModelOperation(FormRecognizerRestClient serviceClient, ClientDiagno
             _serviceClient = serviceClient;
             _diagnostics = diagnostics;
             _targetModelId = targetModelId;
+            _operationInternal = new(_diagnostics, this, rawResponse: null);
 
             // TODO: Use regex to parse ids.
             // https://github.com/Azure/azure-sdk-for-net/issues/11505
@@ -160,7 +142,7 @@ protected CopyModelOperation()
         /// An instance of <see cref="CopyModelOperation"/> sends requests to a server in UpdateStatusAsync, UpdateStatus, and other methods.
         /// Responses from these requests can be accessed using GetRawResponse.
         /// </remarks>
-        public override Response GetRawResponse() => _response;
+        public override Response GetRawResponse() => _operationInternal.RawResponse;
 
         /// <summary>
         /// Periodically calls the server till the long-running operation completes.
@@ -170,8 +152,8 @@ protected CopyModelOperation()
         /// <remarks>
         /// This method will periodically call UpdateStatusAsync till HasCompleted is true, then return the final result of the operation.
         /// </remarks>
-        public override ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(CancellationToken cancellationToken = default) =>
-            this.DefaultWaitForCompletionAsync(cancellationToken);
+        public override async ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(CancellationToken cancellationToken = default) =>
+            await _operationInternal.WaitForCompletionAsync(cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Periodically calls the server till the long-running operation completes.
@@ -186,8 +168,8 @@ public override ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(
         /// <remarks>
         /// This method will periodically call UpdateStatusAsync till HasCompleted is true, then return the final result of the operation.
         /// </remarks>
-        public override ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken = default) =>
-             this.DefaultWaitForCompletionAsync(pollingInterval, cancellationToken);
+        public override async ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(TimeSpan pollingInterval, CancellationToken cancellationToken = default) =>
+            await _operationInternal.WaitForCompletionAsync(pollingInterval, cancellationToken).ConfigureAwait(false);
 
         /// <summary>
         /// Calls the server to get updated status of the long-running operation.
@@ -198,7 +180,7 @@ public override ValueTask<Response<CustomFormModelInfo>> WaitForCompletionAsync(
         /// This operation will update the value returned from GetRawResponse and might update HasCompleted, HasValue, and Value.
         /// </remarks>
         public override Response UpdateStatus(CancellationToken cancellationToken = default) =>
-            UpdateStatusAsync(false, cancellationToken).EnsureCompleted();
+            _operationInternal.UpdateStatus(cancellationToken);
 
         /// <summary>
         /// Calls the server to get updated status of the long-running operation.
@@ -209,52 +191,32 @@ public override Response UpdateStatus(CancellationToken cancellationToken = defa
         /// This operation will update the value returned from GetRawResponse and might update HasCompleted, HasValue, and Value.
         /// </remarks>
         public override async ValueTask<Response> UpdateStatusAsync(CancellationToken cancellationToken = default) =>
-            await UpdateStatusAsync(true, cancellationToken).ConfigureAwait(false);
+            await _operationInternal.UpdateStatusAsync(cancellationToken).ConfigureAwait(false);
 
-        /// <summary>
-        /// Calls the server to get updated status of the long-running operation.
-        /// </summary>
-        /// <param name="async">When <c>true</c>, the method will be executed asynchronously; otherwise, it will execute synchronously.</param>
-        /// <param name="cancellationToken">A <see cref="CancellationToken"/> used for the service call.</param>
-        /// <returns>The HTTP response received from the server.</returns>
-        private async ValueTask<Response> UpdateStatusAsync(bool async, CancellationToken cancellationToken)
+        async ValueTask<OperationState<CustomFormModelInfo>> IOperation<CustomFormModelInfo>.UpdateStateAsync(bool async, CancellationToken cancellationToken)
         {
-            if (!_hasCompleted)
+            Response<CopyOperationResult> response = async
+                ? await _serviceClient.GetCustomModelCopyResultAsync(new Guid(_modelId), new Guid(_resultId), cancellationToken).ConfigureAwait(false)
+                : _serviceClient.GetCustomModelCopyResult(new Guid(_modelId), new Guid(_resultId), cancellationToken);
+
+            OperationStatus status = response.Value.Status;
+            Response rawResponse = response.GetRawResponse();
+
+            if (status == OperationStatus.Succeeded)
+            {
+                return OperationState<CustomFormModelInfo>.Success(rawResponse,
+                    ConvertValue(response.Value, _targetModelId, CustomFormModelStatus.Ready));
+            }
+            else if (status == OperationStatus.Failed)
             {
-                using DiagnosticScope scope = _diagnostics.CreateScope($"{nameof(CopyModelOperation)}.{nameof(UpdateStatus)}");
-                scope.Start();
-
-                try
-                {
-                    Response<CopyOperationResult> update = async
-                        ? await _serviceClient.GetCustomModelCopyResultAsync(new Guid(_modelId), new Guid(_resultId), cancellationToken).ConfigureAwait(false)
-                        : _serviceClient.GetCustomModelCopyResult(new Guid(_modelId), new Guid(_resultId), cancellationToken);
-
-                    _response = update.GetRawResponse();
-
-                    if (update.Value.Status == OperationStatus.Succeeded)
-                    {
-                        // We need to first assign a value and then mark the operation as completed to avoid a race condition with the getter in Value
-                        _value = ConvertValue(update.Value, _targetModelId, CustomFormModelStatus.Ready);
-                        _hasCompleted = true;
-                    }
-                    else if (update.Value.Status == OperationStatus.Failed)
-                    {
-                        _requestFailedException = await ClientCommon
-                            .CreateExceptionForFailedOperationAsync(async, _diagnostics, _response, update.Value.CopyResult.Errors)
-                            .ConfigureAwait(false);
-                        _hasCompleted = true;
-                        throw _requestFailedException;
-                    }
-                }
-                catch (Exception e)
-                {
-                    scope.Failed(e);
-                    throw;
-                }
+                RequestFailedException requestFailedException = await ClientCommon
+                    .CreateExceptionForFailedOperationAsync(async, _diagnostics, rawResponse, response.Value.CopyResult.Errors)
+                    .ConfigureAwait(false);
+
+                return OperationState<CustomFormModelInfo>.Failure(rawResponse, requestFailedException);
             }
 
-            return GetRawResponse();
+            return OperationState<CustomFormModelInfo>.Pending(rawResponse);
         }
 
         private static CustomFormModelInfo ConvertValue(CopyOperationResult result, string modelId, CustomFormModelStatus status)