SERVER-52972 Retry drop role transactions in UMC command

1 files changed, 126 insertions, 65 deletions

diff --git a/src/mongo/db/commands/user_management_commands.cpp b/src/mongo/db/commands/user_management_commands.cpp
index bdbf0d89c2c..5462ae4ccb0 100644
--- a/src/mongo/db/commands/user_management_commands.cpp
+++ b/src/mongo/db/commands/user_management_commands.cpp
@@ -1756,6 +1756,63 @@ void CmdUMCTyped<RevokeRolesFromRoleCommand, void>::Invocation::typedRun(Operati
     uassertStatusOK(status);
 }
 
+/**
+ * Attempt to complete a transaction, retrying up to two times (3 total attempts).
+ * Emit an audit entry prior to the first commit attempt,
+ * but do not repeat the audit entry for retries.
+ */
+using TxnOpsCallback = std::function<Status(UMCTransaction&)>;
+using TxnAuditCallback = std::function<void()>;
+Status retryTransactionOps(OperationContext* opCtx,
+                           StringData forCommand,
+                           TxnOpsCallback ops,
+                           TxnAuditCallback audit) {
+    // In practice this status never makes it to a return
+    // since its populated with the return from ops(),
+    // but guard against bit-rot by pre-populating a generic failure.
+    Status status(ErrorCodes::OperationFailed, "Operation was never attempted");
+    constexpr int kMaxAttempts = 3;
+
+    for (int tries = kMaxAttempts; tries > 0; --tries) {
+        if (tries < kMaxAttempts) {
+            // Emit log on all but the first attempt.
+            LOGV2_DEBUG(5297200,
+                        4,
+                        "Retrying user management command transaction",
+                        "command"_attr = forCommand,
+                        "reason"_attr = status);
+        }
+
+        UMCTransaction txn(opCtx, forCommand);
+        status = ops(txn);
+        if (!status.isOK()) {
+            // A failure in the setup ops is just a failure, abort without retry.
+            return status;
+        }
+
+        if (tries == kMaxAttempts) {
+            // Only emit audit on first attempt.
+            audit();
+        }
+
+        status = txn.commit();
+        if (status.isOK()) {
+            // Success, see ya later!
+            return status;
+        }
+
+        // Try to responsibly abort, but accept not being able to.
+        txn.abort().ignore();
+
+        if ((status != ErrorCodes::LockTimeout) && (status != ErrorCodes::SnapshotUnavailable)) {
+            // Something moved underneath us. Try again.
+            return status;
+        }
+    }
+
+    return status;
+}
+
 CmdUMCTyped<DropRoleCommand, void> cmdDropRole;
 template <>
 void CmdUMCTyped<DropRoleCommand, void>::Invocation::typedRun(OperationContext* opCtx) {
@@ -1783,40 +1840,40 @@ void CmdUMCTyped<DropRoleCommand, void>::Invocation::typedRun(OperationContext*
         }
     });
 
-    UMCTransaction txn(opCtx, DropRoleCommand::kCommandName);
-
-    // Remove this role from all users
-    auto swCount = txn.update(AuthorizationManager::usersCollectionNamespace,
-                              BSON("roles" << BSON("$elemMatch" << roleName.toBSON())),
-                              BSON("$pull" << BSON("roles" << roleName.toBSON())));
-    if (!swCount.isOK()) {
-        uassertStatusOK(useDefaultCode(swCount.getStatus(), ErrorCodes::UserModificationFailed)
-                            .withContext(str::stream()
-                                         << "Failed to remove role " << roleName.getFullName()
-                                         << " from all users"));
-    }
+    const auto dropRoleOps = [&](UMCTransaction& txn) -> Status {
+        // Remove this role from all users
+        auto swCount = txn.update(AuthorizationManager::usersCollectionNamespace,
+                                  BSON("roles" << BSON("$elemMatch" << roleName.toBSON())),
+                                  BSON("$pull" << BSON("roles" << roleName.toBSON())));
+        if (!swCount.isOK()) {
+            return useDefaultCode(swCount.getStatus(), ErrorCodes::UserModificationFailed)
+                .withContext(str::stream() << "Failed to remove role " << roleName.getFullName()
+                                           << " from all users");
+        }
 
-    // Remove this role from all other roles
-    swCount = txn.update(AuthorizationManager::rolesCollectionNamespace,
-                         BSON("roles" << BSON("$elemMatch" << roleName.toBSON())),
-                         BSON("$pull" << BSON("roles" << roleName.toBSON())));
-    if (!swCount.isOK()) {
-        uassertStatusOK(useDefaultCode(swCount.getStatus(), ErrorCodes::RoleModificationFailed)
-                            .withContext(str::stream()
-                                         << "Failed to remove role " << roleName.getFullName()
-                                         << " from all users"));
-    }
+        // Remove this role from all other roles
+        swCount = txn.update(AuthorizationManager::rolesCollectionNamespace,
+                             BSON("roles" << BSON("$elemMatch" << roleName.toBSON())),
+                             BSON("$pull" << BSON("roles" << roleName.toBSON())));
+        if (!swCount.isOK()) {
+            return useDefaultCode(swCount.getStatus(), ErrorCodes::RoleModificationFailed)
+                .withContext(str::stream() << "Failed to remove role " << roleName.getFullName()
+                                           << " from all users");
+        }
 
-    // Finally, remove the actual role document
-    swCount = txn.remove(AuthorizationManager::rolesCollectionNamespace, roleName.toBSON());
-    if (!swCount.isOK()) {
-        uassertStatusOK(swCount.getStatus().withContext(str::stream() << "Failed to remove role "
-                                                                      << roleName.getFullName()));
-    }
+        // Finally, remove the actual role document
+        swCount = txn.remove(AuthorizationManager::rolesCollectionNamespace, roleName.toBSON());
+        if (!swCount.isOK()) {
+            return swCount.getStatus().withContext(str::stream() << "Failed to remove role "
+                                                                 << roleName.getFullName());
+        }
 
-    audit::logDropRole(client, roleName);
+        return Status::OK();
+    };
 
-    auto status = txn.commit();
+    auto status = retryTransactionOps(opCtx, DropRoleCommand::kCommandName, dropRoleOps, [&] {
+        audit::logDropRole(client, roleName);
+    });
     if (!status.isOK()) {
         uassertStatusOK(status.withContext("Failed applying dropRole transaction"));
     }
@@ -1845,50 +1902,54 @@ CmdUMCTyped<DropAllRolesFromDatabaseCommand, DropAllRolesFromDatabaseReply>::Inv
         }
     });
 
-    UMCTransaction txn(opCtx, DropAllRolesFromDatabaseCommand::kCommandName);
-    auto roleMatch = BSON(AuthorizationManager::ROLE_DB_FIELD_NAME << dbname);
-    auto rolesMatch = BSON("roles" << roleMatch);
-
-    // Remove these roles from all users
-    auto swCount = txn.update(
-        AuthorizationManager::usersCollectionNamespace, rolesMatch, BSON("$pull" << rolesMatch));
-    if (!swCount.isOK()) {
-        uassertStatusOK(useDefaultCode(swCount.getStatus(), ErrorCodes::UserModificationFailed)
-                            .withContext(str::stream() << "Failed to remove roles from \"" << dbname
-                                                       << "\" db from all users"));
-    }
+    DropAllRolesFromDatabaseReply reply;
+    const auto dropRoleOps = [&](UMCTransaction& txn) -> Status {
+        auto roleMatch = BSON(AuthorizationManager::ROLE_DB_FIELD_NAME << dbname);
+        auto rolesMatch = BSON("roles" << roleMatch);
+
+        // Remove these roles from all users
+        auto swCount = txn.update(AuthorizationManager::usersCollectionNamespace,
+                                  rolesMatch,
+                                  BSON("$pull" << rolesMatch));
+        if (!swCount.isOK()) {
+            return useDefaultCode(swCount.getStatus(), ErrorCodes::UserModificationFailed)
+                .withContext(str::stream() << "Failed to remove roles from \"" << dbname
+                                           << "\" db from all users");
+        }
 
-    // Remove these roles from all other roles
-    swCount = txn.update(AuthorizationManager::rolesCollectionNamespace,
-                         BSON("roles.db" << dbname),
-                         BSON("$pull" << rolesMatch));
-    if (!swCount.isOK()) {
-        uassertStatusOK(useDefaultCode(swCount.getStatus(), ErrorCodes::RoleModificationFailed)
-                            .withContext(str::stream() << "Failed to remove roles from \"" << dbname
-                                                       << "\" db from all roles"));
-    }
+        // Remove these roles from all other roles
+        swCount = txn.update(AuthorizationManager::rolesCollectionNamespace,
+                             BSON("roles.db" << dbname),
+                             BSON("$pull" << rolesMatch));
+        if (!swCount.isOK()) {
+            return useDefaultCode(swCount.getStatus(), ErrorCodes::RoleModificationFailed)
+                .withContext(str::stream() << "Failed to remove roles from \"" << dbname
+                                           << "\" db from all roles");
+        }
 
-    // Finally, remove the actual role documents
-    swCount = txn.remove(AuthorizationManager::rolesCollectionNamespace, roleMatch);
-    if (!swCount.isOK()) {
-        uassertStatusOK(swCount.getStatus().withContext(
-            str::stream() << "Removed roles from \"" << dbname
-                          << "\" db "
-                             " from all users and roles but failed to actually delete"
-                             " those roles themselves"));
-    }
+        // Finally, remove the actual role documents
+        swCount = txn.remove(AuthorizationManager::rolesCollectionNamespace, roleMatch);
+        if (!swCount.isOK()) {
+            return swCount.getStatus().withContext(
+                str::stream() << "Removed roles from \"" << dbname
+                              << "\" db "
+                                 " from all users and roles but failed to actually delete"
+                                 " those roles themselves");
+        }
 
-    audit::logDropAllRolesFromDatabase(Client::getCurrent(), dbname);
+        reply.setCount(swCount.getValue());
+        return Status::OK();
+    };
 
-    auto status = txn.commit();
+    auto status =
+        retryTransactionOps(opCtx, DropAllRolesFromDatabaseCommand::kCommandName, dropRoleOps, [&] {
+            audit::logDropAllRolesFromDatabase(Client::getCurrent(), dbname);
+        });
     if (!status.isOK()) {
         uassertStatusOK(
             status.withContext("Failed applying dropAllRolesFromDatabase command transaction"));
     }
 
-
-    DropAllRolesFromDatabaseReply reply;
-    reply.setCount(swCount.getValue());
     return reply;
 }