Add Decider argument to retryN and retryWithExponentialBackoff

Summary:
This allows customizing the exceptions for which retry
is triggered. This is useful when we only want to retry
in case of a few specific exceptions but have others
cause the operation to fail immediately.

Reviewed By: lewissbaker

Differential Revision: D24766615

fbshipit-source-id: f97451673f575bef511399cbde6c1ad110f9493a

folly/experimental/coro/Retry.h

@@ -87,13 +87,18 @@ auto retryWhen(Func func, RetryDelayFunc retryDelay)
 
 namespace detail {
 
+template <typename Decider>
 class RetryImmediatelyWithLimit {
  public:
-  explicit RetryImmediatelyWithLimit(uint32_t maxRetries) noexcept
-      : retriesRemaining_(maxRetries) {}
+  template <typename Decider2>
+  explicit RetryImmediatelyWithLimit(
+      uint32_t maxRetries,
+      Decider2&& decider) noexcept
+      : retriesRemaining_(maxRetries),
+        decider_(static_cast<Decider2&&>(decider)) {}
 
   Task<void> operator()(exception_wrapper&& ew) & {
-    if (retriesRemaining_ == 0) {
+    if (retriesRemaining_ == 0 || !decider_(ew)) {
       co_yield folly::coro::co_error(std::move(ew));
     }
 
@@ -107,41 +112,56 @@ class RetryImmediatelyWithLimit {
 
  private:
   uint32_t retriesRemaining_;
+  Decider decider_;
+};
+
+struct AlwaysRetry {
+  bool operator()(const folly::exception_wrapper&) noexcept { return true; }
 };
 
 } // namespace detail
 
 // Executes the operation returned by func(), retrying it up to
 // 'maxRetries' times on failure with no delay between retries.
+template <typename Func, typename Decider>
+auto retryN(uint32_t maxRetries, Func&& func, Decider&& decider) {
+  return folly::coro::retryWhen(
+      static_cast<Func&&>(func),
+      detail::RetryImmediatelyWithLimit<remove_cvref_t<Decider>>{
+          maxRetries, static_cast<Decider&&>(decider)});
+}
+
 template <typename Func>
 auto retryN(uint32_t maxRetries, Func&& func) {
-  return folly::coro::retryWhen(
-      static_cast<Func&&>(func), detail::RetryImmediatelyWithLimit{maxRetries});
+  return folly::coro::retryN(
+      maxRetries, static_cast<Func&&>(func), detail::AlwaysRetry{});
 }
 
 namespace detail {
 
-template <typename URNG>
+template <typename URNG, typename Decider>
 class ExponentialBackoffWithJitter {
  public:
-  template <typename URNG2>
+  template <typename URNG2, typename Decider2>
   explicit ExponentialBackoffWithJitter(
       Timekeeper* tk,
       uint32_t maxRetries,
       Duration minBackoff,
       Duration maxBackoff,
       double relativeJitterStdDev,
-      URNG2&& rng) noexcept
+      URNG2&& rng,
+      Decider2&& decider) noexcept
       : timeKeeper_(tk),
         maxRetries_(maxRetries),
         retryCount_(0),
         minBackoff_(minBackoff),
         maxBackoff_(maxBackoff),
         relativeJitterStdDev_(relativeJitterStdDev),
-        randomGen_(static_cast<URNG2&&>(rng)) {}
+        randomGen_(static_cast<URNG2&&>(rng)),
+        decider_(static_cast<Decider2&&>(decider)) {}
 
   Task<void> operator()(exception_wrapper&& ew) & {
-    if (retryCount_ == maxRetries_) {
+    if (retryCount_ == maxRetries_ || !decider_(ew)) {
       co_yield folly::coro::co_error(std::move(ew));
     }
 
@@ -179,6 +199,7 @@ class ExponentialBackoffWithJitter {
   const Duration maxBackoff_;
   const double relativeJitterStdDev_;
   URNG randomGen_;
+  Decider decider_;
 };
 
 } // namespace detail
@@ -186,8 +207,9 @@ class ExponentialBackoffWithJitter {
 // Executes the operation returned from 'func()', retrying it on failure
 // up to 'maxRetries' times, with an exponential backoff, doubling the backoff
 // on average for each retry, applying some random jitter, up to the specified
-// maximum backoff.
-template <typename Func, typename URNG>
+// maximum backoff, passing each error to decider to decide whether to retry or
+// not.
+template <typename Func, typename URNG, typename Decider>
 auto retryWithExponentialBackoff(
     uint32_t maxRetries,
     Duration minBackoff,
@@ -195,16 +217,39 @@ auto retryWithExponentialBackoff(
     double relativeJitterStdDev,
     Timekeeper* timeKeeper,
     URNG&& rng,
-    Func&& func) {
+    Func&& func,
+    Decider&& decider) {
   return folly::coro::retryWhen(
       static_cast<Func&&>(func),
-      detail::ExponentialBackoffWithJitter<remove_cvref_t<URNG>>{
-          timeKeeper,
+      detail::ExponentialBackoffWithJitter<
+          remove_cvref_t<URNG>,
+          remove_cvref_t<Decider>>{timeKeeper,
+                                   maxRetries,
+                                   minBackoff,
+                                   maxBackoff,
+                                   relativeJitterStdDev,
+                                   static_cast<URNG&&>(rng),
+                                   static_cast<Decider&&>(decider)});
+}
+
+template <typename Func, typename URNG>
+auto retryWithExponentialBackoff(
+    uint32_t maxRetries,
+    Duration minBackoff,
+    Duration maxBackoff,
+    double relativeJitterStdDev,
+    Timekeeper* timeKeeper,
+    URNG&& rng,
+    Func&& func) {
+  return folly::coro::retryWithExponentialBackoff(
       maxRetries,
       minBackoff,
       maxBackoff,
       relativeJitterStdDev,
-          static_cast<URNG&&>(rng)});
+      timeKeeper,
+      static_cast<URNG&&>(rng),
+      static_cast<Func&&>(func),
+      detail::AlwaysRetry{});
 }
 
 template <typename Func>
@@ -241,4 +286,23 @@ auto retryWithExponentialBackoff(
       static_cast<Func&&>(func));
 }
 
+template <typename Func, typename Decider>
+auto retryWithExponentialBackoff(
+    uint32_t maxRetries,
+    Duration minBackoff,
+    Duration maxBackoff,
+    double relativeJitterStdDev,
+    Func&& func,
+    Decider&& decider) {
+  return folly::coro::retryWithExponentialBackoff(
+      maxRetries,
+      minBackoff,
+      maxBackoff,
+      relativeJitterStdDev,
+      static_cast<Timekeeper*>(nullptr),
+      ThreadLocalPRNG(),
+      static_cast<Func&&>(func),
+      static_cast<Decider&&>(decider));
+}
+
 } // namespace folly::coro

folly/experimental/coro/test/RetryTest.cpp

@@ -78,6 +78,23 @@ TEST(RetryN, EventualSuccess) {
   }());
 }
 
+TEST(RetryN, NeverRetry) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    int runCount = 0;
+    folly::Try<void> result =
+        co_await folly::coro::co_awaitTry(folly::coro::retryN(
+            3,
+            [&]() -> folly::coro::Task<void> {
+              ++runCount;
+              if (runCount <= 2) {
+                co_yield folly::coro::co_error(SomeError{runCount});
+              }
+            },
+            [](const folly::exception_wrapper&) { return false; }));
+    EXPECT_EQ(1, runCount);
+  }());
+}
+
 TEST(RetryWithJitter, Success) {
   folly::coro::blockingWait([]() -> folly::coro::Task<void> {
     int runCount = 0;
@@ -146,4 +163,62 @@ TEST(RetryWithJitter, EventualSuccess) {
   }());
 }
 
+TEST(RetryWithDecider, AlwaysRetry) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    int runCount = 0;
+    folly::Try<void> result = co_await folly::coro::co_awaitTry(
+        folly::coro::retryWithExponentialBackoff(
+            5, 0ms, 1ms, 0.0, [&]() -> folly::coro::Task<void> {
+              ++runCount;
+              co_yield folly::coro::co_error(SomeError(1));
+            }));
+    EXPECT_TRUE(result.hasException());
+    EXPECT_EQ(6, runCount);
+  }());
+}
+
+TEST(RetryWithDecider, NeverRetry) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    int runCount = 0;
+    folly::Try<void> result = co_await folly::coro::co_awaitTry(
+        folly::coro::retryWithExponentialBackoff(
+            5,
+            0ms,
+            1ms,
+            0.0,
+            [&]() -> folly::coro::Task<void> {
+              ++runCount;
+              co_yield folly::coro::co_error(SomeError(1));
+            },
+            [](const folly::exception_wrapper&) { return false; }));
+    EXPECT_TRUE(result.hasException());
+    EXPECT_EQ(1, runCount);
+  }());
+}
+
+TEST(RetryWithDecider, SometimesRetry) {
+  folly::coro::blockingWait([]() -> folly::coro::Task<void> {
+    int runCount = 0;
+    folly::Try<void> result = co_await folly::coro::co_awaitTry(
+        folly::coro::retryWithExponentialBackoff(
+            5,
+            0ms,
+            1ms,
+            0.0,
+            [&]() -> folly::coro::Task<void> {
+              ++runCount;
+              co_yield folly::coro::co_error(SomeError(runCount));
+            },
+            [](const folly::exception_wrapper& ew) {
+              try {
+                ew.throw_exception();
+              } catch (const SomeError& e) {
+                return e.value < 3;
+              }
+            }));
+    EXPECT_TRUE(result.hasException());
+    EXPECT_EQ(3, runCount);
+  }());
+}
+
 #endif // FOLLY_HAS_COROUTINES