1 files changed, 19 insertions, 15 deletions

diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp
index ebed6807c7e..96eb011018d 100644
--- a/src/mongo/db/dbhelpers.cpp
+++ b/src/mongo/db/dbhelpers.cpp
@@ -62,8 +62,8 @@
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/storage/data_protector.h"
+#include "mongo/db/storage/encryption_hooks.h"
 #include "mongo/db/storage/storage_options.h"
-#include "mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h"
 #include "mongo/db/write_concern.h"
 #include "mongo/db/write_concern_options.h"
 #include "mongo/s/shard_key_pattern.h"
@@ -270,6 +270,7 @@ long long Helpers::removeRange(OperationContext* txn,
                                const KeyRange& range,
                                BoundInclusion boundInclusion,
                                const WriteConcernOptions& writeConcern,
+                               Milliseconds& replWaitDuration,
                                RemoveSaver* callback,
                                bool fromMigrate,
                                bool onlyRemoveOrphanedDocs) {
@@ -283,6 +284,7 @@ long long Helpers::removeRange(OperationContext* txn,
     BSONObj max;
 
     {
+        ScopedTransaction scopedXact(txn, MODE_IS);
         AutoGetCollectionForRead ctx(txn, ns);
         Collection* collection = ctx.getCollection();
         if (!collection) {
@@ -326,11 +328,12 @@ long long Helpers::removeRange(OperationContext* txn,
 
     long long numDeleted = 0;
 
-    Milliseconds millisWaitingForReplication{0};
+    replWaitDuration = Milliseconds::zero();
 
     while (1) {
         // Scoping for write lock.
         {
+            ScopedTransaction scopedXact(txn, MODE_IX);
             AutoGetCollection ctx(txn, NamespaceString(ns), MODE_IX, MODE_IX);
             Collection* collection = ctx.getCollection();
             if (!collection)
@@ -438,20 +441,21 @@ long long Helpers::removeRange(OperationContext* txn,
                     txn,
                     repl::ReplClientInfo::forClient(txn->getClient()).getLastOp(),
                     writeConcern);
-            if (replStatus.status.code() == ErrorCodes::ExceededTimeLimit) {
+            if (replStatus.status.code() == ErrorCodes::ExceededTimeLimit ||
+                replStatus.status.code() == ErrorCodes::WriteConcernFailed) {
                 warning(LogComponent::kSharding) << "replication to secondaries for removeRange at "
                                                     "least 60 seconds behind";
             } else {
                 uassertStatusOK(replStatus.status);
             }
-            millisWaitingForReplication += replStatus.duration;
+            replWaitDuration += replStatus.duration;
         }
     }
 
     if (writeConcern.shouldWaitForOtherNodes())
         log(LogComponent::kSharding)
             << "Helpers::removeRangeUnlocked time spent waiting for replication: "
-            << durationCount<Milliseconds>(millisWaitingForReplication) << "ms" << endl;
+            << durationCount<Milliseconds>(replWaitDuration) << "ms" << endl;
 
     MONGO_LOG_COMPONENT(1, LogComponent::kSharding) << "end removal of " << min << " to " << max
                                                     << " in " << ns << " (took "
@@ -485,19 +489,19 @@ Helpers::RemoveSaver::RemoveSaver(const string& a, const string& b, const string
     ss << why << "." << terseCurrentTime(false) << "." << NUM++ << ".bson";
     _file /= ss.str();
 
-    auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
-    if (hooks->enabled()) {
-        _protector = hooks->getDataProtector();
-        _file += hooks->getProtectedPathSuffix();
+    auto encryptionHooks = EncryptionHooks::get(getGlobalServiceContext());
+    if (encryptionHooks->enabled()) {
+        _protector = encryptionHooks->getDataProtector();
+        _file += encryptionHooks->getProtectedPathSuffix();
     }
 }
 
 Helpers::RemoveSaver::~RemoveSaver() {
     if (_protector && _out) {
-        auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
-        invariant(hooks->enabled());
+        auto encryptionHooks = EncryptionHooks::get(getGlobalServiceContext());
+        invariant(encryptionHooks->enabled());
 
-        size_t protectedSizeMax = hooks->additionalBytesForProtectedBuffer();
+        size_t protectedSizeMax = encryptionHooks->additionalBytesForProtectedBuffer();
         std::unique_ptr<uint8_t[]> protectedBuffer(new uint8_t[protectedSizeMax]);
 
         size_t resultLen;
@@ -560,10 +564,10 @@ Status Helpers::RemoveSaver::goingToDelete(const BSONObj& o) {
 
     std::unique_ptr<uint8_t[]> protectedBuffer;
     if (_protector) {
-        auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
-        invariant(hooks->enabled());
+        auto encryptionHooks = EncryptionHooks::get(getGlobalServiceContext());
+        invariant(encryptionHooks->enabled());
 
-        size_t protectedSizeMax = dataSize + hooks->additionalBytesForProtectedBuffer();
+        size_t protectedSizeMax = dataSize + encryptionHooks->additionalBytesForProtectedBuffer();
         protectedBuffer.reset(new uint8_t[protectedSizeMax]);
 
         size_t resultLen;