Add abstract UpdateStatus methods to OperationResult

sdk/core/System.ClientModel/CHANGELOG.md

@@ -13,6 +13,7 @@
 ### Other Changes
 
 - Removed `ReturnWhen` enum in favor of using bool `waitUntilCompleted` parameter in third-party client LRO method signatures.
+- Added abstract `UpdateStatus` method to `OperationResult`.
 
 ## 1.1.0-beta.6 (2024-08-01)
 

sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs

@@ -195,8 +195,10 @@ public abstract partial class OperationResult : System.ClientModel.ClientResult
         protected OperationResult(System.ClientModel.Primitives.PipelineResponse response) { }
         public abstract bool IsCompleted { get; protected set; }
         public abstract System.ClientModel.ContinuationToken? RehydrationToken { get; protected set; }
-        public abstract void WaitForCompletion(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));
-        public abstract System.Threading.Tasks.Task WaitForCompletionAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));
+        public abstract System.ClientModel.ClientResult UpdateStatus(System.ClientModel.Primitives.RequestOptions? options = null);
+        public abstract System.Threading.Tasks.ValueTask<System.ClientModel.ClientResult> UpdateStatusAsync(System.ClientModel.Primitives.RequestOptions? options = null);
+        public virtual void WaitForCompletion(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { }
+        public virtual System.Threading.Tasks.ValueTask WaitForCompletionAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
     [System.AttributeUsageAttribute(System.AttributeTargets.Class)]
     public sealed partial class PersistableModelProxyAttribute : System.Attribute

sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs

@@ -194,8 +194,10 @@ public abstract partial class OperationResult : System.ClientModel.ClientResult
         protected OperationResult(System.ClientModel.Primitives.PipelineResponse response) { }
         public abstract bool IsCompleted { get; protected set; }
         public abstract System.ClientModel.ContinuationToken? RehydrationToken { get; protected set; }
-        public abstract void WaitForCompletion(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));
-        public abstract System.Threading.Tasks.Task WaitForCompletionAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));
+        public abstract System.ClientModel.ClientResult UpdateStatus(System.ClientModel.Primitives.RequestOptions? options = null);
+        public abstract System.Threading.Tasks.ValueTask<System.ClientModel.ClientResult> UpdateStatusAsync(System.ClientModel.Primitives.RequestOptions? options = null);
+        public virtual void WaitForCompletion(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { }
+        public virtual System.Threading.Tasks.ValueTask WaitForCompletionAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
     [System.AttributeUsageAttribute(System.AttributeTargets.Class)]
     public sealed partial class PersistableModelProxyAttribute : System.Attribute

sdk/core/System.ClientModel/src/Convenience/OperationResult.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.ClientModel.Internal;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -44,34 +45,108 @@ protected OperationResult(PipelineResponse response)
     /// </summary>
     /// <value>A token that can be used to rehydrate the operation, for example
     /// to monitor its progress or to obtain its final result, from a process
-    /// different thatn the one that started the operation.</value>
+    /// different than the one that started the operation.</value>
     public abstract ContinuationToken? RehydrationToken { get; protected set; }
 
+    /// <summary>
+    /// Sends a request to the service to get the current status of the
+    /// operation and updates <see cref="IsCompleted"/> and other relevant
+    /// properties.
+    /// </summary>
+    /// <param name="options">The <see cref="RequestOptions"/> to be used when
+    /// sending the request to the service.</param>
+    /// <returns>The <see cref="ClientResult"/> returned from the service call.
+    /// </returns>
+    /// <remarks>This method updates the value returned from
+    /// <see cref="ClientResult.GetRawResponse"/> and will update
+    /// <see cref="IsCompleted"/> to <c>true</c> once the operation has finished
+    /// running on the service.  It will also update <c>Value</c> or
+    /// <c>Status</c> properties if present on the <see cref="OperationResult"/>
+    /// derived type.</remarks>
+    public abstract ValueTask<ClientResult> UpdateStatusAsync(RequestOptions? options = default);
+
+    /// <summary>
+    /// Sends a request to the service to get the current status of the
+    /// operation and updates <see cref="IsCompleted"/> and other relevant
+    /// properties.
+    /// </summary>
+    /// <param name="options">The <see cref="RequestOptions"/> to be used when
+    /// sending the request to the service.</param>
+    /// <returns>The <see cref="ClientResult"/> returned from the service call.
+    /// </returns>
+    /// <remarks>This method updates the value returned from
+    /// <see cref="ClientResult.GetRawResponse"/> and will update
+    /// <see cref="IsCompleted"/> to <c>true</c> once the operation has finished
+    /// running on the service.  It will also update <c>Value</c> or
+    /// <c>Status</c> properties if present on the <see cref="OperationResult"/>
+    /// derived type.</remarks>
+    public abstract ClientResult UpdateStatus(RequestOptions? options = default);
+
     /// <summary>
     /// Waits for the operation to complete processing on the service.
     /// </summary>
-    /// <remarks>Derived types may implement <see cref="WaitForCompletionAsync"/>
-    /// using different mechanisms to obtain updates from the service regarding
-    /// the progress of the operation.  If the derived type polls for status
-    /// updates, it provides overloads of <see cref="WaitForCompletionAsync"/>
+    /// <remarks>Derived types may override <see cref="WaitForCompletionAsync"/>
+    /// to implement different mechanisms for obtaining updates from the service
+    /// regarding the progress of the operation. For example, if the derived
+    /// type polls for status updates, it may provides overloads of
+    /// <see cref="WaitForCompletionAsync"/>
     /// that allow the caller to specify the polling interval or delay strategy
-    /// used to wait between sending request for updates.
+    /// used to wait between sending request for updates.  By default,
+    /// <see cref="WaitForCompletionAsync"/> waits a default interval between
+    /// calling <see cref="UpdateStatusAsync"/> to obtain a status updates, so
+    /// if updates are delivered via streaming or another mechanism where a wait
+    /// time is not required, derived types can override this method to update
+    /// the status more frequently.
     /// </remarks>
     /// <exception cref="OperationCanceledException">The <paramref name="cancellationToken"/>
     /// was cancelled.</exception>
-    public abstract Task WaitForCompletionAsync(CancellationToken cancellationToken = default);
+    public virtual async ValueTask WaitForCompletionAsync(CancellationToken cancellationToken = default)
+    {
+        PollingInterval pollingInterval = new();
+
+        while (!IsCompleted)
+        {
+            PipelineResponse response = GetRawResponse();
+
+            await pollingInterval.WaitAsync(response, cancellationToken).ConfigureAwait(false);
+
+            ClientResult result = await UpdateStatusAsync(cancellationToken.ToRequestOptions()).ConfigureAwait(false);
+
+            SetRawResponse(result.GetRawResponse());
+        }
+    }
 
     /// <summary>
     /// Waits for the operation to complete processing on the service.
     /// </summary>
-    /// <remarks>Derived types may implement <see cref="WaitForCompletion"/>
-    /// using different mechanisms to obtain updates from the service regarding
-    /// the progress of the operation.  If the derived type polls for status
-    /// updates, it provides overloads of <see cref="WaitForCompletion"/>
+    /// <remarks>Derived types may override <see cref="WaitForCompletion"/>
+    /// to implement different mechanisms for obtaining updates from the service
+    /// regarding the progress of the operation. For example, if the derived
+    /// type polls for status updates, it may provides overloads of
+    /// <see cref="WaitForCompletion"/>
     /// that allow the caller to specify the polling interval or delay strategy
-    /// used to wait between sending request for updates.
+    /// used to wait between sending request for updates.  By default,
+    /// <see cref="WaitForCompletion"/> waits a default interval between
+    /// calling <see cref="UpdateStatus"/> to obtain a status updates, so
+    /// if updates are delivered via streaming or another mechanism where a wait
+    /// time is not required, derived types can override this method to update
+    /// the status more frequently.
     /// </remarks>
     /// <exception cref="OperationCanceledException">The <paramref name="cancellationToken"/>
     /// was cancelled.</exception>
-    public abstract void WaitForCompletion(CancellationToken cancellationToken = default);
+    public virtual void WaitForCompletion(CancellationToken cancellationToken = default)
+    {
+        PollingInterval pollingInterval = new();
+
+        while (!IsCompleted)
+        {
+            PipelineResponse response = GetRawResponse();
+
+            pollingInterval.Wait(response, cancellationToken);
+
+            ClientResult result = UpdateStatus(cancellationToken.ToRequestOptions());
+
+            SetRawResponse(result.GetRawResponse());
+        }
+    }
 }

sdk/core/System.ClientModel/src/Internal/CancellationTokenExtensions.cs

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.ClientModel.Primitives;
+using System.Threading;
+
+namespace System.ClientModel.Internal;
+
+internal static class CancellationTokenExtensions
+{
+    public static RequestOptions? ToRequestOptions(this CancellationToken cancellationToken)
+    {
+        if (cancellationToken == default)
+        {
+            return null;
+        }
+
+        return new RequestOptions()
+        {
+            CancellationToken = cancellationToken
+        };
+    }
+}

sdk/core/System.ClientModel/src/Internal/PollingInterval.cs

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.ClientModel.Primitives;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.ClientModel.Internal;
+
+internal sealed class PollingInterval
+{
+    private static readonly TimeSpan DefaultDelay = TimeSpan.FromSeconds(1.0);
+
+    private readonly TimeSpan _interval;
+
+    public PollingInterval(TimeSpan? interval = default)
+    {
+        _interval = interval ?? DefaultDelay;
+    }
+
+    public async Task WaitAsync(PipelineResponse response, CancellationToken cancellationToken)
+    {
+        TimeSpan delay = GetDelay(response);
+
+        await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
+    }
+
+    public void Wait(PipelineResponse response, CancellationToken cancellationToken)
+    {
+        TimeSpan delay = GetDelay(response);
+
+        if (!cancellationToken.CanBeCanceled)
+        {
+            Thread.Sleep(delay);
+            return;
+        }
+
+        if (cancellationToken.WaitHandle.WaitOne(delay))
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+        }
+    }
+
+    private TimeSpan GetDelay(PipelineResponse response)
+        => PipelineResponseHeaders.TryGetRetryAfter(response, out TimeSpan retryAfter)
+            && retryAfter > _interval ? retryAfter : _interval;
+}

sdk/core/System.ClientModel/src/Message/PipelineResponseHeaders.cs

@@ -11,6 +11,8 @@ namespace System.ClientModel.Primitives;
 /// </summary>
 public abstract class PipelineResponseHeaders : IEnumerable<KeyValuePair<string, string>>
 {
+    private const string RetryAfterHeaderName = "Retry-After";
+
     /// <summary>
     /// Attempts to retrieve the value associated with the specified header
     /// name.
@@ -37,4 +39,26 @@ public abstract class PipelineResponseHeaders : IEnumerable<KeyValuePair<string,
     public abstract IEnumerator<KeyValuePair<string, string>> GetEnumerator();
 
     IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+    internal static bool TryGetRetryAfter(PipelineResponse response, out TimeSpan value)
+    {
+        // See: https://www.rfc-editor.org/rfc/rfc7231#section-7.1.3
+        if (response.Headers.TryGetValue(RetryAfterHeaderName, out string? retryAfter))
+        {
+            if (int.TryParse(retryAfter, out var delaySeconds))
+            {
+                value = TimeSpan.FromSeconds(delaySeconds);
+                return true;
+            }
+
+            if (DateTimeOffset.TryParse(retryAfter, out DateTimeOffset retryAfterDate))
+            {
+                value = retryAfterDate - DateTimeOffset.Now;
+                return true;
+            }
+        }
+
+        value = default;
+        return false;
+    }
 }

sdk/core/System.ClientModel/src/Pipeline/ClientRetryPolicy.cs

@@ -24,7 +24,6 @@ public class ClientRetryPolicy : PipelinePolicy
 
     private const int DefaultMaxRetries = 3;
     private static readonly TimeSpan DefaultInitialDelay = TimeSpan.FromSeconds(0.8);
-    private const string RetryAfterHeaderName = "Retry-After";
 
     private readonly int _maxRetries;
     private readonly TimeSpan _initialDelay;
@@ -276,7 +275,7 @@ protected virtual TimeSpan GetNextDelay(PipelineMessage message, int tryCount)
         double nextDelayMilliseconds = (1 << (tryCount - 1)) * _initialDelay.TotalMilliseconds;
 
         if (message.Response is not null &&
-            TryGetRetryAfter(message.Response, out TimeSpan retryAfter) &&
+            PipelineResponseHeaders.TryGetRetryAfter(message.Response, out TimeSpan retryAfter) &&
             retryAfter.TotalMilliseconds > nextDelayMilliseconds)
         {
             return retryAfter;
@@ -319,26 +318,4 @@ protected virtual void Wait(TimeSpan time, CancellationToken cancellationToken)
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
         }
     }
-
-    private static bool TryGetRetryAfter(PipelineResponse response, out TimeSpan value)
-    {
-        // See: https://www.rfc-editor.org/rfc/rfc7231#section-7.1.3
-        if (response.Headers.TryGetValue(RetryAfterHeaderName, out string? retryAfter))
-        {
-            if (int.TryParse(retryAfter, out var delaySeconds))
-            {
-                value = TimeSpan.FromSeconds(delaySeconds);
-                return true;
-            }
-
-            if (DateTimeOffset.TryParse(retryAfter, out DateTimeOffset retryAfterDate))
-            {
-                value = retryAfterDate - DateTimeOffset.Now;
-                return true;
-            }
-        }
-
-        value = default;
-        return false;
-    }
 }