ClientModel: Add retry classification to MessageClassifier (Azure#41586)

* Add IsRetriable to MessageClassifier * simplify APIs * fix * fix * nits * Make TryClassify methods abstract
azure-sdk · Jan 26, 2024 · ad9610c · ad9610c
1 parent 5b894af
commit ad9610c
Show file tree

Hide file tree

Showing 10 changed files with 406 additions and 109 deletions.
diff --git a/sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs b/sdk/core/System.ClientModel/api/System.ClientModel.net6.0.cs
@@ -147,7 +147,7 @@ public partial class PipelineMessage : System.IDisposable
         protected internal PipelineMessage(System.ClientModel.Primitives.PipelineRequest request) { }
         public bool BufferResponse { get { throw null; } set { } }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } protected internal set { } }
-        public System.ClientModel.Primitives.PipelineMessageClassifier? MessageClassifier { get { throw null; } set { } }
+        public System.ClientModel.Primitives.PipelineMessageClassifier MessageClassifier { get { throw null; } set { } }
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public System.ClientModel.Primitives.PipelineRequest Request { get { throw null; } }
         public System.ClientModel.Primitives.PipelineResponse? Response { get { throw null; } protected internal set { } }
@@ -157,11 +157,13 @@ protected virtual void Dispose(bool disposing) { }
         public void SetProperty(System.Type type, object value) { }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }
-    public partial class PipelineMessageClassifier
+    public abstract partial class PipelineMessageClassifier
     {
-        protected internal PipelineMessageClassifier() { }
+        protected PipelineMessageClassifier() { }
+        public static System.ClientModel.Primitives.PipelineMessageClassifier Default { get { throw null; } }
         public static System.ClientModel.Primitives.PipelineMessageClassifier Create(System.ReadOnlySpan<ushort> successStatusCodes) { throw null; }
-        public virtual bool IsErrorResponse(System.ClientModel.Primitives.PipelineMessage message) { throw null; }
+        public abstract bool TryClassify(System.ClientModel.Primitives.PipelineMessage message, out bool isError);
+        public abstract bool TryClassify(System.ClientModel.Primitives.PipelineMessage message, System.Exception? exception, out bool isRetriable);
     }
     public abstract partial class PipelinePolicy
     {

diff --git a/sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs b/sdk/core/System.ClientModel/api/System.ClientModel.netstandard2.0.cs
@@ -146,7 +146,7 @@ public partial class PipelineMessage : System.IDisposable
         protected internal PipelineMessage(System.ClientModel.Primitives.PipelineRequest request) { }
         public bool BufferResponse { get { throw null; } set { } }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } protected internal set { } }
-        public System.ClientModel.Primitives.PipelineMessageClassifier? MessageClassifier { get { throw null; } set { } }
+        public System.ClientModel.Primitives.PipelineMessageClassifier MessageClassifier { get { throw null; } set { } }
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public System.ClientModel.Primitives.PipelineRequest Request { get { throw null; } }
         public System.ClientModel.Primitives.PipelineResponse? Response { get { throw null; } protected internal set { } }
@@ -156,11 +156,13 @@ protected virtual void Dispose(bool disposing) { }
         public void SetProperty(System.Type type, object value) { }
         public bool TryGetProperty(System.Type type, out object? value) { throw null; }
     }
-    public partial class PipelineMessageClassifier
+    public abstract partial class PipelineMessageClassifier
     {
-        protected internal PipelineMessageClassifier() { }
+        protected PipelineMessageClassifier() { }
+        public static System.ClientModel.Primitives.PipelineMessageClassifier Default { get { throw null; } }
         public static System.ClientModel.Primitives.PipelineMessageClassifier Create(System.ReadOnlySpan<ushort> successStatusCodes) { throw null; }
-        public virtual bool IsErrorResponse(System.ClientModel.Primitives.PipelineMessage message) { throw null; }
+        public abstract bool TryClassify(System.ClientModel.Primitives.PipelineMessage message, out bool isError);
+        public abstract bool TryClassify(System.ClientModel.Primitives.PipelineMessage message, System.Exception? exception, out bool isRetriable);
     }
     public abstract partial class PipelinePolicy
     {

diff --git a/sdk/core/System.ClientModel/src/Message/PipelineMessage.cs b/sdk/core/System.ClientModel/src/Message/PipelineMessage.cs
@@ -19,6 +19,7 @@ protected internal PipelineMessage(PipelineRequest request)
         _propertyBag = new ArrayBackedPropertyBag<ulong, object>();
 
         BufferResponse = true;
+        MessageClassifier = PipelineMessageClassifier.Default;
     }
 
     public PipelineRequest Request { get; }
@@ -64,7 +65,7 @@ public CancellationToken CancellationToken
     // the client-provided classifier or compose a chain of classification
     // handlers that preserve the functionality of the client-provided classifier
     // at the end of the chain.
-    public PipelineMessageClassifier? MessageClassifier { get; set; }
+    public PipelineMessageClassifier MessageClassifier { get; set; }
 
     public void Apply(RequestOptions options)
     {

diff --git a/sdk/core/System.ClientModel/src/Options/PipelineMessageClassifier.cs b/sdk/core/System.ClientModel/src/Options/PipelineMessageClassifier.cs
@@ -1,25 +1,83 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.IO;
+
 namespace System.ClientModel.Primitives;
 
-public class PipelineMessageClassifier
+public abstract class PipelineMessageClassifier
 {
-    internal static PipelineMessageClassifier Default { get; } = new PipelineMessageClassifier();
+    public static PipelineMessageClassifier Default { get; } = new EndOfChainClassifier();
 
     public static PipelineMessageClassifier Create(ReadOnlySpan<ushort> successStatusCodes)
         => new ResponseStatusClassifier(successStatusCodes);
 
-    protected internal PipelineMessageClassifier() { }
+    protected PipelineMessageClassifier() { }
+
+    public abstract bool TryClassify(PipelineMessage message, out bool isError);
 
-    /// <summary>
-    /// Specifies if the response contained in the <paramref name="message"/> is not successful.
-    /// </summary>
-    public virtual bool IsErrorResponse(PipelineMessage message)
+    public abstract bool TryClassify(PipelineMessage message, Exception? exception, out bool isRetriable);
+
+    internal class EndOfChainClassifier : PipelineMessageClassifier
     {
-        message.AssertResponse();
+        public override bool TryClassify(PipelineMessage message, out bool isError)
+        {
+            message.AssertResponse();
+
+            int statusKind = message.Response!.Status / 100;
+            isError = statusKind == 4 || statusKind == 5;
+
+            // Always classify the message
+            return true;
+        }
+
+        public override bool TryClassify(PipelineMessage message, Exception? exception, out bool isRetriable)
+        {
+            isRetriable = exception is null ?
+                IsRetriable(message) :
+                IsRetriable(message, exception);
+
+            // Always classify the message
+            return true;
+        }
+
+        private static bool IsRetriable(PipelineMessage message)
+        {
+            message.AssertResponse();
+
+            return message.Response!.Status switch
+            {
+                // Request Timeout
+                408 => true,
+
+                // Too Many Requests
+                429 => true,
+
+                // Internal Server Error
+                500 => true,
+
+                // Bad Gateway
+                502 => true,
+
+                // Service Unavailable
+                503 => true,
+
+                // Gateway Timeout
+                504 => true,
+
+                // Default case
+                _ => false
+            };
+        }
+
+        private static bool IsRetriable(PipelineMessage message, Exception exception)
+            => IsRetriable(exception) ||
+                // Retry non-user initiated cancellations
+                (exception is OperationCanceledException &&
+                !message.CancellationToken.IsCancellationRequested);
 
-        int statusKind = message.Response!.Status / 100;
-        return statusKind == 4 || statusKind == 5;
+        private static bool IsRetriable(Exception exception)
+            => (exception is IOException) ||
+                (exception is ClientResultException ex && ex.Status == 0);
     }
 }
diff --git a/sdk/core/System.ClientModel/src/Options/RequestOptions.cs b/sdk/core/System.ClientModel/src/Options/RequestOptions.cs
@@ -106,12 +106,6 @@ internal void Apply(PipelineMessage message)
         // cancellation token will be set again in HttpPipeline.Send.
         message.CancellationToken = CancellationToken;
 
-        // We don't overwrite the classifier on the message if it's already set.
-        // This preserves any values set by the client author, and is also
-        // needed for Azure.Core-based clients so we don't overwrite a default
-        // Azure.Core ResponseClassifier.
-        message.MessageClassifier ??= PipelineMessageClassifier.Default;
-
         // Copy custom pipeline policies to the message.
         message.PerCallPolicies = _perCallPolicies;
         message.PerTryPolicies = _perTryPolicies;

diff --git a/sdk/core/System.ClientModel/src/Options/ResponseStatusClassifier.cs b/sdk/core/System.ClientModel/src/Options/ResponseStatusClassifier.cs
@@ -2,40 +2,53 @@
 // Licensed under the MIT License.
 
 using System.ClientModel.Internal;
+using System.Diagnostics;
 
-namespace System.ClientModel.Primitives
+namespace System.ClientModel.Primitives;
+
+internal class ResponseStatusClassifier : PipelineMessageClassifier
 {
-    internal class ResponseStatusClassifier : PipelineMessageClassifier
+    private BitVector640 _successCodes;
+
+    /// <summary>
+    /// Creates a new instance of <see cref="ResponseStatusClassifier"/>.
+    /// </summary>
+    /// <param name="successStatusCodes">The status codes that this classifier
+    /// will consider not to be errors.</param>
+    public ResponseStatusClassifier(ReadOnlySpan<ushort> successStatusCodes)
     {
-        private BitVector640 _successCodes;
-
-        /// <summary>
-        /// Creates a new instance of <see cref="ResponseStatusClassifier"/>.
-        /// </summary>
-        /// <param name="successStatusCodes">The status codes that this classifier
-        /// will consider not to be errors.</param>
-        public ResponseStatusClassifier(ReadOnlySpan<ushort> successStatusCodes)
-        {
-            _successCodes = new();
+        _successCodes = new();
 
-            foreach (int statusCode in successStatusCodes)
-            {
-                AddClassifier(statusCode, isError: false);
-            }
+        foreach (int statusCode in successStatusCodes)
+        {
+            AddClassifier(statusCode, isError: false);
         }
+    }
 
-        public sealed override bool IsErrorResponse(PipelineMessage message)
-        {
-            message.AssertResponse();
+    public override bool TryClassify(PipelineMessage message, out bool isError)
+    {
+        message.AssertResponse();
 
-            return !_successCodes[message.Response!.Status];
-        }
+        isError = !_successCodes[message.Response!.Status];
 
-        private void AddClassifier(int statusCode, bool isError)
-        {
-            Argument.AssertInRange(statusCode, 0, 639, nameof(statusCode));
+        // BitVector-based classifiers should always end any chain.
+        return true;
+    }
 
-            _successCodes[statusCode] = !isError;
-        }
+    public override bool TryClassify(PipelineMessage message, Exception? exception, out bool isRetriable)
+    {
+        bool classified = Default.TryClassify(message, exception, out isRetriable);
+
+        Debug.Assert(classified);
+
+        // BitVector-based classifiers should always end any chain.
+        return true;
+    }
+
+    private void AddClassifier(int statusCode, bool isError)
+    {
+        Argument.AssertInRange(statusCode, 0, 639, nameof(statusCode));
+
+        _successCodes[statusCode] = !isError;
     }
-}
+}
diff --git a/sdk/core/System.ClientModel/src/Pipeline/ClientRetryPolicy.cs b/sdk/core/System.ClientModel/src/Pipeline/ClientRetryPolicy.cs
@@ -3,7 +3,7 @@
 
 using System.ClientModel.Internal;
 using System.Collections.Generic;
-using System.IO;
+using System.Diagnostics;
 using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
@@ -165,9 +165,14 @@ protected virtual bool ShouldRetryCore(PipelineMessage message, Exception? excep
             return false;
         }
 
-        return exception is null ?
-            IsRetriable(message) :
-            IsRetriable(message, exception);
+        if (!message.MessageClassifier.TryClassify(message, exception, out bool isRetriable))
+        {
+            bool classified = PipelineMessageClassifier.Default.TryClassify(message, exception, out isRetriable);
+
+            Debug.Assert(classified);
+        }
+
+        return isRetriable;
     }
 
     protected virtual ValueTask<bool> ShouldRetryCoreAsync(PipelineMessage message, Exception? exception)
@@ -200,48 +205,4 @@ protected virtual void WaitCore(TimeSpan time, CancellationToken cancellationTok
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
         }
     }
-
-    #region Retry Classifier
-
-    // Overriding response-retriable classification will be added in a later ClientModel release.
-    private static bool IsRetriable(PipelineMessage message)
-    {
-        message.AssertResponse();
-
-        return message.Response!.Status switch
-        {
-            // Request Timeout
-            408 => true,
-
-            // Too Many Requests
-            429 => true,
-
-            // Internal Server Error
-            500 => true,
-
-            // Bad Gateway
-            502 => true,
-
-            // Service Unavailable
-            503 => true,
-
-            // Gateway Timeout
-            504 => true,
-
-            // Default case
-            _ => false
-        };
-    }
-
-    private static bool IsRetriable(PipelineMessage message, Exception exception)
-        => IsRetriable(exception) ||
-            // Retry non-user initiated cancellations
-            (exception is OperationCanceledException &&
-            !message.CancellationToken.IsCancellationRequested);
-
-    private static bool IsRetriable(Exception exception)
-        => (exception is IOException) ||
-            (exception is ClientResultException ex && ex.Status == 0);
-
-    #endregion
 }
diff --git a/sdk/core/System.ClientModel/src/Pipeline/PipelineTransport.cs b/sdk/core/System.ClientModel/src/Pipeline/PipelineTransport.cs
@@ -41,9 +41,17 @@ public async ValueTask ProcessAsync(PipelineMessage message)
         message.Response.SetIsError(ClassifyResponse(message));
     }
 
-    private static bool ClassifyResponse(PipelineMessage message) =>
-        message.MessageClassifier?.IsErrorResponse(message) ??
-        PipelineMessageClassifier.Default.IsErrorResponse(message);
+    private static bool ClassifyResponse(PipelineMessage message)
+    {
+        if (!message.MessageClassifier.TryClassify(message, out bool isError))
+        {
+            bool classified = PipelineMessageClassifier.Default.TryClassify(message, out isError);
+
+            Debug.Assert(classified);
+        }
+
+        return isError;
+    }
 
     protected abstract void ProcessCore(PipelineMessage message);