SERVER-35215 Future::onError<ErrorCodes::Error>()

(cherry picked from commit fa9d7223cf59ee1b5681ea129fe86f55c608c632)

2 files changed, 156 insertions, 9 deletions

diff --git a/src/mongo/util/future.h b/src/mongo/util/future.h
index 7a8fece88ab..fc1fda53c23 100644
--- a/src/mongo/util/future.h
+++ b/src/mongo/util/future.h
@@ -895,13 +895,11 @@ public:
      * general error handler for the entire chain.
      */
     template <typename Func,  // Status -> T or Status -> StatusWith<T>
-              typename RawResult = NormalizedCallResult<Func, Status>,
-              typename = std::enable_if_t<!isFuture<RawResult>>>
+              typename Result = RawNormalizedCallResult<Func, Status>,
+              typename = std::enable_if_t<!isFuture<Result>>>
         Future<T> onError(Func&& func) && noexcept {
         static_assert(
-            std::is_same<RawResult, T>::value || std::is_same<RawResult, StatusWith<T>>::value ||
-                (std::is_same<T, FakeVoid>::value &&
-                 (std::is_same<RawResult, void>::value || std::is_same<RawResult, Status>::value)),
+            std::is_same<Result, T>::value,
             "func passed to Future<T>::onError must return T, StatusWith<T>, or Future<T>");
 
         return generalImpl(
@@ -927,13 +925,13 @@ public:
      * Same as above onError() but for case where func returns a Future that needs to be unwrapped.
      */
     template <typename Func,  // Status -> Future<T>
-              typename RawResult = NormalizedCallResult<Func, Status>,
-              typename = std::enable_if_t<isFuture<RawResult>>,
+              typename Result = RawNormalizedCallResult<Func, Status>,
+              typename = std::enable_if_t<isFuture<Result>>,
               typename = void>
         Future<T> onError(Func&& func) && noexcept {
         static_assert(
-            std::is_same<RawResult, Future<T>>::value ||
-                (std::is_same<T, FakeVoid>::value && std::is_same<RawResult, Future<void>>::value),
+            std::is_same<Result, Future<T>>::value ||
+                (std::is_same<T, FakeVoid>::value && std::is_same<Result, Future<void>>::value),
             "func passed to Future<T>::onError must return T, StatusWith<T>, or Future<T>");
 
         return generalImpl(
@@ -964,6 +962,31 @@ public:
     }
 
     /**
+     * Same as the other two onErrors but only calls the callback if the code matches the template
+     * parameter. Otherwise lets the error propagate unchanged.
+     */
+    template <ErrorCodes::Error code, typename Func>
+        Future<T> onError(Func&& func) && noexcept {
+        using Result = RawNormalizedCallResult<Func, Status>;
+        static_assert(
+            std::is_same<Result, T>::value || std::is_same<Result, Future<T>>::value ||
+                (std::is_same<T, FakeVoid>::value && std::is_same<Result, Future<void>>::value),
+            "func passed to Future<T>::onError must return T, StatusWith<T>, or Future<T>");
+
+        if (immediate || (isReady() && shared->status.isOK()))
+            return std::move(*this);  // Avoid copy/moving func if we know we won't call it.
+
+        // TODO in C++17 with constexpr if this can be done cleaner and more efficiently by not
+        // throwing.
+        return std::move(*this).onError([func =
+                                             std::forward<Func>(func)](Status && status) mutable {
+            if (status != code)
+                uassertStatusOK(status);
+            return throwingCall(func, std::move(status));
+        });
+    }
+
+    /**
      * TODO do we need a version of then/onError like onCompletion() that handles both success and
      * Failure, but doesn't end the chain like getAsync()? Right now we don't, and we can add one if
      * we do.
@@ -1224,6 +1247,11 @@ public:
         return std::move(inner).onError(std::forward<Func>(func));
     }
 
+    template <ErrorCodes::Error code, typename Func>  // Status -> T or StatusWith<T> or Future<T>
+        Future<void> onError(Func&& func) && noexcept {
+        return std::move(inner).onError<code>(std::forward<Func>(func));
+    }
+
     template <typename Func>  // () -> void
         Future<void> tap(Func&& func) && noexcept {
         return std::move(inner).tap(std::forward<Func>(func));
diff --git a/src/mongo/util/future_test.cpp b/src/mongo/util/future_test.cpp
index a5cbebd9ada..caf30dbb019 100644
--- a/src/mongo/util/future_test.cpp
+++ b/src/mongo/util/future_test.cpp
@@ -513,6 +513,64 @@ TEST(Future, Fail_onErrorFutureAsync) {
     });
 }
 
+TEST(Future, Success_onErrorCode) {
+    FUTURE_SUCCESS_TEST([] { return 1; },
+                        [](Future<int>&& fut) {
+                            ASSERT_EQ(std::move(fut)
+                                          .onError<ErrorCodes::InternalError>([](Status) {
+                                              FAIL("onError<code>() callback was called");
+                                              return 0;
+                                          })
+                                          .then([](int i) { return i + 2; })
+                                          .get(),
+                                      3);
+                        });
+}
+
+TEST(Future, Fail_onErrorCodeMatch) {
+    FUTURE_FAIL_TEST<int>([](Future<int>&& fut) {
+        auto res =
+            std::move(fut)
+                .onError([](Status s) {
+                    ASSERT_EQ(s, failStatus);
+                    return StatusWith<int>(ErrorCodes::InternalError, "");
+                })
+                .onError<ErrorCodes::InternalError>([](Status&&) { return StatusWith<int>(3); })
+                .getNoThrow();
+        ASSERT_EQ(res, 3);
+    });
+}
+
+TEST(Future, Fail_onErrorCodeMatchFuture) {
+    FUTURE_FAIL_TEST<int>([](Future<int>&& fut) {
+        auto res = std::move(fut)
+                       .onError([](Status s) {
+                           ASSERT_EQ(s, failStatus);
+                           return StatusWith<int>(ErrorCodes::InternalError, "");
+                       })
+                       .onError<ErrorCodes::InternalError>([](Status&&) { return Future<int>(3); })
+                       .getNoThrow();
+        ASSERT_EQ(res, 3);
+    });
+}
+
+TEST(Future, Fail_onErrorCodeMismatch) {
+    FUTURE_FAIL_TEST<int>([](Future<int>&& fut) {
+        ASSERT_EQ(std::move(fut)
+                      .onError<ErrorCodes::InternalError>([](Status s) -> int {
+                          FAIL("Why was this called?") << s;
+                          MONGO_UNREACHABLE;
+                      })
+                      .onError([](Status s) {
+                          ASSERT_EQ(s, failStatus);
+                          return 3;
+                      })
+                      .getNoThrow(),
+                  3);
+    });
+}
+
+
 TEST(Future, Success_tap) {
     FUTURE_SUCCESS_TEST([] { return 1; },
                         [](Future<int>&& fut) {
@@ -904,6 +962,7 @@ TEST(Future_Void, Fail_onErrorSimple) {
                   3);
     });
 }
+
 TEST(Future_Void, Fail_onErrorError_throw) {
     FUTURE_FAIL_TEST<void>([](Future<void>&& fut) {
         auto fut2 = std::move(fut).onError([](Status s) {
@@ -966,6 +1025,66 @@ TEST(Future_Void, Fail_onErrorFutureAsync) {
     });
 }
 
+TEST(Future_Void, Success_onErrorCode) {
+    FUTURE_SUCCESS_TEST([] {},
+                        [](Future<void>&& fut) {
+                            ASSERT_EQ(std::move(fut)
+                                          .onError<ErrorCodes::InternalError>([](Status) {
+                                              FAIL("onError<code>() callback was called");
+                                          })
+                                          .then([] { return 3; })
+                                          .get(),
+                                      3);
+                        });
+}
+
+TEST(Future_Void, Fail_onErrorCodeMatch) {
+    FUTURE_FAIL_TEST<void>([](Future<void>&& fut) {
+        bool called = false;
+        auto res = std::move(fut)
+                       .onError([](Status s) {
+                           ASSERT_EQ(s, failStatus);
+                           return Status(ErrorCodes::InternalError, "");
+                       })
+                       .onError<ErrorCodes::InternalError>([&](Status&&) { called = true; })
+                       .then([] { return 3; })
+                       .getNoThrow();
+        ASSERT_EQ(res, 3);
+        ASSERT(called);
+    });
+}
+
+TEST(Future_Void, Fail_onErrorCodeMatchFuture) {
+    FUTURE_FAIL_TEST<void>([](Future<void>&& fut) {
+        bool called = false;
+        auto res = std::move(fut)
+                       .onError([](Status s) {
+                           ASSERT_EQ(s, failStatus);
+                           return Status(ErrorCodes::InternalError, "");
+                       })
+                       .onError<ErrorCodes::InternalError>([&](Status&&) {
+                           called = true;
+                           return Future<void>::makeReady();
+                       })
+                       .then([] { return 3; })
+                       .getNoThrow();
+        ASSERT_EQ(res, 3);
+        ASSERT(called);
+    });
+}
+
+TEST(Future_Void, Fail_onErrorCodeMismatch) {
+    FUTURE_FAIL_TEST<void>([](Future<void>&& fut) {
+        ASSERT_EQ(std::move(fut)
+                      .onError<ErrorCodes::InternalError>(
+                          [](Status s) { FAIL("Why was this called?") << s; })
+                      .onError([](Status s) { ASSERT_EQ(s, failStatus); })
+                      .then([] { return 3; })
+                      .getNoThrow(),
+                  3);
+    });
+}
+
 TEST(Future_Void, Success_tap) {
     FUTURE_SUCCESS_TEST(
         [] {},