impl(v3): add attempt predicate to AsyncRetryLoop (#15915)

Author: scotthart

google/cloud/internal/async_retry_loop.h

@@ -170,16 +170,23 @@ struct FutureValueType<future<T>> {
  * functions.  If the value is visible, the retry loop will stop on the next
  * callback and/or before the next request or timer is issued.
  */
-template <typename Functor, typename Request, typename RetryPolicyType>
+template <
+    typename Functor, typename Request, typename RetryPolicyType,
+    typename ReturnType = google::cloud::internal::invoke_result_t<
+        Functor, google::cloud::CompletionQueue&,
+        std::shared_ptr<grpc::ClientContext>, ImmutableOptions, Request const&>>
 class AsyncRetryLoopImpl
     : public std::enable_shared_from_this<
           AsyncRetryLoopImpl<Functor, Request, RetryPolicyType>> {
  public:
-  AsyncRetryLoopImpl(std::unique_ptr<RetryPolicyType> retry_policy,
-                     std::unique_ptr<BackoffPolicy> backoff_policy,
-                     Idempotency idempotency, google::cloud::CompletionQueue cq,
-                     Functor&& functor, ImmutableOptions options,
-                     Request request, char const* location)
+  AsyncRetryLoopImpl(
+      std::unique_ptr<RetryPolicyType> retry_policy,
+      std::unique_ptr<BackoffPolicy> backoff_policy, Idempotency idempotency,
+      google::cloud::CompletionQueue cq, Functor&& functor,
+      ImmutableOptions options, Request request, char const* location,
+      std::function<
+          bool(typename FutureValueType<ReturnType>::value_type const&)>
+          attempt_predicate = {})
       : retry_policy_(std::move(retry_policy)),
         backoff_policy_(std::move(backoff_policy)),
         idempotency_(idempotency),
@@ -188,11 +195,9 @@ class AsyncRetryLoopImpl
         functor_(std::forward<Functor>(functor)),
         request_(std::move(request)),
         location_(location),
-        call_context_(std::move(options)) {}
+        call_context_(std::move(options)),
+        attempt_predicate_(std::move(attempt_predicate)) {}
 
-  using ReturnType = ::google::cloud::internal::invoke_result_t<
-      Functor, google::cloud::CompletionQueue&,
-      std::shared_ptr<grpc::ClientContext>, ImmutableOptions, Request const&>;
   using T = typename FutureValueType<ReturnType>::value_type;
 
   future<T> Start() {
@@ -256,8 +261,11 @@ class AsyncRetryLoopImpl
   }
 
   void OnAttempt(T result) {
-    // A successful attempt, set the value and finish the loop.
-    if (result.ok()) return SetDone(std::move(result));
+    // If the attempt is successful and satisfies the attempt predicate, if
+    // provided, set the value and finish the loop.
+    if (result.ok() && (!attempt_predicate_ || attempt_predicate_(result))) {
+      return SetDone(std::move(result));
+    }
     // Some kind of failure, first verify that it is retryable.
     last_status_ = GetResultStatus(std::move(result));
     auto delay =
@@ -325,6 +333,8 @@ class AsyncRetryLoopImpl
   CallContext call_context_;
   Status last_status_;
   promise<T> result_;
+  std::function<bool(typename FutureValueType<ReturnType>::value_type const&)>
+      attempt_predicate_;
 
   // Only the following variables require synchronization, as they coordinate
   // the work between the retry loop (which would be lock-free) and the cancel
@@ -339,17 +349,23 @@ class AsyncRetryLoopImpl
 /**
  * Create the right AsyncRetryLoopImpl object and start the retry loop on it.
  */
-template <typename Functor, typename Request, typename RetryPolicyType,
-          std::enable_if_t<google::cloud::internal::is_invocable<
-                               Functor, google::cloud::CompletionQueue&,
-                               std::shared_ptr<grpc::ClientContext>,
-                               ImmutableOptions, Request const&>::value,
-                           int> = 0>
-auto AsyncRetryLoop(std::unique_ptr<RetryPolicyType> retry_policy,
-                    std::unique_ptr<BackoffPolicy> backoff_policy,
-                    Idempotency idempotency, google::cloud::CompletionQueue cq,
-                    Functor&& functor, ImmutableOptions options,
-                    Request request, char const* location)
+template <
+    typename Functor, typename Request, typename RetryPolicyType,
+    std::enable_if_t<google::cloud::internal::is_invocable<
+                         Functor, google::cloud::CompletionQueue&,
+                         std::shared_ptr<grpc::ClientContext>, ImmutableOptions,
+                         Request const&>::value,
+                     int> = 0,
+    typename ReturnType = google::cloud::internal::invoke_result_t<
+        Functor, google::cloud::CompletionQueue&,
+        std::shared_ptr<grpc::ClientContext>, ImmutableOptions, Request const&>>
+auto AsyncRetryLoop(
+    std::unique_ptr<RetryPolicyType> retry_policy,
+    std::unique_ptr<BackoffPolicy> backoff_policy, Idempotency idempotency,
+    google::cloud::CompletionQueue cq, Functor&& functor,
+    ImmutableOptions options, Request request, char const* location,
+    std::function<bool(typename FutureValueType<ReturnType>::value_type const&)>
+        attempt_predicate = {})
     -> google::cloud::internal::invoke_result_t<
         Functor, google::cloud::CompletionQueue&,
         std::shared_ptr<grpc::ClientContext>, ImmutableOptions,
@@ -358,7 +374,7 @@ auto AsyncRetryLoop(std::unique_ptr<RetryPolicyType> retry_policy,
       std::make_shared<AsyncRetryLoopImpl<Functor, Request, RetryPolicyType>>(
           std::move(retry_policy), std::move(backoff_policy), idempotency,
           std::move(cq), std::forward<Functor>(functor), options,
-          std::move(request), location);
+          std::move(request), location, std::move(attempt_predicate));
   return loop->Start();
 }
 

google/cloud/internal/async_retry_loop_test.cc

@@ -160,6 +160,32 @@ TEST(AsyncRetryLoopTest, TransientThenSuccess) {
   EXPECT_EQ(84, *actual);
 }
 
+TEST(AsyncRetryLoopTest, TransientPredicateThenSuccess) {
+  AutomaticallyCreatedBackgroundThreads background;
+  ::testing::MockFunction<bool(StatusOr<int>)> mock_predicate;
+  EXPECT_CALL(mock_predicate, Call)
+      .WillOnce([](StatusOr<int> const&) { return false; })
+      .WillOnce([](StatusOr<int> const&) { return false; })
+      .WillOnce([](StatusOr<int> const&) { return true; });
+
+  auto pending = AsyncRetryLoop(
+      TestRetryPolicy(), TestBackoffPolicy(), Idempotency::kIdempotent,
+      background.cq(),
+      [&](google::cloud::CompletionQueue&, auto,
+          ImmutableOptions const& options, int request) {
+        EXPECT_EQ(options->get<TestOption>(), "TransientPredicateThenSuccess");
+        return make_ready_future(StatusOr<int>(2 * request));
+      },
+      MakeImmutableOptions(
+          Options{}.set<TestOption>("TransientPredicateThenSuccess")),
+      42, "error message", mock_predicate.AsStdFunction());
+
+  OptionsSpan overlay(Options{}.set<TestOption>("uh-oh"));
+  StatusOr<int> actual = pending.get();
+  ASSERT_THAT(actual.status(), IsOk());
+  EXPECT_EQ(84, *actual);
+}
+
 TEST(AsyncRetryLoopTest, ReturnJustStatus) {
   int counter = 0;
   AutomaticallyCreatedBackgroundThreads background;