[LLC] Implement RequestOptions.AddPolicy (Azure#25214)

annelo-msft · web-flow · commit e3cce3398f7e · 2021-11-18T00:13:00.000Z
* add test for AddPolicy

* share proposed refactoring to HttpPipelineExtensions

* ideas in the direction of implementation

* mark retry policy

* revert mark retry policy

* shifting API around

* experiment with using MemoryPool

* use ArrayPool

* use ArrayPool

* add tests

* update API

* throw if didn't use pipeline builder

* pr feedback and merge follow-up

* pr fb and merge follow-up

* pr fb

* simplify request pipeline creation

* move AddPolicy tests to Core; pr fb

* nit

* pr fb

* pr fb &amp; fix build

* pr fb

* pr fb &amp; build fix

* pr fb

* pr fb

* pr fb &amp; fix core ci
diff --git a/sdk/core/Azure.Core/api/Azure.Core.net461.cs b/sdk/core/Azure.Core/api/Azure.Core.net461.cs
@@ -194,6 +194,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -764,6 +765,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs b/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs
@@ -194,6 +194,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -764,6 +765,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs b/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs
@@ -194,6 +194,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -764,6 +765,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }
diff --git a/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs b/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs
@@ -194,6 +194,7 @@ public partial class RequestContext
         public RequestContext() { }
         public System.Threading.CancellationToken CancellationToken { get { throw null; } set { } }
         public Azure.ErrorOptions ErrorOptions { get { throw null; } set { } }
+        public void AddPolicy(Azure.Core.Pipeline.HttpPipelinePolicy policy, Azure.Core.HttpPipelinePosition position) { }
         public static implicit operator Azure.RequestContext (Azure.ErrorOptions options) { throw null; }
     }
     public partial class RequestFailedException : System.Exception, System.Runtime.Serialization.ISerializable
@@ -764,6 +765,7 @@ public HttpPipeline(Azure.Core.Pipeline.HttpPipelineTransport transport, Azure.C
         public static System.IDisposable CreateClientRequestIdScope(string? clientRequestId) { throw null; }
         public static System.IDisposable CreateHttpMessagePropertiesScope(System.Collections.Generic.IDictionary<string, object?> messageProperties) { throw null; }
         public Azure.Core.HttpMessage CreateMessage() { throw null; }
+        public Azure.Core.HttpMessage CreateMessage(Azure.RequestContext context) { throw null; }
         public Azure.Core.Request CreateRequest() { throw null; }
         public void Send(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { }
         public System.Threading.Tasks.ValueTask SendAsync(Azure.Core.HttpMessage message, System.Threading.CancellationToken cancellationToken) { throw null; }
diff --git a/sdk/core/Azure.Core/src/HttpMessage.cs b/sdk/core/Azure.Core/src/HttpMessage.cs
@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
-using System.Net.Http;
 using System.Threading;
 using Azure.Core.Pipeline;
 
@@ -13,7 +12,7 @@ namespace Azure.Core
     /// <summary>
     /// Represents a context flowing through the <see cref="HttpPipeline"/>.
     /// </summary>
-    public sealed class HttpMessage: IDisposable
+    public sealed class HttpMessage : IDisposable
     {
         private Dictionary<string, object>? _properties;
 
@@ -81,6 +80,19 @@ public Response Response
         /// </summary>
         public TimeSpan? NetworkTimeout { get; set; }
 
+        internal void AddPolicies(RequestContext context)
+        {
+            if (context == null || context.Policies == null || context.Policies.Count == 0)
+            {
+                return;
+            }
+
+            Policies ??= new(context.Policies.Count);
+            Policies.AddRange(context.Policies);
+        }
+
+        internal List<(HttpPipelinePosition Position, HttpPipelinePolicy Policy)>? Policies { get; set; }
+
         /// <summary>
         /// Gets a property that modifies the pipeline behavior. Please refer to individual policies documentation on what properties it supports.
         /// </summary>
@@ -115,7 +127,7 @@ public void SetProperty(string name, object value)
             {
                 case ResponseShouldNotBeUsedStream responseContent:
                     return responseContent.Original;
-                case Stream stream :
+                case Stream stream:
                     _response.ContentStream = new ResponseShouldNotBeUsedStream(_response.ContentStream);
                     return stream;
                 default:
@@ -132,7 +144,7 @@ public void Dispose()
             _response?.Dispose();
         }
 
-        private class ResponseShouldNotBeUsedStream: Stream
+        private class ResponseShouldNotBeUsedStream : Stream
         {
             public Stream Original { get; }
 
diff --git a/sdk/core/Azure.Core/src/Pipeline/DisposableHttpPipeline.cs b/sdk/core/Azure.Core/src/Pipeline/DisposableHttpPipeline.cs
@@ -14,11 +14,14 @@ public sealed class DisposableHttpPipeline : HttpPipeline, IDisposable
         /// Creates a new instance of <see cref="HttpPipeline"/> with the provided transport, policies and response classifier.
         /// </summary>
         /// <param name="transport">The <see cref="HttpPipelineTransport"/> to use for sending the requests.</param>
+        /// <param name="perCallIndex"></param>
+        /// <param name="perRetryIndex"></param>
         /// <param name="policies">Policies to be invoked as part of the pipeline in order.</param>
         /// <param name="responseClassifier">The response classifier to be used in invocations.</param>
-        internal DisposableHttpPipeline(HttpPipelineTransport transport, HttpPipelinePolicy[]? policies = null, ResponseClassifier? responseClassifier = null)
-            : base(transport, policies, responseClassifier)
-        { }
+        internal DisposableHttpPipeline(HttpPipelineTransport transport, int perCallIndex, int perRetryIndex, HttpPipelinePolicy[]? policies = null, ResponseClassifier? responseClassifier = null)
+            : base(transport, perCallIndex, perRetryIndex, policies, responseClassifier)
+        {
+        }
 
         /// <summary>
         /// Calls Dispose on the underlying <see cref="HttpPipelineTransport"/>.
diff --git a/sdk/core/Azure.Core/src/Pipeline/HttpPipeline.cs b/sdk/core/Azure.Core/src/Pipeline/HttpPipeline.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -19,6 +20,25 @@ public class HttpPipeline
 
         private readonly ReadOnlyMemory<HttpPipelinePolicy> _pipeline;
 
+        /// <summary>
+        /// Indicates whether or not the pipeline was created using its internal constructor.
+        /// If it was, we know the indices where we can add per-request policies at positions
+        /// <see cref="HttpPipelinePosition.PerCall"/> and <see cref="HttpPipelinePosition.PerRetry"/>.
+        /// </summary>
+        private readonly bool _internallyConstructed;
+
+        /// <summary>
+        /// The pipeline index where <see cref="HttpPipelinePosition.PerCall"/> policies will be added,
+        /// if any are specified using <see cref="RequestContext.AddPolicy(HttpPipelinePolicy, HttpPipelinePosition)"/>.
+        /// </summary>
+        private readonly int _perCallIndex;
+
+        /// <summary>
+        /// The pipeline index where <see cref="HttpPipelinePosition.PerRetry"/> policies will be added,
+        /// if any are specified using <see cref="RequestContext.AddPolicy(HttpPipelinePolicy, HttpPipelinePosition)"/>.
+        /// </summary>
+        private readonly int _perRetryIndex;
+
         /// <summary>
         /// Creates a new instance of <see cref="HttpPipeline"/> with the provided transport, policies and response classifier.
         /// </summary>
@@ -39,6 +59,13 @@ public HttpPipeline(HttpPipelineTransport transport, HttpPipelinePolicy[]? polic
             _pipeline = all;
         }
 
+        internal HttpPipeline(HttpPipelineTransport transport, int perCallIndex, int perRetryIndex, HttpPipelinePolicy[]? policies = null, ResponseClassifier? responseClassifier = null) : this(transport, policies, responseClassifier)
+        {
+            _perCallIndex = perCallIndex;
+            _perRetryIndex = perRetryIndex;
+            _internallyConstructed = true;
+        }
+
         /// <summary>
         /// Creates a new <see cref="Request"/> instance.
         /// </summary>
@@ -55,6 +82,18 @@ public HttpMessage CreateMessage()
             return new HttpMessage(CreateRequest(), ResponseClassifier);
         }
 
+        /// <summary>
+        /// Creates a new <see cref="HttpMessage"/> instance.
+        /// </summary>
+        /// <param name="context">Context specifying the message options.</param>
+        /// <returns>The message.</returns>
+        public HttpMessage CreateMessage(RequestContext context)
+        {
+            var message = CreateMessage();
+            message.AddPolicies(context);
+            return message;
+        }
+
         /// <summary>
         /// The <see cref="ResponseClassifier"/> instance used in this pipeline invocations.
         /// </summary>
@@ -70,7 +109,28 @@ public ValueTask SendAsync(HttpMessage message, CancellationToken cancellationTo
         {
             message.CancellationToken = cancellationToken;
             AddHttpMessageProperties(message);
-            return _pipeline.Span[0].ProcessAsync(message, _pipeline.Slice(1));
+
+            if (message.Policies == null || message.Policies.Count == 0)
+            {
+                return _pipeline.Span[0].ProcessAsync(message, _pipeline.Slice(1));
+            }
+
+            return SendAsync(message);
+        }
+
+        private async ValueTask SendAsync(HttpMessage message)
+        {
+            var length = _pipeline.Length + message.Policies!.Count;
+            var policies = ArrayPool<HttpPipelinePolicy>.Shared.Rent(length);
+            try
+            {
+                var pipeline = CreateRequestPipeline(policies, message.Policies);
+                await pipeline.Span[0].ProcessAsync(message, pipeline.Slice(1)).ConfigureAwait(false);
+            }
+            finally
+            {
+                ArrayPool<HttpPipelinePolicy>.Shared.Return(policies);
+            }
         }
 
         /// <summary>
@@ -82,8 +142,27 @@ public void Send(HttpMessage message, CancellationToken cancellationToken)
         {
             message.CancellationToken = cancellationToken;
             AddHttpMessageProperties(message);
-            _pipeline.Span[0].Process(message, _pipeline.Slice(1));
+
+            if (message.Policies == null || message.Policies.Count == 0)
+            {
+                _pipeline.Span[0].Process(message, _pipeline.Slice(1));
+            }
+            else
+            {
+                var length = _pipeline.Length + message.Policies.Count;
+                var policies = ArrayPool<HttpPipelinePolicy>.Shared.Rent(length);
+                try
+                {
+                    var pipeline = CreateRequestPipeline(policies, message.Policies);
+                    pipeline.Span[0].Process(message, pipeline.Slice(1));
+                }
+                finally
+                {
+                    ArrayPool<HttpPipelinePolicy>.Shared.Return(policies);
+                }
+            }
         }
+
         /// <summary>
         /// Invokes the pipeline asynchronously with the provided request.
         /// </summary>
@@ -144,6 +223,60 @@ public static IDisposable CreateHttpMessagePropertiesScope(IDictionary<string, o
             return CurrentHttpMessagePropertiesScope.Value;
         }
 
+        private ReadOnlyMemory<HttpPipelinePolicy> CreateRequestPipeline(HttpPipelinePolicy[] policies, List<(HttpPipelinePosition Position, HttpPipelinePolicy Policy)> customPolicies)
+        {
+            if (!_internallyConstructed)
+            {
+                throw new InvalidOperationException("Cannot send messages with per-request policies if the pipeline wasn't constructed with HttpPipelineBuilder.");
+            }
+
+            // Copy over client policies and splice in custom policies at designated indices
+            var pipeline = _pipeline.Span;
+            int transportIndex = pipeline.Length - 1;
+
+            pipeline.Slice(0, _perCallIndex).CopyTo(policies);
+
+            int index = _perCallIndex;
+            int count = AddCustomPolicies(customPolicies, policies, HttpPipelinePosition.PerCall, index);
+
+            index += count;
+            count = _perRetryIndex - _perCallIndex;
+            pipeline.Slice(_perCallIndex, count).CopyTo(policies.AsSpan(index, count));
+
+            index += count;
+            count = AddCustomPolicies(customPolicies, policies, HttpPipelinePosition.PerRetry, index);
+
+            index += count;
+            count = transportIndex - _perRetryIndex;
+            pipeline.Slice(_perRetryIndex, count).CopyTo(policies.AsSpan(index, count));
+
+            index += count;
+            count = AddCustomPolicies(customPolicies, policies, HttpPipelinePosition.BeforeTransport, index);
+
+            index += count;
+            policies[index] = pipeline[transportIndex];
+
+            return new ReadOnlyMemory<HttpPipelinePolicy>(policies, 0, index + 1);
+        }
+
+        private static int AddCustomPolicies(List<(HttpPipelinePosition Position, HttpPipelinePolicy Policy)> source, HttpPipelinePolicy[] target, HttpPipelinePosition position, int start)
+        {
+            int count = 0;
+            if (source != null)
+            {
+                foreach (var policy in source)
+                {
+                    if (policy.Position == position)
+                    {
+                        target[start + count] = policy.Policy;
+                        count++;
+                    }
+                }
+            }
+
+            return count;
+        }
+
         private static void AddHttpMessageProperties(HttpMessage message)
         {
             if (CurrentHttpMessagePropertiesScope.Value != null)
diff --git a/sdk/core/Azure.Core/src/Pipeline/HttpPipelineBuilder.cs b/sdk/core/Azure.Core/src/Pipeline/HttpPipelineBuilder.cs
@@ -53,6 +53,8 @@ public static HttpPipeline Build(
         /// <returns>A new instance of <see cref="HttpPipeline"/></returns>
         public static HttpPipeline Build(ClientOptions options, HttpPipelinePolicy[] perCallPolicies, HttpPipelinePolicy[] perRetryPolicies, ResponseClassifier responseClassifier, HttpPipelineTransportOptions? defaultTransportOptions)
         {
+            int perCallIndex;
+            int perRetryIndex;
             if (perCallPolicies == null)
             {
                 throw new ArgumentNullException(nameof(perCallPolicies));
@@ -94,6 +96,9 @@ void AddCustomerPolicies(HttpPipelinePosition position)
 
             AddCustomerPolicies(HttpPipelinePosition.PerCall);
 
+            policies.RemoveAll(static policy => policy == null);
+            perCallIndex = policies.Count;
+
             policies.Add(ClientRequestIdPolicy.Shared);
 
             if (diagnostics.IsTelemetryEnabled)
@@ -110,6 +115,9 @@ void AddCustomerPolicies(HttpPipelinePosition position)
 
             AddCustomerPolicies(HttpPipelinePosition.PerRetry);
 
+            policies.RemoveAll(static policy => policy == null);
+            perRetryIndex = policies.Count;
+
             if (diagnostics.IsLoggingEnabled)
             {
                 string assemblyName = options.GetType().Assembly!.GetName().Name!;
@@ -122,7 +130,6 @@ void AddCustomerPolicies(HttpPipelinePosition position)
             policies.Add(new RequestActivityPolicy(isDistributedTracingEnabled, ClientDiagnostics.GetResourceProviderNamespace(options.GetType().Assembly), sanitizer));
 
             AddCustomerPolicies(HttpPipelinePosition.BeforeTransport);
-
             policies.RemoveAll(static policy => policy == null);
 
             // Override the provided Transport with the provided transport options if the transport has not been set after default construction and options are not null.
@@ -141,12 +148,16 @@ void AddCustomerPolicies(HttpPipelinePosition position)
                 {
                     transport = HttpPipelineTransport.Create(defaultTransportOptions);
                     return new DisposableHttpPipeline(transport,
+                        perCallIndex,
+                        perRetryIndex,
                         policies.ToArray(),
                         responseClassifier);
                 }
             }
 
             return new HttpPipeline(transport,
+                perCallIndex,
+                perRetryIndex,
                 policies.ToArray(),
                 responseClassifier);
         }
diff --git a/sdk/core/Azure.Core/src/Pipeline/HttpPipelineTransport.cs b/sdk/core/Azure.Core/src/Pipeline/HttpPipelineTransport.cs
@@ -25,7 +25,7 @@ public abstract class HttpPipelineTransport
 
         /// <summary>
         /// Creates a new transport specific instance of <see cref="Request"/>. This should not be called directly, <see cref="HttpPipeline.CreateRequest"/> or
-        /// <see cref="HttpPipeline.CreateMessage"/> should be used instead.
+        /// <see cref="HttpPipeline.CreateMessage()"/> should be used instead.
         /// </summary>
         /// <returns></returns>
         public abstract Request CreateRequest();
diff --git a/sdk/core/Azure.Core/src/Request.cs b/sdk/core/Azure.Core/src/Request.cs
diff --git a/sdk/core/Azure.Core/src/RequestContext.cs b/sdk/core/Azure.Core/src/RequestContext.cs
diff --git a/sdk/core/Azure.Core/tests/HttpPipelineBuilderTest.cs b/sdk/core/Azure.Core/tests/HttpPipelineBuilderTest.cs
diff --git a/sdk/core/Azure.Core/tests/HttpPipelineTests.cs b/sdk/core/Azure.Core/tests/HttpPipelineTests.cs
