SERVER-20384 Add simplified interface to perform an action on commit or rollback

3 files changed, 56 insertions, 46 deletions

diff --git a/src/mongo/db/catalog/index_catalog.cpp b/src/mongo/db/catalog/index_catalog.cpp
index 80ae8e720b2..940b2fa9810 100644
--- a/src/mongo/db/catalog/index_catalog.cpp
+++ b/src/mongo/db/catalog/index_catalog.cpp
@@ -126,35 +126,6 @@ Status IndexCatalog::init(OperationContext* txn) {
     return Status::OK();
 }
 
-namespace {
-class IndexCleanupOnRollback : public RecoveryUnit::Change {
-public:
-    /**
-     * None of these pointers are owned by this class.
-     */
-    IndexCleanupOnRollback(OperationContext* txn,
-                           Collection* collection,
-                           IndexCatalogEntryContainer* entries,
-                           const IndexDescriptor* desc)
-        : _txn(txn), _collection(collection), _entries(entries), _desc(desc) {}
-
-    virtual void commit() {}
-
-    virtual void rollback() {
-        // Need to preserve indexName as _desc no longer exists after remove().
-        const std::string indexName = _desc->indexName();
-        _entries->remove(_desc);
-        _collection->infoCache()->droppedIndex(_txn, indexName);
-    }
-
-private:
-    OperationContext* _txn;
-    Collection* _collection;
-    IndexCatalogEntryContainer* _entries;
-    const IndexDescriptor* _desc;
-};
-}  // namespace
-
 IndexCatalogEntry* IndexCatalog::_setupInMemoryStructures(OperationContext* txn,
                                                           IndexDescriptor* descriptor,
                                                           bool initFromDisk) {
@@ -179,8 +150,12 @@ IndexCatalogEntry* IndexCatalog::_setupInMemoryStructures(OperationContext* txn,
     _entries.add(entry.release());
 
     if (!initFromDisk) {
-        txn->recoveryUnit()->registerChange(
-            new IndexCleanupOnRollback(txn, _collection, &_entries, descriptor));
+        txn->recoveryUnit()->onRollback([this, txn, descriptor] {
+            // Need to preserve indexName as descriptor no longer exists after remove().
+            const std::string indexName = descriptor->indexName();
+            _entries.remove(descriptor);
+            _collection->infoCache()->droppedIndex(txn, indexName);
+        });
     }
 
     invariant(save == _entries.find(descriptor));
diff --git a/src/mongo/db/commands/dbhash.cpp b/src/mongo/db/commands/dbhash.cpp
index b428514fc29..fe42aea4124 100644
--- a/src/mongo/db/commands/dbhash.cpp
+++ b/src/mongo/db/commands/dbhash.cpp
@@ -218,24 +218,15 @@ bool DBHashCmd::run(OperationContext* txn,
     return 1;
 }
 
-class DBHashCmd::DBHashLogOpHandler : public RecoveryUnit::Change {
-public:
-    DBHashLogOpHandler(DBHashCmd* dCmd, StringData ns) : _dCmd(dCmd), _ns(ns.toString()) {}
-    void commit() {
-        stdx::lock_guard<stdx::mutex> lk(_dCmd->_cachedHashedMutex);
-        _dCmd->_cachedHashed.erase(_ns);
-    }
-    void rollback() {}
-
-private:
-    DBHashCmd* _dCmd;
-    const std::string _ns;
-};
-
 void DBHashCmd::wipeCacheForCollection(OperationContext* txn, StringData ns) {
     if (!isCachable(ns))
         return;
-    txn->recoveryUnit()->registerChange(new DBHashLogOpHandler(this, ns));
+
+    std::string nsOwned = ns.toString();
+    txn->recoveryUnit()->onCommit([this, txn, nsOwned] {
+        stdx::lock_guard<stdx::mutex> lk(_cachedHashedMutex);
+        _cachedHashed.erase(nsOwned);
+    });
 }
 
 bool DBHashCmd::isCachable(StringData ns) const {
diff --git a/src/mongo/db/storage/recovery_unit.h b/src/mongo/db/storage/recovery_unit.h
index 494ae463af9..df167f12dcf 100644
--- a/src/mongo/db/storage/recovery_unit.h
+++ b/src/mongo/db/storage/recovery_unit.h
@@ -165,6 +165,50 @@ public:
      */
     virtual void registerChange(Change* change) = 0;
 
+    /**
+     * Registers a callback to be called if the current WriteUnitOfWork rolls back.
+     *
+     * Be careful about the lifetimes of all variables captured by the callback!
+     */
+    template <typename Callback>
+    void onRollback(Callback callback) {
+        class OnRollbackChange final : public Change {
+        public:
+            OnRollbackChange(Callback&& callback) : _callback(std::move(callback)) {}
+            void rollback() final {
+                _callback();
+            }
+            void commit() final {}
+
+        private:
+            Callback _callback;
+        };
+
+        registerChange(new OnRollbackChange(std::move(callback)));
+    }
+
+    /**
+     * Registers a callback to be called if the current WriteUnitOfWork commits.
+     *
+     * Be careful about the lifetimes of all variables captured by the callback!
+     */
+    template <typename Callback>
+    void onCommit(Callback callback) {
+        class OnCommitChange final : public Change {
+        public:
+            OnCommitChange(Callback&& callback) : _callback(std::move(callback)) {}
+            void rollback() final {}
+            void commit() final {
+                _callback();
+            }
+
+        private:
+            Callback _callback;
+        };
+
+        registerChange(new OnCommitChange(std::move(callback)));
+    }
+
     //
     // The remaining methods probably belong on DurRecoveryUnit rather than on the interface.
     //