From ab90ef68b03684def81492e13542b02406160a85 Mon Sep 17 00:00:00 2001 From: Ahson Khan Date: Mon, 13 May 2024 18:09:02 -0700 Subject: [PATCH] Add a virtual ShouldRetry method to the RetryPolicy for customization. (#5584) * Add a virtual ShouldTry method to the RetryPolicy for customization. * Make test hooks virtual again. * Add unit test for ShouldRetry. * Fix typo. * Address PR feedback. * Revert making the mock'd functions private for non-test builds. * Add new line. --- .../inc/azure/core/http/policies/policy.hpp | 26 ++++-- sdk/core/azure-core/src/http/retry_policy.cpp | 14 +++- .../azure-core/test/ut/retry_policy_test.cpp | 81 +++++++++++++++++++ 3 files changed, 115 insertions(+), 6 deletions(-) diff --git a/sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp b/sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp index 511a1c3a7a..e626dbc48b 100644 --- a/sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp +++ b/sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp @@ -363,11 +363,7 @@ namespace Azure { namespace Core { namespace Http { namespace Policies { /** * @brief HTTP retry policy. */ - class RetryPolicy -#if !defined(_azure_TESTING_BUILD) - final -#endif - : public HttpPolicy { + class RetryPolicy : public HttpPolicy { private: RetryOptions m_retryOptions; @@ -415,6 +411,26 @@ namespace Azure { namespace Core { namespace Http { namespace Policies { int32_t attempt, std::chrono::milliseconds& retryAfter, double jitterFactor = -1) const; + + /** + * @brief Overriding this method customizes the logic of when the RetryPolicy will re-attempt + * a request, based on the returned HTTP response. + * + * @remark A null response pointer means there was no response received from the corresponding + * request. Custom implementations of this method that override the retry behavior, should + * handle that error case, if that needs to be customized. + * + * @remark Unless overriden, the default implementation is to always return true. The + * non-retriable errors, including those specified in the RetryOptions, remain evaluated + * before calling ShouldRetry. + * + * @param response An HTTP response returned corresponding to the request sent by the policy. + * @param retryOptions The set of options provided to the RetryPolicy. + * @return Whether or not the HTTP request should be sent again through the pipeline. + */ + virtual bool ShouldRetry( + std::unique_ptr const& response, + RetryOptions const& retryOptions) const; }; /** diff --git a/sdk/core/azure-core/src/http/retry_policy.cpp b/sdk/core/azure-core/src/http/retry_policy.cpp index 217e0bd898..29eeae8e4a 100644 --- a/sdk/core/azure-core/src/http/retry_policy.cpp +++ b/sdk/core/azure-core/src/http/retry_policy.cpp @@ -118,6 +118,11 @@ int32_t RetryPolicy::GetRetryCount(Context const& context) return *ptr; } +bool RetryPolicy::ShouldRetry(std::unique_ptr const&, RetryOptions const&) const +{ + return true; +} + std::unique_ptr RetryPolicy::Send( Request& request, NextHttpPolicy nextPolicy, @@ -141,13 +146,20 @@ std::unique_ptr RetryPolicy::Send( auto response = nextPolicy.Send(request, retryContext); // If we are out of retry attempts, if a response is non-retriable (or simply 200 OK, i.e - // doesn't need to be retried), then ShouldRetry returns false. + // doesn't need to be retried), then ShouldRetryOnResponse returns false. if (!ShouldRetryOnResponse(*response.get(), m_retryOptions, attempt, retryAfter)) { // If this is the second attempt and StartTry was called, we need to stop it. Otherwise // trying to perform same request would use last retry query/headers return response; } + + // Service SDKs can inject custom logic to define whether the request should be retried, + // based on the response. The default is true. + if (!ShouldRetry(response, m_retryOptions)) + { + return response; + } } catch (const TransportException& e) { diff --git a/sdk/core/azure-core/test/ut/retry_policy_test.cpp b/sdk/core/azure-core/test/ut/retry_policy_test.cpp index b482418ce5..2679d78e0a 100644 --- a/sdk/core/azure-core/test/ut/retry_policy_test.cpp +++ b/sdk/core/azure-core/test/ut/retry_policy_test.cpp @@ -124,8 +124,89 @@ class RetryPolicyTest final : public RetryPolicy { return m_shouldRetryOnResponse(response, retryOptions, attempt, retryAfter, jitterFactor); } }; + +class RetryPolicyTest_CustomRetry_False final : public RetryPolicy { +public: + RetryPolicyTest_CustomRetry_False(RetryOptions const& retryOptions) : RetryPolicy(retryOptions) {} + + std::unique_ptr Clone() const override + { + return std::make_unique(*this); + } + +protected: + bool ShouldRetry(std::unique_ptr const& response, RetryOptions const& retryOptions) + const override + { + (void)response; + (void)retryOptions; + return false; + } +}; } // namespace +TEST(RetryPolicy, ShouldRetry) +{ + using namespace std::chrono_literals; + RetryOptions const retryOptions{3, 10ms, 60s, {HttpStatusCode::Ok}}; + + // The default ShouldRetry implementation is to always return true, + // which means we will retry up to the retry count in the options. + { + Azure::Core::Context context = Azure::Core::Context::ApplicationContext; + auto policy = RetryPolicy(retryOptions); + + RawResponse const* responsePtrSent = nullptr; + int32_t retryCounter = -1; + + std::vector> policies; + policies.emplace_back(std::make_unique(policy)); + policies.emplace_back(std::make_unique([&]() { + auto response = std::make_unique(1, 1, HttpStatusCode::Ok, "Test"); + + responsePtrSent = response.get(); + retryCounter++; + return response; + })); + + Azure::Core::Http::_internal::HttpPipeline pipeline(policies); + + Request request(HttpMethod::Get, Azure::Core::Url("https://www.microsoft.com")); + pipeline.Send(request, context); + + EXPECT_NE(responsePtrSent, nullptr); + EXPECT_EQ(retryCounter, 3); + } + + // Overriding ShouldRetry to return false will mean we won't retry. + { + Azure::Core::Context context = Azure::Core::Context(); + auto policy = RetryPolicyTest_CustomRetry_False(retryOptions); + + RawResponse const* responsePtrSent = nullptr; + int32_t retryCounter = -1; + + std::vector> policies; + policies.emplace_back(std::make_unique(policy)); + + policies.emplace_back(std::make_unique([&]() { + auto response = std::make_unique(1, 1, HttpStatusCode::Ok, "Test"); + + responsePtrSent = response.get(); + retryCounter++; + return response; + })); + + Azure::Core::Http::_internal::HttpPipeline pipeline(policies); + + Request request(HttpMethod::Get, Azure::Core::Url("https://www.microsoft.com")); + pipeline.Send(request, context); + + EXPECT_NE(responsePtrSent, nullptr); + EXPECT_EQ(retryCounter, 0); + } +} + TEST(RetryPolicy, ShouldRetryOnResponse) { using namespace std::chrono_literals;