Add a virtual ShouldRetry method to the RetryPolicy for customization.

sdk/core/azure-core/inc/azure/core/http/policies/policy.hpp

@@ -363,14 +363,28 @@ 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;
 
+#if defined(_azure_TESTING_BUILD)
+    protected:
+#else
+    private:
+#endif
+      virtual bool ShouldRetryOnTransportFailure(
+          RetryOptions const& retryOptions,
+          int32_t attempt,
+          std::chrono::milliseconds& retryAfter,
+          double jitterFactor = -1) const;
+
+      virtual bool ShouldRetryOnResponse(
+          RawResponse const& response,
+          RetryOptions const& retryOptions,
+          int32_t attempt,
+          std::chrono::milliseconds& retryAfter,
+          double jitterFactor = -1) const;
+
     public:
       /**
        * Constructs HTTP retry policy with the provided #Azure::Core::Http::Policies::RetryOptions.
@@ -403,18 +417,25 @@ namespace Azure { namespace Core { namespace Http { namespace Policies {
       static int32_t GetRetryCount(Context const& context);
 
     protected:
-      virtual bool ShouldRetryOnTransportFailure(
-          RetryOptions const& retryOptions,
-          int32_t attempt,
-          std::chrono::milliseconds& retryAfter,
-          double jitterFactor = -1) const;
-
-      virtual bool ShouldRetryOnResponse(
-          RawResponse const& response,
-          RetryOptions const& retryOptions,
-          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,15 @@ int32_t RetryPolicy::GetRetryCount(Context const& context)
   return *ptr;
 }
 
+bool RetryPolicy::ShouldRetry(
+    std::unique_ptr<RawResponse> const& response,
+    RetryOptions const& retryOptions) const
+{
+  (void)response;
+  (void)retryOptions;
+  return true;
+}
+
 std::unique_ptr<RawResponse> RetryPolicy::Send(
     Request& request,
     NextHttpPolicy nextPolicy,
@@ -141,13 +150,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;