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.

Author: ahsonkhan

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<RawResponse> const& response,
+          RetryOptions const& retryOptions) const;
     };
 
     /**

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<RawResponse> const&, RetryOptions const&) const
+{
+  return true;
+}
+
 std::unique_ptr<RawResponse> RetryPolicy::Send(
     Request& request,
     NextHttpPolicy nextPolicy,
@@ -141,13 +146,20 @@ std::unique_ptr<RawResponse> 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)
     {

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<HttpPolicy> Clone() const override
+  {
+    return std::make_unique<RetryPolicyTest_CustomRetry_False>(*this);
+  }
+
+protected:
+  bool ShouldRetry(std::unique_ptr<RawResponse> 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<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policies;
+    policies.emplace_back(std::make_unique<RetryPolicy>(policy));
+    policies.emplace_back(std::make_unique<TestTransportPolicy>([&]() {
+      auto response = std::make_unique<RawResponse>(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<std::unique_ptr<Azure::Core::Http::Policies::HttpPolicy>> policies;
+    policies.emplace_back(std::make_unique<RetryPolicyTest_CustomRetry_False>(policy));
+
+    policies.emplace_back(std::make_unique<TestTransportPolicy>([&]() {
+      auto response = std::make_unique<RawResponse>(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;