SERVER-38368 Allow handling ErrorCategory types for Futures error handling

2 files changed, 67 insertions, 0 deletions

diff --git a/src/mongo/util/future.h b/src/mongo/util/future.h
index 3033bec5687..c32731c1986 100644
--- a/src/mongo/util/future.h
+++ b/src/mongo/util/future.h
@@ -1097,6 +1097,29 @@ public:
     }
 
     /**
+     * Similar to the first two onErrors, but only calls the callback if the category matches
+     * the template parameter. Otherwise lets the error propagate unchanged.
+     */
+    template <ErrorCategory category, typename Func>
+        Future<T> onErrorCategory(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>::onErrorCategory must return T, StatusWith<T>, or Future<T>");
+
+        if (_immediate || (isReady() && _shared->status.isOK()))
+            return std::move(*this);
+
+        return std::move(*this).onError([func =
+                                             std::forward<Func>(func)](Status && status) mutable {
+            if (!ErrorCodes::isA<category>(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.
diff --git a/src/mongo/util/future_test_future_int.cpp b/src/mongo/util/future_test_future_int.cpp
index 48f6335421c..6e4395952a7 100644
--- a/src/mongo/util/future_test_future_int.cpp
+++ b/src/mongo/util/future_test_future_int.cpp
@@ -529,6 +529,50 @@ TEST(Future, Fail_onErrorCodeMismatch) {
     });
 }
 
+TEST(Future, Success_onErrorCategory) {
+    FUTURE_SUCCESS_TEST([] { return 1; },
+                        [](Future<int>&& fut) {
+                            ASSERT_EQ(
+                                std::move(fut)
+                                    .onErrorCategory<ErrorCategory::NetworkError>([](Status) {
+                                        FAIL("onErrorCategory<category>() callback was called");
+                                        return 0;
+                                    })
+                                    .then([](int i) { return i + 2; })
+                                    .get(),
+                                3);
+                        });
+}
+
+TEST(Future, Fail_onErrorCategoryMatch) {
+    FUTURE_FAIL_TEST<int>([](Future<int>&& fut) {
+        auto res = std::move(fut)
+                       .onError([](Status s) {
+                           ASSERT_EQ(s, failStatus());
+                           return StatusWith<int>(ErrorCodes::HostUnreachable, "");
+                       })
+                       .onErrorCategory<ErrorCategory::NetworkError>(
+                           [](Status&&) { return StatusWith<int>(3); })
+                       .getNoThrow();
+        ASSERT_EQ(res, 3);
+    });
+}
+
+TEST(Future, Fail_onErrorCategoryMismatch) {
+    FUTURE_FAIL_TEST<int>([](Future<int>&& fut) {
+        ASSERT_EQ(std::move(fut)
+                      .onErrorCategory<ErrorCategory::NetworkError>([](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; },