diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/mongo/util/future.h | 23 | ||||
-rw-r--r-- | src/mongo/util/future_test_future_int.cpp | 44 |
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; }, |