Adding Test-Proxy support to [Perf Fw] tests (#3257)

* cp

* cp

* cp

* Enable test proxy for perf tests

* unwanted

* format

* fix docs

* cspell

* parallel support and multi proxy

* update arg name

* more logs

* fix request redirection

* Apply suggestions from code review

Co-authored-by: Ahson Khan <ahkha@microsoft.com>

* run once before record, then after recording on

Co-authored-by: Ahson Khan <ahkha@microsoft.com>

sdk/core/perf/CMakeLists.txt

@@ -28,6 +28,7 @@ set(
 set(
   AZURE_PERFORMANCE_SOURCE
   src/arg_parser.cpp
+  src/base_test.cpp
   src/options.cpp
   src/program.cpp
   src/random_stream.cpp

sdk/core/perf/inc/azure/perf/base_test.hpp

@@ -9,21 +9,69 @@
 
 #pragma once
 
-#include <azure/core/context.hpp>
-
+#include "azure/perf/options.hpp"
 #include "azure/perf/test_options.hpp"
+#include <azure/core/context.hpp>
+#include <azure/core/internal/client_options.hpp>
+#include <azure/core/url.hpp>
 
 #include <string>
 #include <vector>
 
+namespace {
+class ProxyPolicy;
+}
+
 namespace Azure { namespace Perf {
+  class Program;
 
   /**
    * @brief The base interface for a performance test.
    *
    */
-  struct BaseTest
-  {
+  class BaseTest {
+    // Provides private access so a test program can run PostSetup.
+    friend class Program;
+    friend class ::ProxyPolicy;
+
+  private:
+    std::string m_recordId;
+    std::string m_proxy;
+    bool m_isPlayBackMode = false;
+
+    void SetTestProxy(std::string const& proxy) { m_proxy = proxy; }
+
+    /**
+     * @brief Define actions to run after test set up and before the actual test.
+     *
+     * @details This method enables the performance framework to set the proxy server for recordings
+     * or any other configuration to happen after a test set up definition.
+     *
+     */
+    void PostSetUp();
+
+    /**
+     * @brief Define actions to run after each test run.
+     *
+     * @details This method enabled the performance framework to remove test proxy forwarding before
+     * letting test do clean up.
+     *
+     */
+    void PreCleanUp();
+
+    /**
+     * @brief Set the client options depending on the test options.
+     *
+     * @param clientOptions ref to the client options that contains the http pipeline policies.
+     */
+    void ConfigureCoreClientOptions(Azure::Core::_internal::ClientOptions* clientOptions);
+
+  protected:
+    Azure::Perf::TestOptions m_options;
+
+  public:
+    BaseTest(Azure::Perf::TestOptions options) : m_options(options) {}
+
     /**
      * @brief Run one time at the beggining and before any test.
      *
@@ -77,5 +125,30 @@ namespace Azure { namespace Perf {
      *
      */
     virtual void GlobalCleanup() {}
+
+    /**
+     * @brief Update an existing \p clientOptions with the test configuration set by the
+     * environment.
+     *
+     * @note If test proxy env var is set, the proxy policy is added to the \p clientOptions.
+     *
+     * @param clientOptions Ref to the client options that contains the Http client policies.
+     */
+    template <class T> void ConfigureClientOptions(T* clientOptions)
+    {
+      ConfigureCoreClientOptions(clientOptions);
+    }
+
+    /**
+     * @brief Create and return client options with test configuration set in the environment.
+     *
+     * @note If test proxy env var is set, the proxy policy is added to the \p clientOptions.
+     */
+    template <class T> T InitClientOptions()
+    {
+      T options;
+      ConfigureClientOptions(&options);
+      return options;
+    }
   };
 }} // namespace Azure::Perf

sdk/core/perf/inc/azure/perf/options.hpp

@@ -96,6 +96,19 @@ namespace Azure { namespace Perf {
      */
     int Warmup = 5;
 
+    /**
+     * @brief Redirect test requests through this server proxy.
+     *
+     * @details More than one proxy address can be added using semicolon separated format. Do not
+     * use spaces after a semicolon as it would be considered as another command argument. When
+     * multiple proxies are set, each server is assigned to a performance test run on round-robin.
+     *
+     * @note Only the requests from the test are redirected. Any request from set up won't be
+     * redirected.
+     *
+     */
+    std::vector<std::string> TestProxies;
+
     /**
      * @brief Create an array of the performance framework options.
      *

sdk/core/perf/inc/azure/perf/test.hpp

@@ -11,19 +11,15 @@
 
 #include "azure/perf/base_test.hpp"
 #include "azure/perf/dynamic_test_options.hpp"
-#include "azure/perf/options.hpp"
 
 #include <memory>
 
 namespace Azure { namespace Perf {
-
   /**
    * @brief Define a performance test with options.
    *
    */
   class PerfTest : public Azure::Perf::BaseTest {
-  protected:
-    Azure::Perf::TestOptions m_options;
 
   public:
     /**
@@ -31,7 +27,7 @@ namespace Azure { namespace Perf {
      *
      * @param options The command-line parsed options.
      */
-    PerfTest(Azure::Perf::TestOptions options) : m_options(options) {}
+    PerfTest(Azure::Perf::TestOptions options) : BaseTest(options) {}
 
     /**
      * @brief Destroy the Performance Test object.

sdk/core/perf/inc/azure/perf/test_options.hpp

@@ -41,18 +41,18 @@ namespace Azure { namespace Perf {
      * @brief The number of arguments expected after the sentinel for the test option.
      *
      */
-    uint16_t expectedArgs;
+    uint16_t ExpectedArgs;
 
     /**
      * @brief Make an option to be mandatory to run the test.
      *
      */
-    bool required = false;
+    bool Required = false;
 
     /**
      * @brief Make the option to be replaced with **** on all outputs
      *
      */
-    bool sensitiveData = false;
+    bool SensitiveData = false;
   };
 }} // namespace Azure::Perf

sdk/core/perf/src/arg_parser.cpp

@@ -5,6 +5,7 @@
 #include "azure/perf/program.hpp"
 
 #include <stdexcept>
+#include <string>
 #include <vector>
 
 #define GET_ARG(Name, Is)
@@ -20,12 +21,12 @@ argagg::parser_results Azure::Perf::Program::ArgParser::Parse(
   for (auto option : testOptions)
   {
     argParser.definitions.push_back(
-        {option.Name, option.Activators, option.DisplayMessage, option.expectedArgs});
+        {option.Name, option.Activators, option.DisplayMessage, option.ExpectedArgs});
   }
   for (auto option : optionsMetadata)
   {
     argParser.definitions.push_back(
-        {option.Name, option.Activators, option.DisplayMessage, option.expectedArgs});
+        {option.Name, option.Activators, option.DisplayMessage, option.ExpectedArgs});
   }
 
   // Will throw on fail
@@ -93,6 +94,15 @@ Azure::Perf::GlobalTestOptions Azure::Perf::Program::ArgParser::Parse(
   {
     options.Warmup = parsedArgs["Warmup"];
   }
+  if (parsedArgs["TestProxies"])
+  {
+    std::string proxy;
+    std::istringstream fullArg(parsedArgs["TestProxies"].as<std::string>());
+    while (std::getline(fullArg, proxy, ';'))
+    {
+      options.TestProxies.push_back(proxy);
+    }
+  }
 
   return options;
 }

sdk/core/perf/src/base_test.cpp

@@ -0,0 +1,206 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include "azure/perf/base_test.hpp"
+#include "azure/core/http/policies/policy.hpp"
+#include "azure/core/internal/http/pipeline.hpp"
+
+#include <functional>
+#include <string>
+#include <vector>
+
+using namespace Azure::Core::Http;
+using namespace Azure::Core::Http::Policies;
+using namespace Azure::Core;
+namespace {
+
+class ProxyPolicy final : public HttpPolicy {
+private:
+  Azure::Perf::BaseTest* m_testContext;
+
+public:
+  ProxyPolicy(Azure::Perf::BaseTest* proxyManager) : m_testContext(proxyManager) {}
+
+  // copy
+  ProxyPolicy(ProxyPolicy const& other) : ProxyPolicy{other.m_testContext} {}
+
+  // move
+  ProxyPolicy(ProxyPolicy&& other) : m_testContext{other.m_testContext} {}
+
+  std::unique_ptr<RawResponse> Send(
+      Request& request,
+      NextHttpPolicy nextPolicy,
+      Context const& context) const override
+  {
+    std::string const recordId(m_testContext->m_recordId);
+    if (recordId.empty())
+    {
+      return nextPolicy.Send(request, context);
+    }
+
+    // Use a new request to redirect
+    auto redirectRequest = Azure::Core::Http::Request(
+        request.GetMethod(), Azure::Core::Url(m_testContext->m_proxy), request.GetBodyStream());
+    if (!request.ShouldBufferResponse())
+    {
+      // This is a download with keep connection open. Let's switch the request
+      redirectRequest = Azure::Core::Http::Request(
+          request.GetMethod(), Azure::Core::Url(m_testContext->m_proxy), false);
+    }
+    redirectRequest.GetUrl().SetPath(request.GetUrl().GetPath());
+
+    // Copy all headers
+    for (auto& header : request.GetHeaders())
+    {
+      redirectRequest.SetHeader(header.first, header.second);
+    }
+    // QP
+    for (auto const& qp : request.GetUrl().GetQueryParameters())
+    {
+      redirectRequest.GetUrl().AppendQueryParameter(qp.first, qp.second);
+    }
+    // Set x-recording-upstream-base-uri
+    {
+      auto const& url = request.GetUrl();
+      auto const port = url.GetPort();
+      auto const host
+          = url.GetScheme() + "://" + url.GetHost() + (port != 0 ? ":" + std::to_string(port) : "");
+      redirectRequest.SetHeader("x-recording-upstream-base-uri", host);
+    }
+    // Set recording-id
+    redirectRequest.SetHeader("x-recording-id", recordId);
+    redirectRequest.SetHeader("x-recording-remove", "false");
+
+    // Using recordId, find out MODE
+    if (m_testContext->m_isPlayBackMode)
+    {
+      // PLAYBACK mode
+      redirectRequest.SetHeader("x-recording-mode", "playback");
+    }
+    else
+    {
+      // RECORDING mode
+      redirectRequest.SetHeader("x-recording-mode", "record");
+    }
+
+    return nextPolicy.Send(redirectRequest, context);
+  }
+
+  std::unique_ptr<HttpPolicy> Clone() const override
+  {
+    return std::make_unique<ProxyPolicy>(*this);
+  }
+};
+
+} // namespace
+
+namespace Azure { namespace Perf {
+
+  void BaseTest::ConfigureCoreClientOptions(Azure::Core::_internal::ClientOptions* clientOptions)
+  {
+    if (!m_proxy.empty())
+    {
+      clientOptions->PerRetryPolicies.push_back(std::make_unique<ProxyPolicy>(this));
+    }
+  }
+
+  void BaseTest::PostSetUp()
+  {
+
+    if (!m_proxy.empty())
+    {
+      Azure::Core::_internal::ClientOptions clientOp;
+      clientOp.Retry.MaxRetries = 0;
+      std::vector<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policiesOp;
+      std::vector<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policiesRe;
+      Azure::Core::Http::_internal::HttpPipeline pipeline(
+          clientOp, "PerfFw", "na", std::move(policiesRe), std::move(policiesOp));
+      Azure::Core::Context ctx;
+
+      //  Make one call to Run() before starting recording, to avoid capturing one-time setup
+      //  like authorization requests.
+      this->Run(ctx);
+
+      // Send start-record call
+      {
+        Azure::Core::Url startRecordReq(m_proxy);
+        startRecordReq.AppendPath("record");
+        startRecordReq.AppendPath("start");
+        Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Post, startRecordReq);
+        auto response = pipeline.Send(request, ctx);
+
+        auto const& headers = response->GetHeaders();
+        auto findHeader = std::find_if(
+            headers.begin(),
+            headers.end(),
+            [](std::pair<std::string const&, std::string const&> h) {
+              return h.first == "x-recording-id";
+            });
+        m_recordId = findHeader->second;
+      }
+
+      // Record one call to re-use response on all test runs
+      this->Run(ctx);
+
+      // Stop recording
+      {
+        Azure::Core::Url stopRecordReq(m_proxy);
+        stopRecordReq.AppendPath("record");
+        stopRecordReq.AppendPath("stop");
+        Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Post, stopRecordReq);
+        request.SetHeader("x-recording-id", m_recordId);
+        pipeline.Send(request, ctx);
+      }
+
+      // Start playback
+      {
+        Azure::Core::Url startPlayback(m_proxy);
+        startPlayback.AppendPath("playback");
+        startPlayback.AppendPath("start");
+        Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Post, startPlayback);
+        request.SetHeader("x-recording-id", m_recordId);
+        auto response = pipeline.Send(request, ctx);
+
+        auto const& headers = response->GetHeaders();
+        auto findHeader = std::find_if(
+            headers.begin(),
+            headers.end(),
+            [](std::pair<std::string const&, std::string const&> h) {
+              return h.first == "x-recording-id";
+            });
+        m_recordId = findHeader->second;
+        m_isPlayBackMode = true;
+      }
+    }
+  }
+
+  void BaseTest::PreCleanUp()
+  {
+    if (!m_recordId.empty())
+    {
+      Azure::Core::_internal::ClientOptions clientOp;
+      clientOp.Retry.MaxRetries = 0;
+      std::vector<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policiesOp;
+      std::vector<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policiesRe;
+      Azure::Core::Http::_internal::HttpPipeline pipeline(
+          clientOp, "PerfFw", "na", std::move(policiesRe), std::move(policiesOp));
+      Azure::Core::Context ctx;
+
+      // Stop playback
+      {
+        Azure::Core::Url stopPlaybackReq(m_proxy);
+        stopPlaybackReq.AppendPath("playback");
+        stopPlaybackReq.AppendPath("stop");
+        Azure::Core::Http::Request request(Azure::Core::Http::HttpMethod::Post, stopPlaybackReq);
+        request.SetHeader("x-recording-id", m_recordId);
+        request.SetHeader("x-purge-inmemory-recording", "true"); // cspell:disable-line
+
+        pipeline.Send(request, ctx);
+
+        m_recordId.clear();
+        m_isPlayBackMode = false;
+      }
+    }
+  }
+
+}} // namespace Azure::Perf

sdk/core/perf/src/options.cpp

@@ -33,6 +33,14 @@ void Azure::Perf::to_json(Azure::Core::Json::_internal::json& j, const GlobalTes
   {
     j["Rate"] = nullptr;
   }
+  if (p.TestProxies.empty())
+  {
+    j["TestProxies"] = "N/A";
+  }
+  else
+  {
+    j["TestProxies"] = p.TestProxies;
+  }
 }
 
 std::vector<Azure::Perf::TestOption> Azure::Perf::GlobalTestOptions::GetOptionMetadata()
@@ -50,6 +58,7 @@ std::vector<Azure::Perf::TestOption> Azure::Perf::GlobalTestOptions::GetOptionMe
     [Option('r', "rate", HelpText = "Target throughput (ops/sec)")]
     [Option("sync", HelpText = "Runs sync version of test")]  -- Not supported
     [Option('w', "warmup", Default = 5, HelpText = "Duration of warmup in seconds")]
+    [Option('x', "proxy", Default = "", HelpText = "Proxy server")]
   */
   return {
       {"Duration",
@@ -75,5 +84,6 @@ std::vector<Azure::Perf::TestOption> Azure::Perf::GlobalTestOptions::GetOptionMe
       {"Port", {"--port"}, "Port to redirect HTTP requests. Default to no redirection.", 1},
       {"Rate", {"-r", "--rate"}, "Target throughput (ops/sec). Default to no throughput.", 1},
       {"Warmup", {"-w", "--warmup"}, "Duration of warmup in seconds. Default to 5 seconds.", 1},
+      {"TestProxies", {"-x", "--test-proxies"}, "URIs of TestProxy Servers (separated by ';')", 1},
       {"help", {"-h", "--help"}, "Display help information.", 0}};
 }

sdk/core/perf/src/program.cpp

@@ -84,11 +84,11 @@ inline void PrintOptions(
       try
       {
         optionsAsJson[option.Name]
-            = option.sensitiveData ? "***" : parsedArgs[option.Name].as<std::string>();
+            = option.SensitiveData ? "***" : parsedArgs[option.Name].as<std::string>();
       }
       catch (std::out_of_range const&)
       {
-        if (!option.required)
+        if (!option.Required)
         {
           // arg was not parsed
           optionsAsJson[option.Name] = "default value";
@@ -203,7 +203,8 @@ inline void RunTests(
   uint64_t lastCompleted = 0;
   auto progressThread = std::thread(
       [&title, &completedOperations, &lastCompletionTimes, &lastCompleted, &progresToken]() {
-        std::cout << "=== " << title << " ===" << std::endl
+        std::cout << std::endl
+                  << "=== " << title << " ===" << std::endl
                   << "Current\t\tTotal\t\tAverage" << std::endl;
         while (!progresToken.IsCancelled())
         {
@@ -293,7 +294,7 @@ void Azure::Perf::Program::Run(
   argResults = Azure::Perf::Program::ArgParser::Parse(argc, argv, testOptions);
   // ReCreate Test with parsed results
   test = testGenerator(Azure::Perf::TestOptions(argResults));
-  auto options = Azure::Perf::Program::ArgParser::Parse(argResults);
+  GlobalTestOptions options = Azure::Perf::Program::ArgParser::Parse(argResults);
 
   if (options.JobStatistics)
   {
@@ -313,11 +314,20 @@ void Azure::Perf::Program::Run(
   for (int i = 0; i < parallelTasks; i++)
   {
     parallelTest[i] = testGenerator(Azure::Perf::TestOptions(argResults));
+    // Let the test know it should use a proxy
+    if (!options.TestProxies.empty())
+    {
+      // Take the corresponding proxy from the list in round robin
+      parallelTest[i]->SetTestProxy(options.TestProxies[i % options.TestProxies.size()]);
+    }
   }
 
   /******************** Global Set up ******************************/
+  std::cout << std::endl << "=== Global Setup ===" << std::endl;
   test->GlobalSetup();
 
+  std::cout << std::endl << "=== Test Setup ===" << std::endl;
+
   /******************** Set up ******************************/
   {
     std::vector<std::thread> tasks(parallelTasks);
@@ -332,6 +342,27 @@ void Azure::Perf::Program::Run(
     }
   }
 
+  // instrument test for recordings if the env is set up.
+  std::cout << std::endl << "=== Post Setup ===" << std::endl;
+  {
+    if (!options.TestProxies.empty())
+    {
+      std::cout << " - Creating test recordgins for each test using test-proxies..." << std::endl;
+      std::cout << " - Enabling test-proxy playback" << std::endl;
+    }
+
+    std::vector<std::thread> tasks(parallelTasks);
+    for (int i = 0; i < parallelTasks; i++)
+    {
+      tasks[i] = std::thread([&parallelTest, i]() { parallelTest[i]->PostSetUp(); });
+    }
+    // Wait for all tests to complete setUp
+    for (auto& t : tasks)
+    {
+      t.join();
+    }
+  }
+
   /******************** WarmUp ******************************/
   if (options.Warmup)
   {
@@ -356,6 +387,26 @@ void Azure::Perf::Program::Run(
     std::cout << "Error: " << error.what();
   }
 
+  std::cout << std::endl << "=== Pre-Cleanup ===" << std::endl;
+  {
+    if (!options.TestProxies.empty())
+    {
+      std::cout << " - Deleting test recordings from test-proxies..." << std::endl;
+      std::cout << " - Disabling test-proxy playback" << std::endl;
+    }
+
+    std::vector<std::thread> tasks(parallelTasks);
+    for (int i = 0; i < parallelTasks; i++)
+    {
+      tasks[i] = std::thread([&parallelTest, i]() { parallelTest[i]->PreCleanUp(); });
+    }
+    // Wait for all tests to complete setUp
+    for (auto& t : tasks)
+    {
+      t.join();
+    }
+  }
+
   /******************** Clean up ******************************/
   if (!options.NoCleanup)
   {

sdk/identity/azure-identity/test/perf/inc/azure/identity/test/secret_credential_test.hpp

@@ -43,7 +43,10 @@ namespace Azure { namespace Identity { namespace Test {
       m_secret = m_options.GetMandatoryOption<std::string>("Secret");
       m_tokenRequestContext.Scopes.push_back(m_options.GetMandatoryOption<std::string>("Scope"));
       m_credential = std::make_unique<Azure::Identity::ClientSecretCredential>(
-          m_tenantId, m_clientId, m_secret);
+          m_tenantId,
+          m_clientId,
+          m_secret,
+          InitClientOptions<Azure::Core::Credentials::TokenCredentialOptions>());
     }
 
     /**

sdk/keyvault/azure-security-keyvault-keys/test/perf/inc/azure/keyvault/keys/test/get_key_test.hpp

@@ -48,8 +48,10 @@ namespace Azure { namespace Security { namespace KeyVault { namespace Keys { nam
       m_secret = m_options.GetMandatoryOption<std::string>("Secret");
       m_credential = std::make_shared<Azure::Identity::ClientSecretCredential>(
           m_tenantId, m_clientId, m_secret);
-      m_client
-          = std::make_unique<Azure::Security::KeyVault::Keys::KeyClient>(m_vaultUrl, m_credential);
+      m_client = std::make_unique<Azure::Security::KeyVault::Keys::KeyClient>(
+          m_vaultUrl,
+          m_credential,
+          InitClientOptions<Azure::Security::KeyVault::Keys::KeyClientOptions>());
     }
 
     /**

sdk/storage/azure-storage-blobs/test/perf/inc/azure/storage/blobs/test/blob_base_test.hpp

@@ -73,7 +73,8 @@ namespace Azure { namespace Storage { namespace Blobs { namespace Test {
 
       // Create client, container and blobClient
       m_serviceClient = std::make_unique<Azure::Storage::Blobs::BlobServiceClient>(
-          Azure::Storage::Blobs::BlobServiceClient::CreateFromConnectionString(m_connectionString));
+          Azure::Storage::Blobs::BlobServiceClient::CreateFromConnectionString(
+              m_connectionString, InitClientOptions<Azure::Storage::Blobs::BlobClientOptions>()));
       m_containerClient = std::make_unique<Azure::Storage::Blobs::BlobContainerClient>(
           m_serviceClient->GetBlobContainerClient(m_containerName));
       m_containerClient->CreateIfNotExists();