SERVER-31622 Support catching DBException by code or category

8 files changed, 313 insertions, 38 deletions

diff --git a/src/mongo/base/error_codes.tpl.cpp b/src/mongo/base/error_codes.tpl.cpp
index c1188c3bd15..3c8d777dc45 100644
--- a/src/mongo/base/error_codes.tpl.cpp
+++ b/src/mongo/base/error_codes.tpl.cpp
@@ -74,4 +74,28 @@ bool ErrorCodes::is${cat.name}(Error err) {
 }
 //#end for
 
+void error_details::throwExceptionForStatus(const Status& status) {
+    /**
+     * This type is used for all exceptions that don't have a more specific type. It is defined
+     * locally in this function to prevent anyone from catching it specifically separately from
+     * AssertionException.
+     */
+    class NonspecificAssertionException final : public AssertionException {
+    public:
+        using AssertionException::AssertionException;
+
+    private:
+        void defineOnlyInFinalSubclassToPreventSlicing() final {}
+    };
+
+    switch (status.code()) {
+        //#for $ec in $codes
+        case ErrorCodes::$ec.name:
+            throw ExceptionFor<ErrorCodes::$ec.name>(status);
+        //#end for
+        default:
+            throw NonspecificAssertionException(status);
+    }
+}
+
 }  // namespace mongo
diff --git a/src/mongo/base/error_codes.tpl.h b/src/mongo/base/error_codes.tpl.h
index a5ebd55042b..a2b2cbb7f06 100644
--- a/src/mongo/base/error_codes.tpl.h
+++ b/src/mongo/base/error_codes.tpl.h
@@ -33,9 +33,18 @@
 #include <string>
 
 #include "mongo/base/string_data.h"
+#include "mongo/platform/compiler.h"
 
 namespace mongo {
 
+class Status;
+
+enum class ErrorCategory {
+    //#for $cat in $categories
+    ${cat.name},
+    //#end for
+};
+
 /**
  * This is a generated class containing a table of error codes and their corresponding error
  * strings. The class is derived from the definitions in src/mongo/base/error_codes.err file and the
@@ -72,11 +81,69 @@ public:
         return static_cast<Error>(code);
     }
 
+    /**
+     * Generic predicate to test if a given error code is in a category.
+     *
+     * This version is intended to simplify forwarding by Status and DBException. Non-generic
+     * callers should just use the specific isCategoryName() methods instead.
+     */
+    template <ErrorCategory category>
+    static bool isA(Error code);
+
     //#for $cat in $categories
-    static bool is${cat.name}(Error err);
+    static bool is${cat.name}(Error code);
     //#end for
 };
 
 std::ostream& operator<<(std::ostream& stream, ErrorCodes::Error code);
 
+//#for $cat in $categories
+template <>
+inline bool ErrorCodes::isA<ErrorCategory::$cat.name>(Error code) {
+    return is${cat.name}(code);
+}
+//#end for
+
+/**
+ * This namespace contains implementation details for our error handling code and should not be used
+ * directly in general code.
+ */
+namespace error_details {
+
+template <int32_t code>
+constexpr bool isNamedCode = false;
+//#for $ec in $codes
+template <>
+constexpr bool isNamedCode<ErrorCodes::$ec.name> = true;
+//#end for
+
+MONGO_COMPILER_NORETURN void throwExceptionForStatus(const Status& status);
+
+template <ErrorCategory... categories>
+struct CategoryList;
+
+template <ErrorCodes::Error code>
+struct ErrorCategoriesForImpl {
+    using type = CategoryList<>;
+};
+
+//#for $ec in $codes:
+//#if $ec.categories
+template <>
+struct ErrorCategoriesForImpl<ErrorCodes::$ec.name> {
+    using type = CategoryList<
+        //#for $i, $cat in enumerate($ec.categories)
+        //#set $comma = '' if i == len($ec.categories) - 1 else ', '
+        ErrorCategory::$cat$comma
+        //#end for
+        >;
+};
+//#end if
+//#end for
+
+template <ErrorCodes::Error code>
+using ErrorCategoriesFor = typename ErrorCategoriesForImpl<code>::type;
+
+}  // namespace error_details
+
 }  // namespace mongo
diff --git a/src/mongo/base/generate_error_codes.py b/src/mongo/base/generate_error_codes.py
index d3da9970692..420ee964ffe 100644
--- a/src/mongo/base/generate_error_codes.py
+++ b/src/mongo/base/generate_error_codes.py
@@ -56,13 +56,16 @@ def render_template(template_path, **kw):
             useCache=False)
     return str(template(**kw))
 
-ErrorCode = namedtuple('ErrorCode', ['name', 'code'])
-def makeErrorCode(name, code):
-    return ErrorCode(name, code)
+class ErrorCode:
+    def __init__(self, name, code):
+        self.name = name
+        self.code = code
+        self.categories = []
 
-ErrorClass = namedtuple('ErrorClass', ['name', 'codes'])
-def makeErrorClass(name, codes):
-    return ErrorClass(name, codes)
+class ErrorClass:
+    def __init__(self, name, codes):
+        self.name = name
+        self.codes = codes
 
 def main(argv):
     # Parse and validate argv.
@@ -107,9 +110,10 @@ def parse_error_definitions_from_file(errors_filename):
     error_codes = []
     error_classes = []
     eval(errors_code,
-            dict(error_code=lambda *args, **kw: error_codes.append(makeErrorCode(*args, **kw)),
-                 error_class=lambda *args: error_classes.append(makeErrorClass(*args))))
+            dict(error_code=lambda *args, **kw: error_codes.append(ErrorCode(*args, **kw)),
+                 error_class=lambda *args: error_classes.append(ErrorClass(*args))))
     error_codes.sort(key=lambda x: x.code)
+
     return error_codes, error_classes
 
 def check_for_conflicts(error_codes, error_classes):
@@ -157,13 +161,16 @@ def has_duplicate_error_classes(error_classes):
     return failed
 
 def has_missing_error_codes(error_codes, error_classes):
-    code_names = set(ec[0] for ec in error_codes)
+    code_names = dict((ec.name, ec) for ec in error_codes)
     failed = False
-    for class_name, class_code_names in error_classes:
-        for name in class_code_names:
-            if name not in code_names:
-                sys.stdout.write('Undeclared error code %s in class %s\n' % (name, class_name))
+    for category in error_classes:
+        for name in category.codes:
+            try:
+                code_names[name].categories.append(category.name)
+            except KeyError:
+                sys.stdout.write('Undeclared error code %s in class %s\n' % (name, category.name))
                 failed = True
+
     return failed
 
 if __name__ == '__main__':
diff --git a/src/mongo/base/status.h b/src/mongo/base/status.h
index 3a4462d057d..7c9a3c0ed3c 100644
--- a/src/mongo/base/status.h
+++ b/src/mongo/base/status.h
@@ -131,6 +131,14 @@ public:
     std::string toString() const;
 
     /**
+     * Returns true if this Status's code is a member of the given category.
+     */
+    template <ErrorCategory category>
+    bool isA() const {
+        return ErrorCodes::isA<category>(code());
+    }
+
+    /**
      * Call this method to indicate that it is your intention to ignore a returned status. Ignoring
      * is only possible if the value being ignored is an xvalue -- it is not appropriate to create a
      * status variable and then ignore it.
diff --git a/src/mongo/base/status_test.cpp b/src/mongo/base/status_test.cpp
index c473c24b3e0..57b2ac8a2d1 100644
--- a/src/mongo/base/status_test.cpp
+++ b/src/mongo/base/status_test.cpp
@@ -36,17 +36,21 @@
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/assert_util.h"
 
+namespace mongo {
 namespace {
 
-using mongo::ErrorCodes;
-using mongo::Status;
-
 TEST(Basic, Accessors) {
     Status status(ErrorCodes::MaxError, "error");
     ASSERT_EQUALS(status.code(), ErrorCodes::MaxError);
     ASSERT_EQUALS(status.reason(), "error");
 }
 
+TEST(Basic, IsA) {
+    ASSERT(!Status(ErrorCodes::BadValue, "").isA<ErrorCategory::Interruption>());
+    ASSERT(Status(ErrorCodes::Interrupted, "").isA<ErrorCategory::Interruption>());
+    ASSERT(!Status(ErrorCodes::Interrupted, "").isA<ErrorCategory::ShutdownError>());
+}
+
 TEST(Basic, OKIsAValidStatus) {
     Status status = Status::OK();
     ASSERT_EQUALS(status.code(), ErrorCodes::OK);
@@ -249,4 +253,5 @@ TEST(Transformers, ExceptionToStatus) {
     ASSERT_TRUE(fromBoostExcept.reason().find("boost::exception") != std::string::npos);
 }
 
-}  // unnamed namespace
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/util/assert_util.cpp b/src/mongo/util/assert_util.cpp
index bd1bc9f986a..29a0019dc35 100644
--- a/src/mongo/util/assert_util.cpp
+++ b/src/mongo/util/assert_util.cpp
@@ -55,22 +55,6 @@ using namespace std;
 
 namespace mongo {
 
-namespace {
-
-/**
- * This type is used for all exceptions that don't have a more specific type. It is defined locally
- * in this file to prevent anyone from catching it specifically separately from AssertionException.
- */
-class NonspecificAssertionException final : public AssertionException {
-public:
-    using AssertionException::AssertionException;
-
-private:
-    void defineOnlyInFinalSubclassToPreventSlicing() final {}
-};
-
-}  // namespace
-
 AssertionCount assertionCount;
 
 AssertionCount::AssertionCount() : regular(0), warning(0), msg(0), user(0), rollovers(0) {}
@@ -136,7 +120,7 @@ NOINLINE_DECL void verifyFailed(const char* expr, const char* file, unsigned lin
     severe() << "\n\n***aborting after verify() failure as this is a debug/test build\n\n" << endl;
     std::abort();
 #endif
-    throw NonspecificAssertionException(ErrorCodes::UnknownError, temp.str());
+    error_details::throwExceptionForStatus(Status(ErrorCodes::UnknownError, temp.str()));
 }
 
 NOINLINE_DECL void invariantFailed(const char* expr, const char* file, unsigned line) noexcept {
@@ -202,7 +186,7 @@ NOINLINE_DECL void uassertedWithLocation(int msgid,
     assertionCount.condrollover(++assertionCount.user);
     LOG(1) << "User Assertion: " << msgid << ":" << redact(msg) << ' ' << file << ' ' << dec << line
            << endl;
-    throw NonspecificAssertionException(msgid, msg);
+    error_details::throwExceptionForStatus(Status(ErrorCodes::Error(msgid), msg));
 }
 
 NOINLINE_DECL void msgassertedWithLocation(int msgid,
@@ -212,7 +196,7 @@ NOINLINE_DECL void msgassertedWithLocation(int msgid,
     assertionCount.condrollover(++assertionCount.msg);
     error() << "Assertion: " << msgid << ":" << redact(msg) << ' ' << file << ' ' << dec << line
             << endl;
-    throw NonspecificAssertionException(msgid, msg);
+    error_details::throwExceptionForStatus(Status(ErrorCodes::Error(msgid), msg));
 }
 
 std::string causedBy(StringData e) {
diff --git a/src/mongo/util/assert_util.h b/src/mongo/util/assert_util.h
index d562799a47e..1a3852cd328 100644
--- a/src/mongo/util/assert_util.h
+++ b/src/mongo/util/assert_util.h
@@ -96,6 +96,14 @@ public:
         return _status.codeString();
     }
 
+    /**
+     * Returns true if this DBException's code is a member of the given category.
+     */
+    template <ErrorCategory category>
+    bool isA() const {
+        return ErrorCodes::isA<category>(code());
+    }
+
     static AtomicBool traceExceptions;
 
 protected:
@@ -125,6 +133,68 @@ public:
     AssertionException(int code, StringData msg) : DBException(code, msg) {}
 };
 
+/**
+ * The base class of all DBExceptions for codes of the given ErrorCategory to allow catching by
+ * category.
+ */
+template <ErrorCategory kCategory>
+class ExceptionForCat : public virtual AssertionException {
+protected:
+    // This will only be called by subclasses, and they are required to instantiate
+    // AssertionException themselves since it is a virtual base. Therefore, the AssertionException
+    // construction here should never actually execute, but it is required to be present to allow
+    // subclasses to construct us.
+    ExceptionForCat() : AssertionException((std::abort(), Status::OK())) {
+        invariant(isA<kCategory>());
+    }
+};
+
+
+/**
+ * This namespace contains implementation details for our error handling code and should not be used
+ * directly in general code.
+ */
+namespace error_details {
+
+template <ErrorCodes::Error kCode, typename... Bases>
+class ExceptionForImpl final : public Bases... {
+public:
+    MONGO_STATIC_ASSERT(isNamedCode<kCode>);
+
+    ExceptionForImpl(const Status& status) : AssertionException(status) {
+        invariant(status.code() == kCode);
+    }
+
+private:
+    void defineOnlyInFinalSubclassToPreventSlicing() final {}
+};
+
+template <ErrorCodes::Error code, typename categories = ErrorCategoriesFor<code>>
+struct ExceptionForDispatcher;
+
+template <ErrorCodes::Error code, ErrorCategory... categories>
+struct ExceptionForDispatcher<code, CategoryList<categories...>> {
+    using type = std::conditional_t<sizeof...(categories) == 0,
+                                    ExceptionForImpl<code, AssertionException>,
+                                    ExceptionForImpl<code, ExceptionForCat<categories>...>>;
+};
+
+}  // namespace error_details
+
+
+/**
+ * Resolves to the concrete exception type for the given error code.
+ *
+ * It will be a subclass of both AssertionException, along with ExceptionForCat<> of every category
+ * that the code belongs to.
+ *
+ * TODO in C++17 we can combine this with ExceptionForCat by doing something like:
+ * template <auto codeOrCategory> using ExceptionFor = typename
+ *      error_details::ExceptionForDispatcher<decltype(codeOrCategory)>::type;
+ */
+template <ErrorCodes::Error code>
+using ExceptionFor = typename error_details::ExceptionForDispatcher<code>::type;
+
 MONGO_COMPILER_NORETURN void verifyFailed(const char* expr, const char* file, unsigned line);
 MONGO_COMPILER_NORETURN void invariantOKFailed(const char* expr,
                                                const Status& status,
diff --git a/src/mongo/util/assert_util_test.cpp b/src/mongo/util/assert_util_test.cpp
index c369922c538..315e4200d8b 100644
--- a/src/mongo/util/assert_util_test.cpp
+++ b/src/mongo/util/assert_util_test.cpp
@@ -30,12 +30,121 @@
 
 #include "mongo/platform/basic.h"
 
+#include <type_traits>
+
+#include "mongo/base/static_assert.h"
 #include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/assert_util.h"
 
+namespace mongo {
 namespace {
-using namespace mongo;
+
+#define ASSERT_CATCHES(code, Type)                                         \
+    ([] {                                                                  \
+        try {                                                              \
+            uasserted(code, "");                                           \
+        } catch (const Type&) {                                            \
+            /* Success - ignore*/                                          \
+        } catch (const std::exception& ex) {                               \
+            FAIL("Expected to be able to catch " #code " as a " #Type)     \
+                << " actual exception type: " << demangleName(typeid(ex)); \
+        }                                                                  \
+    }())
+
+#define ASSERT_NOT_CATCHES(code, Type)            \
+    ([] {                                         \
+        try {                                     \
+            uasserted(code, "");                  \
+        } catch (const Type&) {                   \
+            FAIL("Caught " #code " as a " #Type); \
+        } catch (const DBException&) {            \
+            /* Success - ignore*/                 \
+        }                                         \
+    }())
+
+// BadValue - no categories
+MONGO_STATIC_ASSERT(std::is_same<error_details::ErrorCategoriesFor<ErrorCodes::BadValue>,
+                                 error_details::CategoryList<>>());
+MONGO_STATIC_ASSERT(std::is_base_of<AssertionException, ExceptionFor<ErrorCodes::BadValue>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::NetworkError>,
+                                     ExceptionFor<ErrorCodes::BadValue>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::NotMasterError>,
+                                     ExceptionFor<ErrorCodes::BadValue>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::Interruption>,
+                                     ExceptionFor<ErrorCodes::BadValue>>());
+
+TEST(AssertUtils, UassertNamedCodeWithoutCategories) {
+    ASSERT_CATCHES(ErrorCodes::BadValue, DBException);
+    ASSERT_CATCHES(ErrorCodes::BadValue, AssertionException);
+    ASSERT_CATCHES(ErrorCodes::BadValue, ExceptionFor<ErrorCodes::BadValue>);
+    ASSERT_NOT_CATCHES(ErrorCodes::BadValue, ExceptionFor<ErrorCodes::DuplicateKey>);
+    ASSERT_NOT_CATCHES(ErrorCodes::BadValue, ExceptionForCat<ErrorCategory::NetworkError>);
+    ASSERT_NOT_CATCHES(ErrorCodes::BadValue, ExceptionForCat<ErrorCategory::NotMasterError>);
+    ASSERT_NOT_CATCHES(ErrorCodes::BadValue, ExceptionForCat<ErrorCategory::Interruption>);
+}
+
+// NotMaster - just NotMasterError
+MONGO_STATIC_ASSERT(std::is_same<error_details::ErrorCategoriesFor<ErrorCodes::NotMaster>,
+                                 error_details::CategoryList<ErrorCategory::NotMasterError>>());
+MONGO_STATIC_ASSERT(std::is_base_of<AssertionException, ExceptionFor<ErrorCodes::NotMaster>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::NetworkError>,
+                                     ExceptionFor<ErrorCodes::NotMaster>>());
+MONGO_STATIC_ASSERT(std::is_base_of<ExceptionForCat<ErrorCategory::NotMasterError>,
+                                    ExceptionFor<ErrorCodes::NotMaster>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::Interruption>,
+                                     ExceptionFor<ErrorCodes::NotMaster>>());
+
+TEST(AssertUtils, UassertNamedCodeWithOneCategory) {
+    ASSERT_CATCHES(ErrorCodes::NotMaster, DBException);
+    ASSERT_CATCHES(ErrorCodes::NotMaster, AssertionException);
+    ASSERT_CATCHES(ErrorCodes::NotMaster, ExceptionFor<ErrorCodes::NotMaster>);
+    ASSERT_NOT_CATCHES(ErrorCodes::NotMaster, ExceptionFor<ErrorCodes::DuplicateKey>);
+    ASSERT_NOT_CATCHES(ErrorCodes::NotMaster, ExceptionForCat<ErrorCategory::NetworkError>);
+    ASSERT_CATCHES(ErrorCodes::NotMaster, ExceptionForCat<ErrorCategory::NotMasterError>);
+    ASSERT_NOT_CATCHES(ErrorCodes::NotMaster, ExceptionForCat<ErrorCategory::Interruption>);
+}
+
+// InterruptedDueToReplStateChange - NotMasterError and Interruption
+MONGO_STATIC_ASSERT(
+    std::is_same<
+        error_details::ErrorCategoriesFor<ErrorCodes::InterruptedDueToReplStateChange>,
+        error_details::CategoryList<ErrorCategory::Interruption, ErrorCategory::NotMasterError>>());
+MONGO_STATIC_ASSERT(std::is_base_of<AssertionException,
+                                    ExceptionFor<ErrorCodes::InterruptedDueToReplStateChange>>());
+MONGO_STATIC_ASSERT(!std::is_base_of<ExceptionForCat<ErrorCategory::NetworkError>,
+                                     ExceptionFor<ErrorCodes::InterruptedDueToReplStateChange>>());
+MONGO_STATIC_ASSERT(std::is_base_of<ExceptionForCat<ErrorCategory::NotMasterError>,
+                                    ExceptionFor<ErrorCodes::InterruptedDueToReplStateChange>>());
+MONGO_STATIC_ASSERT(std::is_base_of<ExceptionForCat<ErrorCategory::Interruption>,
+                                    ExceptionFor<ErrorCodes::InterruptedDueToReplStateChange>>());
+
+TEST(AssertUtils, UassertNamedCodeWithTwoCategories) {
+    ASSERT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange, DBException);
+    ASSERT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange, AssertionException);
+    ASSERT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange,
+                   ExceptionFor<ErrorCodes::InterruptedDueToReplStateChange>);
+    ASSERT_NOT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange,
+                       ExceptionFor<ErrorCodes::DuplicateKey>);
+    ASSERT_NOT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange,
+                       ExceptionForCat<ErrorCategory::NetworkError>);
+    ASSERT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange,
+                   ExceptionForCat<ErrorCategory::NotMasterError>);
+    ASSERT_CATCHES(ErrorCodes::InterruptedDueToReplStateChange,
+                   ExceptionForCat<ErrorCategory::Interruption>);
+}
+
+MONGO_STATIC_ASSERT(!error_details::isNamedCode<19999>);
+// ExceptionFor<ErrorCodes::Error(19999)> invalidType;  // Must not compile.
+
+TEST(AssertUtils, UassertNumericCode) {
+    ASSERT_CATCHES(19999, DBException);
+    ASSERT_CATCHES(19999, AssertionException);
+    ASSERT_NOT_CATCHES(19999, ExceptionFor<ErrorCodes::DuplicateKey>);
+    ASSERT_NOT_CATCHES(19999, ExceptionForCat<ErrorCategory::NetworkError>);
+    ASSERT_NOT_CATCHES(19999, ExceptionForCat<ErrorCategory::NotMasterError>);
+    ASSERT_NOT_CATCHES(19999, ExceptionForCat<ErrorCategory::Interruption>);
+}
 
 // uassert and its friends
 DEATH_TEST(UassertionTerminationTest, uassert, "Terminating with uassert") {
@@ -111,3 +220,4 @@ DEATH_TEST(MassertionTerminationTest, msgasserted, "Terminating with msgasserted
 }
 
 }  // namespace
+}  // namespace mongo