Add support for response classification to RequestOptions (Azure#22923)

pakrym · web-flow · commit f4f4d32e0a9c · 2021-08-06T13:25:13.000-07:00
diff --git a/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.netstandard2.0.cs b/sdk/core/Azure.Core.Experimental/api/Azure.Core.Experimental.netstandard2.0.cs
@@ -5,11 +5,23 @@ public partial class RequestOptions
         public RequestOptions() { }
         public RequestOptions(Azure.ResponseStatusOption statusOption) { }
         public RequestOptions(System.Action<Azure.Core.HttpMessage> perCall) { }
+        public RequestOptions(params int[] treatAsSuccess) { }
+        public RequestOptions(int[] statusCodes, Azure.ResponseClassification classification) { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.Core.Pipeline.HttpPipelinePolicy? PerCallPolicy { get { throw null; } set { } }
         public Azure.ResponseStatusOption StatusOption { get { throw null; } set { } }
+        public void AddClassifier(System.Func<Azure.Core.HttpMessage, Azure.ResponseClassification?> classifier) { }
+        public void AddClassifier(int[] statusCodes, Azure.ResponseClassification classification) { }
+        public static void Apply(Azure.RequestOptions requestOptions, Azure.Core.HttpMessage message) { }
         public static implicit operator Azure.RequestOptions (Azure.ResponseStatusOption option) { throw null; }
     }
+    public enum ResponseClassification
+    {
+        Retry = 0,
+        DontRetry = 1,
+        Throw = 2,
+        Success = 3,
+    }
     public enum ResponseStatusOption
     {
         Default = 0,
diff --git a/sdk/core/Azure.Core.Experimental/src/RequestOptions.cs b/sdk/core/Azure.Core.Experimental/src/RequestOptions.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using Azure.Core;
 using Azure.Core.Pipeline;
@@ -13,10 +14,14 @@ namespace Azure
     /// </summary>
     public class RequestOptions
     {
+        private List<HttpMessageClassifier>? _classifiers;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="RequestOptions"/> class.
         /// </summary>
-        public RequestOptions() { }
+        public RequestOptions()
+        {
+        }
 
         /// <summary>
         /// Initializes a new instance of the <see cref="RequestOptions"/> class using the given <see cref="RequestOptions"/>.
@@ -30,6 +35,48 @@ public RequestOptions() { }
         /// <param name="perCall"></param>
         public RequestOptions(Action<HttpMessage> perCall) => PerCallPolicy = new ActionPolicy(perCall);
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="RequestOptions"/> class.
+        /// </summary>
+        /// <param name="treatAsSuccess">The status codes to treat as successful.</param>
+        public RequestOptions(params int[] treatAsSuccess) : this(treatAsSuccess, ResponseClassification.Success)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="RequestOptions"/> class.
+        /// Applying provided classification to a set of status codes.
+        /// </summary>
+        /// <param name="statusCodes">The status codes to classify.</param>
+        /// <param name="classification">The classification.</param>
+        public RequestOptions(int[] statusCodes, ResponseClassification classification)
+        {
+            AddClassifier(statusCodes, classification);
+        }
+
+        /// <summary>
+        /// Adds the classification for provided status codes.
+        /// </summary>
+        /// <param name="statusCodes">The status codes to classify.</param>
+        /// <param name="classification">The classification.</param>
+        public void AddClassifier(int[] statusCodes, ResponseClassification classification)
+        {
+            foreach (var statusCode in statusCodes)
+            {
+                AddClassifier(message => message.Response.Status == statusCode ? classification : null);
+            }
+        }
+
+        /// <summary>
+        /// Adds a function that allows to specify how response would be processed by the pipeline.
+        /// </summary>
+        /// <param name="classifier"></param>
+        public void AddClassifier(Func<HttpMessage, ResponseClassification?> classifier)
+        {
+            _classifiers ??= new();
+            _classifiers.Add(new FuncHttpMessageClassifier(classifier));
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="RequestOptions"/> class using the given <see cref="ResponseStatusOption"/>.
         /// </summary>
@@ -52,6 +99,24 @@ public RequestOptions() { }
         /// </summary>
         public HttpPipelinePolicy? PerCallPolicy { get; set; }
 
+        /// <summary>
+        /// Applies options from <see cref="RequestOptions"/> instance to a <see cref="HttpMessage"/>.
+        /// </summary>
+        /// <param name="requestOptions"></param>
+        /// <param name="message"></param>
+        public static void Apply(RequestOptions requestOptions, HttpMessage message)
+        {
+            if (requestOptions.PerCallPolicy != null)
+            {
+                message.SetProperty("RequestOptionsPerCallPolicyCallback", requestOptions.PerCallPolicy);
+            }
+
+            if (requestOptions._classifiers != null)
+            {
+                message.ResponseClassifier = new PerCallResponseClassifier(message.ResponseClassifier, requestOptions._classifiers);
+            }
+        }
+
         /// <summary>
         /// An <see cref="HttpPipelineSynchronousPolicy"/> which invokes an action when a request is being sent.
         /// </summary>
@@ -63,5 +128,112 @@ internal class ActionPolicy : HttpPipelineSynchronousPolicy
 
             public override void OnSendingRequest(HttpMessage message) => Action.Invoke(message);
         }
+
+        private class PerCallResponseClassifier : ResponseClassifier
+        {
+            private readonly ResponseClassifier _inner;
+            private readonly List<HttpMessageClassifier> _classifiers;
+
+            public PerCallResponseClassifier(ResponseClassifier inner, List<HttpMessageClassifier> classifiers)
+            {
+                _inner = inner;
+                _classifiers = classifiers;
+            }
+
+            public override bool IsRetriableResponse(HttpMessage message)
+            {
+                if (Applies(message, ResponseClassification.DontRetry)) return false;
+                if (Applies(message, ResponseClassification.Retry)) return true;
+
+                return _inner.IsRetriableResponse(message);
+            }
+
+            public override bool IsRetriableException(Exception exception)
+            {
+                return _inner.IsRetriableException(exception);
+            }
+
+            public override bool IsRetriable(HttpMessage message, Exception exception)
+            {
+                if (Applies(message, ResponseClassification.DontRetry)) return false;
+
+                return _inner.IsRetriable(message, exception);
+            }
+
+            public override bool IsErrorResponse(HttpMessage message)
+            {
+                if (Applies(message, ResponseClassification.Throw)) return true;
+                if (Applies(message, ResponseClassification.Success)) return false;
+
+                return _inner.IsErrorResponse(message);
+            }
+
+            private bool Applies(HttpMessage message, ResponseClassification responseClassification)
+            {
+                foreach (var classifier in _classifiers)
+                {
+                    if (classifier.TryClassify(message, null, out var c) &&
+                        c == responseClassification)
+                    {
+                        return true;
+                    }
+                }
+
+                return false;
+            }
+        }
+
+        private abstract class HttpMessageClassifier
+        {
+            public abstract bool TryClassify(HttpMessage message, Exception? exception, out ResponseClassification classification);
+        }
+
+        private class FuncHttpMessageClassifier : HttpMessageClassifier
+        {
+            private readonly Func<HttpMessage, ResponseClassification?> _func;
+
+            public FuncHttpMessageClassifier(Func<HttpMessage, ResponseClassification?> func)
+            {
+                _func = func;
+            }
+
+            public override bool TryClassify(HttpMessage message, Exception? exception, out ResponseClassification classification)
+            {
+                if (_func(message) is ResponseClassification c)
+                {
+                    classification = c;
+                    return true;
+                }
+
+                classification = default;
+                return false;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Specifies how response would be processed by the pipeline and the client.
+    /// </summary>
+    public enum ResponseClassification
+    {
+        /// <summary>
+        /// The response would be retried.
+        /// </summary>
+        Retry,
+
+        /// <summary>
+        /// The response would be retried.
+        /// </summary>
+        DontRetry,
+
+        /// <summary>
+        /// The client would throw an exception for the response.
+        /// </summary>
+        Throw,
+
+        /// <summary>
+        /// The client would tread the response a successful.
+        /// </summary>
+        Success,
     }
-}
+}
diff --git a/sdk/core/Azure.Core.Experimental/tests/RequestOptionsTest.cs b/sdk/core/Azure.Core.Experimental/tests/RequestOptionsTest.cs
@@ -0,0 +1,88 @@
+﻿// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Core.TestFramework;
+using NUnit.Framework;
+
+namespace Azure.Core.Tests
+{
+    public class RequestOptionsTest
+    {
+        [Test]
+        public void CanOverrideDefaultClassificationSuccess()
+        {
+            var m = CreateTestMessage(404);
+
+            var options = new RequestOptions();
+            options.AddClassifier(new[] { 404 }, ResponseClassification.Success);
+
+            RequestOptions.Apply(options, m);
+
+            Assert.False(m.ResponseClassifier.IsErrorResponse(m));
+        }
+
+        [Test]
+        public void CanOverrideDefaultClassificationThrow()
+        {
+            var m = CreateTestMessage(200);
+
+            var options = new RequestOptions();
+            options.AddClassifier(new[] { 200 }, ResponseClassification.Throw);
+
+            RequestOptions.Apply(options, m);
+
+            Assert.True(m.ResponseClassifier.IsErrorResponse(m));
+        }
+
+        [Test]
+        public void CanOverrideDefaultClassificationRetry()
+        {
+            var m = CreateTestMessage(200);
+
+            var options = new RequestOptions();
+            options.AddClassifier(new[] { 200 }, ResponseClassification.Retry);
+
+            RequestOptions.Apply(options, m);
+
+            Assert.True(m.ResponseClassifier.IsRetriableResponse(m));
+        }
+
+        [Test]
+        public void CanOverrideDefaultClassificationNoRetry()
+        {
+            var m = CreateTestMessage(500);
+
+            var options = new RequestOptions();
+            options.AddClassifier(new[] { 500 }, ResponseClassification.DontRetry);
+
+            RequestOptions.Apply(options, m);
+
+            Assert.False(m.ResponseClassifier.IsRetriableResponse(m));
+        }
+
+        [Test]
+        public void CanOverrideDefaultClassificationWithFunc()
+        {
+            HttpMessage m = new HttpMessage(new MockRequest(), new ResponseClassifier())
+            {
+                Response = new MockResponse(500)
+            };
+
+            var options = new RequestOptions();
+            options.AddClassifier(_ => ResponseClassification.Success);
+
+            RequestOptions.Apply(options, m);
+
+            Assert.False(m.ResponseClassifier.IsErrorResponse(m));
+        }
+
+        private static HttpMessage CreateTestMessage(int status)
+        {
+            HttpMessage m = new HttpMessage(new MockRequest(), new ResponseClassifier())
+            {
+                Response = new MockResponse(status)
+            };
+            return m;
+        }
+    }
+}
diff --git a/sdk/core/Azure.Core/api/Azure.Core.net461.cs b/sdk/core/Azure.Core/api/Azure.Core.net461.cs
@@ -312,7 +312,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public Azure.Core.Request Request { get { throw null; } }
         public Azure.Response Response { get { throw null; } set { } }
-        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } }
+        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
         public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs b/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs
@@ -312,7 +312,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public Azure.Core.Request Request { get { throw null; } }
         public Azure.Response Response { get { throw null; } set { } }
-        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } }
+        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
         public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs b/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs
@@ -312,7 +312,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public Azure.Core.Request Request { get { throw null; } }
         public Azure.Response Response { get { throw null; } set { } }
-        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } }
+        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
         public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs b/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs
@@ -312,7 +312,7 @@ public HttpMessage(Azure.Core.Request request, Azure.Core.ResponseClassifier res
         public System.TimeSpan? NetworkTimeout { get { throw null; } set { } }
         public Azure.Core.Request Request { get { throw null; } }
         public Azure.Response Response { get { throw null; } set { } }
-        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } }
+        public Azure.Core.ResponseClassifier ResponseClassifier { get { throw null; } set { } }
         public void Dispose() { }
         public System.IO.Stream? ExtractResponseContent() { throw null; }
         public void SetProperty(string name, object value) { }
diff --git a/sdk/core/Azure.Core/src/HttpMessage.cs b/sdk/core/Azure.Core/src/HttpMessage.cs
@@ -68,7 +68,7 @@ public Response Response
         /// <summary>
         /// The <see cref="ResponseClassifier"/> instance to use for response classification during pipeline invocation.
         /// </summary>
-        public ResponseClassifier ResponseClassifier { get; }
+        public ResponseClassifier ResponseClassifier { get; set; }
 
         /// <summary>
         /// Gets or sets the value indicating if response would be buffered as part of the pipeline. Defaults to true.
