[Perf] Add support for test-proxy (#22621)

Author: mikeharder

common/Perf/Azure.Test.Perf/IPerfTest.cs

@@ -11,8 +11,10 @@ internal interface IPerfTest : IDisposable, IAsyncDisposable
     {
         Task GlobalSetupAsync();
         Task SetupAsync();
+        Task RecordAndStartPlayback();
         void Run(CancellationToken cancellationToken);
         Task RunAsync(CancellationToken cancellationToken);
+        Task StopPlayback();
         Task CleanupAsync();
         Task GlobalCleanupAsync();
     }

common/Perf/Azure.Test.Perf/PerfOptions.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using CommandLine;
+using System;
 
 namespace Azure.Test.Perf
 {
@@ -55,6 +56,9 @@ public class PerfOptions
         [Option("sync", HelpText = "Runs sync version of test")]
         public bool Sync { get; set; }
 
+        [Option('x', "test-proxy", HelpText = "URI of TestProxy Server")]
+        public Uri TestProxy { get; set; }
+
         [Option('w', "warmup", Default = 5, HelpText = "Duration of warmup in seconds")]
         public int Warmup { get; set; }
     }

common/Perf/Azure.Test.Perf/PerfProgram.cs

@@ -105,12 +105,26 @@ private static async Task Run(Type testType, PerfOptions options)
                 {
                     await tests[0].GlobalSetupAsync();
 
+                    var startedPlayback = false;
+
                     try
                     {
                         await Task.WhenAll(tests.Select(t => t.SetupAsync()));
                         setupStatusCts.Cancel();
                         setupStatusThread.Join();
 
+                        if (options.TestProxy != null)
+                        {
+                            using var recordStatusCts = new CancellationTokenSource();
+                            var recordStatusThread = PerfStressUtilities.PrintStatus("=== Record and Start Playback ===", () => ".", newLine: false, recordStatusCts.Token);
+
+                            await Task.WhenAll(tests.Select(t => t.RecordAndStartPlayback()));
+                            startedPlayback = true;
+
+                            recordStatusCts.Cancel();
+                            recordStatusThread.Join();
+                        }
+
                         if (options.Warmup > 0)
                         {
                             await RunTestsAsync(tests, options, "Warmup", warmup: true);
@@ -141,14 +155,28 @@ private static async Task Run(Type testType, PerfOptions options)
                     }
                     finally
                     {
-                        if (!options.NoCleanup)
+                        try
                         {
-                            if (cleanupStatusThread == null)
+                            if (startedPlayback)
                             {
-                                cleanupStatusThread = PerfStressUtilities.PrintStatus("=== Cleanup ===", () => ".", newLine: false, cleanupStatusCts.Token);
+                                using var playbackStatusCts = new CancellationTokenSource();
+                                var playbackStatusThread = PerfStressUtilities.PrintStatus("=== Stop Playback ===", () => ".", newLine: false, playbackStatusCts.Token);
+                                await Task.WhenAll(tests.Select(t => t.StopPlayback()));
+                                playbackStatusCts.Cancel();
+                                playbackStatusThread.Join();
                             }
+                        }
+                        finally
+                        {
+                            if (!options.NoCleanup)
+                            {
+                                if (cleanupStatusThread == null)
+                                {
+                                    cleanupStatusThread = PerfStressUtilities.PrintStatus("=== Cleanup ===", () => ".", newLine: false, cleanupStatusCts.Token);
+                                }
 
-                            await Task.WhenAll(tests.Select(t => t.CleanupAsync()));
+                                await Task.WhenAll(tests.Select(t => t.CleanupAsync()));
+                            }
                         }
                     }
                 }

common/Perf/Azure.Test.Perf/PerfTest.cs

@@ -1,19 +1,87 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using Azure.Core;
+using Azure.Core.Pipeline;
 using System;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Azure.Test.Perf
 {
     public abstract class PerfTest<TOptions> : IPerfTest where TOptions : PerfOptions
     {
+        private readonly HttpPipelineTransport _insecureTransport;
+
+        private readonly HttpClient _recordPlaybackHttpClient;
+        private readonly TestProxyPolicy _testProxyPolicy;
+
+        private string _recordingId;
+
         protected TOptions Options { get; private set; }
 
         public PerfTest(TOptions options)
         {
             Options = options;
+
+            if (Options.Insecure)
+            {
+                var transport = (new PerfClientOptions()).Transport;
+
+                // Disable SSL validation
+                if (transport is HttpClientTransport)
+                {
+                    _insecureTransport = new HttpClientTransport(new HttpClient(new HttpClientHandler()
+                    {
+                        ServerCertificateCustomValidationCallback = (message, cert, chain, errors) => true
+                    }));
+                }
+                else
+                {
+                    // Assume _transport is HttpWebRequestTransport (currently internal class)
+                    ServicePointManager.ServerCertificateValidationCallback = (message, cert, chain, errors) => true;
+
+                    _insecureTransport = transport;
+                }
+            }
+
+            if (Options.TestProxy != null)
+            {
+                if (Options.Insecure)
+                {
+                    _recordPlaybackHttpClient = new HttpClient(new HttpClientHandler()
+                    {
+                        ServerCertificateCustomValidationCallback = (message, cert, chain, errors) => true
+                    });
+                }
+                else
+                {
+                    _recordPlaybackHttpClient = new HttpClient();
+                }
+
+                _testProxyPolicy = new TestProxyPolicy(Options.TestProxy);
+            }
+        }
+
+        protected TClientOptions ConfigureClientOptions<TClientOptions>(TClientOptions clientOptions) where TClientOptions : ClientOptions
+        {
+            if (_insecureTransport != null)
+            {
+                clientOptions.Transport = _insecureTransport;
+            }
+
+            if (_testProxyPolicy != null)
+            {
+                // TestProxyPolicy should be per-retry to run as late as possible in the pipeline.  For example, some
+                // clients compute a request signature as a per-retry policy, and TestProxyPolicy should run after the
+                // signature is computed to avoid altering the signature.
+                clientOptions.AddPolicy(_testProxyPolicy, HttpPipelinePosition.PerRetry);
+            }
+
+            return clientOptions;
         }
 
         public virtual Task GlobalSetupAsync()
@@ -26,10 +94,48 @@ public virtual Task SetupAsync()
             return Task.CompletedTask;
         }
 
+        public async Task RecordAndStartPlayback()
+        {
+            await StartRecording();
+
+            _testProxyPolicy.RecordingId = _recordingId;
+            _testProxyPolicy.Mode = "record";
+
+            // Record one call to Run()
+            if (Options.Sync)
+            {
+                Run(CancellationToken.None);
+            }
+            else
+            {
+                await RunAsync(CancellationToken.None);
+            }
+
+            await StopRecording();
+
+            await StartPlayback();
+
+            _testProxyPolicy.Mode = "playback";
+            _testProxyPolicy.RecordingId = _recordingId;
+        }
+
         public abstract void Run(CancellationToken cancellationToken);
 
         public abstract Task RunAsync(CancellationToken cancellationToken);
 
+        public async Task StopPlayback()
+        {
+            var message = new HttpRequestMessage(HttpMethod.Post, new Uri(Options.TestProxy, "/playback/stop"));
+            message.Headers.Add("x-recording-id", _recordingId);
+            message.Headers.Add("x-purge-inmemory-recording", bool.TrueString);
+
+            await _recordPlaybackHttpClient.SendAsync(message);
+
+            // Stop redirecting requests to test proxy
+            _testProxyPolicy.Mode = null;
+            _testProxyPolicy.RecordingId = null;
+        }
+
         public virtual Task CleanupAsync()
         {
             return Task.CompletedTask;
@@ -73,5 +179,33 @@ protected static string GetEnvironmentVariable(string name)
             }
             return value;
         }
+
+        private async Task StartRecording()
+        {
+            var message = new HttpRequestMessage(HttpMethod.Post, new Uri(Options.TestProxy, "/record/start"));
+
+            var response = await _recordPlaybackHttpClient.SendAsync(message);
+            _recordingId = response.Headers.GetValues("x-recording-id").Single();
+        }
+
+        private async Task StopRecording()
+        {
+            var message = new HttpRequestMessage(HttpMethod.Post, new Uri(Options.TestProxy, "/record/stop"));
+            message.Headers.Add("x-recording-id", _recordingId);
+
+            await _recordPlaybackHttpClient.SendAsync(message);
+        }
+
+        private async Task StartPlayback()
+        {
+            var message = new HttpRequestMessage(HttpMethod.Post, new Uri(Options.TestProxy, "/playback/start"));
+            message.Headers.Add("x-recording-id", _recordingId);
+
+            var response = await _recordPlaybackHttpClient.SendAsync(message);
+            _recordingId = response.Headers.GetValues("x-recording-id").Single();
+        }
+
+        // Dummy class used to fetch default HttpPipelineTransport
+        private class PerfClientOptions : ClientOptions { }
     }
 }