SERVER-18168: Get rid of nested units of work in the RecoveryUnit

21 files changed, 150 insertions, 214 deletions

diff --git a/src/mongo/db/catalog/index_catalog_entry.cpp b/src/mongo/db/catalog/index_catalog_entry.cpp
index 67c41f54fd0..24c0087b2c9 100644
--- a/src/mongo/db/catalog/index_catalog_entry.cpp
+++ b/src/mongo/db/catalog/index_catalog_entry.cpp
@@ -167,14 +167,13 @@ namespace mongo {
         RecoveryUnitSwap(OperationContext* txn, RecoveryUnit* newRecoveryUnit)
             : _txn(txn),
               _oldRecoveryUnit(_txn->releaseRecoveryUnit()),
-              _newRecoveryUnit(newRecoveryUnit) {
-
-            _txn->setRecoveryUnit(_newRecoveryUnit.get());
-        }
+              _oldRecoveryUnitState(_txn->setRecoveryUnit(newRecoveryUnit,
+                                                          OperationContext::kNotInUnitOfWork)),
+              _newRecoveryUnit(newRecoveryUnit) { }
 
         ~RecoveryUnitSwap() {
             _txn->releaseRecoveryUnit();
-            _txn->setRecoveryUnit(_oldRecoveryUnit);
+            _txn->setRecoveryUnit(_oldRecoveryUnit, _oldRecoveryUnitState);
         }
 
     private:
@@ -183,6 +182,7 @@ namespace mongo {
 
         // Owned, but life-time is not controlled
         RecoveryUnit* const _oldRecoveryUnit;
+        OperationContext::RecoveryUnitState const _oldRecoveryUnitState;
 
         // Owned and life-time is controlled
         const boost::scoped_ptr<RecoveryUnit> _newRecoveryUnit;
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index c4d56be099c..73264590f42 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -336,7 +336,8 @@ namespace mongo {
                     txn->recoveryUnit()->abandonSnapshot();
                     cursor->setOwnedRecoveryUnit(txn->releaseRecoveryUnit());
                     StorageEngine* engine = getGlobalServiceContext()->getGlobalStorageEngine();
-                    txn->setRecoveryUnit(engine->newRecoveryUnit());
+                    txn->setRecoveryUnit(engine->newRecoveryUnit(),
+                                         OperationContext::kNotInUnitOfWork);
                 }
             }
             else {
diff --git a/src/mongo/db/commands/pipeline_command.cpp b/src/mongo/db/commands/pipeline_command.cpp
index afc01a8c81a..571c3087254 100644
--- a/src/mongo/db/commands/pipeline_command.cpp
+++ b/src/mongo/db/commands/pipeline_command.cpp
@@ -140,7 +140,9 @@ namespace mongo {
                 txn->recoveryUnit()->abandonSnapshot();
                 cursor->setOwnedRecoveryUnit(txn->releaseRecoveryUnit());
                 StorageEngine* storageEngine = getGlobalServiceContext()->getGlobalStorageEngine();
-                txn->setRecoveryUnit(storageEngine->newRecoveryUnit());
+                invariant(txn->setRecoveryUnit(storageEngine->newRecoveryUnit(),
+                                               OperationContext::kNotInUnitOfWork)
+                          == OperationContext::kNotInUnitOfWork);
             }
 
             // Cursor needs to be in a saved state while we yield locks for getmore. State
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index 5c9de5fd52f..121b3774772 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -42,20 +42,32 @@ namespace mongo {
     class CurOp;
     class ProgressMeter;
     class StringData;
+    class WriteUnitOfWork;
+
     /**
-     * This class encompasses the state required by an operation.
-     *
-     * TODO(HK): clarify what this means.  There's one OperationContext for one user operation...
-     *           but is this true for getmore?  Also what about things like fsyncunlock / internal
-     *           users / etc.?
-     *
-     * On construction, an OperationContext associates itself with the current client, and only on
-     * destruction it deassociates itself. At any time a client can be associated with at most one
-     * OperationContext.
+     * This class encompasses the state required by an operation and lives from the time a network
+     * peration is dispatched until its execution is finished. Note that each "getmore" on a cursor
+     * is a separate operation. On construction, an OperationContext associates itself with the
+     * current client, and only on destruction it deassociates itself. At any time a client can be
+     * associated with at most one OperationContext. Each OperationContext has a RecoveryUnit
+     * associated with it, though the lifetime is not necesarily the same, see releaseRecoveryUnit 
+     * and setRecoveryUnit. The operation context also keeps track of some transaction state
+     * (RecoveryUnitState) to reduce complexity and duplication in the storage-engine specific
+     * RecoveryUnit and to allow better invariant checking.
      */
     class OperationContext : public Decorable<OperationContext> {
         MONGO_DISALLOW_COPYING(OperationContext);
+        
     public:
+        /**
+         * The RecoveryUnitState is used by WriteUnitOfWork to ensure valid state transitions.
+         */
+        enum RecoveryUnitState {
+            kNotInUnitOfWork, // not in a unit of work, no writes allowed
+            kActiveUnitOfWork, // in a unit of work that still may either commit or abort
+            kFailedUnitOfWork  // in a unit of work that has failed and must be aborted
+        };
+
         virtual ~OperationContext() { }
 
         /**
@@ -63,6 +75,7 @@ namespace mongo {
          */
         virtual RecoveryUnit* recoveryUnit() const = 0;
 
+
         /**
          * Returns the RecoveryUnit (same return value as recoveryUnit()) but the caller takes
          * ownership of the returned RecoveryUnit, and the OperationContext instance relinquishes
@@ -77,7 +90,13 @@ namespace mongo {
          */
         virtual RecoveryUnit* releaseRecoveryUnit() = 0;
 
-        virtual void setRecoveryUnit(RecoveryUnit* unit) = 0;
+        /**
+         * Associates the OperatingContext with a different RecoveryUnit for getMore or
+         * subtransactions, see RecoveryUnitSwap. The new state is passed and the old state is
+         * returned separately even though the state logically belongs to the RecoveryUnit,
+         * as it is managed by the OperationContext.
+         */
+        virtual RecoveryUnitState setRecoveryUnit(RecoveryUnit* unit, RecoveryUnitState state) = 0;
 
         /**
          * Interface for locking.  Caller DOES NOT own pointer.
@@ -158,7 +177,10 @@ namespace mongo {
     protected:
         OperationContext() { }
 
+        RecoveryUnitState _ruState = kNotInUnitOfWork;
+
     private:
+        friend class WriteUnitOfWork;
         WriteConcernOptions _writeConcern;
     };
 
@@ -167,33 +189,45 @@ namespace mongo {
     public:
         WriteUnitOfWork(OperationContext* txn)
                  : _txn(txn),
-                   _ended(false) {
-
+                   _committed(false),
+                   _toplevel(txn->_ruState == OperationContext::kNotInUnitOfWork) {
             _txn->lockState()->beginWriteUnitOfWork();
-            _txn->recoveryUnit()->beginUnitOfWork(_txn);
+            if (_toplevel) {
+                _txn->recoveryUnit()->beginUnitOfWork(_txn);
+                _txn->_ruState = OperationContext::kActiveUnitOfWork;
+            }
         }
 
         ~WriteUnitOfWork() {
-            _txn->recoveryUnit()->endUnitOfWork();
-
-            if (!_ended) {
+            if (!_committed) {
+                invariant(_txn->_ruState != OperationContext::kNotInUnitOfWork);
+                if (_toplevel) {
+                    _txn->recoveryUnit()->abortUnitOfWork();
+                    _txn->_ruState = OperationContext::kNotInUnitOfWork;
+                }
+                else {
+                    _txn->_ruState = OperationContext::kFailedUnitOfWork;
+                }
                 _txn->lockState()->endWriteUnitOfWork();
             }
         }
 
         void commit() {
-            invariant(!_ended);
-
-            _txn->recoveryUnit()->commitUnitOfWork();
+            invariant(!_committed);
+            invariant (_txn->_ruState == OperationContext::kActiveUnitOfWork);
+            if (_toplevel) {
+                _txn->recoveryUnit()->commitUnitOfWork();
+                _txn->_ruState = OperationContext::kNotInUnitOfWork;
+            }
             _txn->lockState()->endWriteUnitOfWork();
-
-            _ended = true;
+            _committed = true;
         }
 
     private:
         OperationContext* const _txn;
 
-        bool _ended;
+        bool _committed;
+        bool _toplevel;
     };
 
 
diff --git a/src/mongo/db/operation_context_impl.cpp b/src/mongo/db/operation_context_impl.cpp
index 6936e32df02..00f89395527 100644
--- a/src/mongo/db/operation_context_impl.cpp
+++ b/src/mongo/db/operation_context_impl.cpp
@@ -101,10 +101,14 @@ namespace {
         return _recovery.release();
     }
 
-    void OperationContextImpl::setRecoveryUnit(RecoveryUnit* unit) {
+    OperationContext::RecoveryUnitState OperationContextImpl::setRecoveryUnit(RecoveryUnit* unit,
+                                                                         RecoveryUnitState state) {
         _recovery.reset(unit);
+        RecoveryUnitState oldState = _ruState;
+        _ruState = state;
         if ( unit )
             unit->beingSetOnOperationContext();
+        return oldState;
     }
 
     Locker* OperationContextImpl::lockState() const {
diff --git a/src/mongo/db/operation_context_impl.h b/src/mongo/db/operation_context_impl.h
index 5f9f9f1903a..86a5f3af870 100644
--- a/src/mongo/db/operation_context_impl.h
+++ b/src/mongo/db/operation_context_impl.h
@@ -44,7 +44,8 @@ namespace mongo {
 
         virtual RecoveryUnit* releaseRecoveryUnit() override;
 
-        virtual void setRecoveryUnit(RecoveryUnit* unit) override;
+        virtual RecoveryUnitState setRecoveryUnit(RecoveryUnit* unit,
+                                                  RecoveryUnitState state) override;
 
         virtual Locker* lockState() const override;
 
diff --git a/src/mongo/db/operation_context_noop.h b/src/mongo/db/operation_context_noop.h
index f27f3d1ef96..aaf9cc012bc 100644
--- a/src/mongo/db/operation_context_noop.h
+++ b/src/mongo/db/operation_context_noop.h
@@ -71,8 +71,12 @@ namespace mongo {
             return _recoveryUnit.release();
         }
 
-        virtual void setRecoveryUnit(RecoveryUnit* unit) override {
+        virtual RecoveryUnitState setRecoveryUnit(RecoveryUnit* unit,
+                                                  RecoveryUnitState state) override {
+            RecoveryUnitState oldState = _ruState;
             _recoveryUnit.reset(unit);
+            _ruState = state;
+            return oldState;
         }
 
         virtual Locker* lockState() const override {
diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp
index cc09944bfc6..4ada32aacc5 100644
--- a/src/mongo/db/query/find.cpp
+++ b/src/mongo/db/query/find.cpp
@@ -82,7 +82,8 @@ namespace mongo {
 
         // Transfer ownership of the RecoveryUnit from the ClientCursor to the OpCtx.
         RecoveryUnit* ccRecoveryUnit = cc->releaseOwnedRecoveryUnit();
-        txn->setRecoveryUnit(ccRecoveryUnit);
+        _txnPreviousRecoveryUnitState = txn->setRecoveryUnit(ccRecoveryUnit,
+                                                             OperationContext::kNotInUnitOfWork);
     }
 
     void ScopedRecoveryUnitSwapper::dismiss() {
@@ -101,7 +102,7 @@ namespace mongo {
             _cc->setOwnedRecoveryUnit(_txn->releaseRecoveryUnit());
         }
 
-        _txn->setRecoveryUnit(_txnPreviousRecoveryUnit.release());
+        _txn->setRecoveryUnit(_txnPreviousRecoveryUnit.release(), _txnPreviousRecoveryUnitState);
     }
 
     /**
@@ -709,7 +710,9 @@ namespace mongo {
                 txn->recoveryUnit()->abandonSnapshot();
                 cc->setOwnedRecoveryUnit(txn->releaseRecoveryUnit());
                 StorageEngine* storageEngine = getGlobalServiceContext()->getGlobalStorageEngine();
-                txn->setRecoveryUnit(storageEngine->newRecoveryUnit());
+                invariant(txn->setRecoveryUnit(storageEngine->newRecoveryUnit(),
+                                               OperationContext::kNotInUnitOfWork)
+                          == OperationContext::kNotInUnitOfWork);
             }
 
             LOG(5) << "caching executor with cursorid " << ccId
diff --git a/src/mongo/db/query/find.h b/src/mongo/db/query/find.h
index f62b89f5c0a..929b964520c 100644
--- a/src/mongo/db/query/find.h
+++ b/src/mongo/db/query/find.h
@@ -58,6 +58,7 @@ namespace mongo {
         bool _dismissed;
 
         std::unique_ptr<RecoveryUnit> _txnPreviousRecoveryUnit;
+        OperationContext::RecoveryUnitState _txnPreviousRecoveryUnitState;
     };
 
     /**
diff --git a/src/mongo/db/storage/in_memory/in_memory_recovery_unit.cpp b/src/mongo/db/storage/in_memory/in_memory_recovery_unit.cpp
index 79715f592e7..7ccf4574d47 100644
--- a/src/mongo/db/storage/in_memory/in_memory_recovery_unit.cpp
+++ b/src/mongo/db/storage/in_memory/in_memory_recovery_unit.cpp
@@ -36,18 +36,8 @@
 #include "mongo/util/log.h"
 
 namespace mongo {
-    InMemoryRecoveryUnit::~InMemoryRecoveryUnit() {
-        invariant(_depth == 0);
-    }
-
-    void InMemoryRecoveryUnit::beginUnitOfWork(OperationContext* opCtx) {
-        _depth++;
-    }
 
     void InMemoryRecoveryUnit::commitUnitOfWork() {
-        if ( _depth > 1 )
-            return;
-
         try {
             for (Changes::iterator it = _changes.begin(), end = _changes.end(); it != end; ++it) {
                 (*it)->commit();
@@ -59,11 +49,7 @@ namespace mongo {
         }
     }
 
-    void InMemoryRecoveryUnit::endUnitOfWork() {
-         _depth--;
-         if (_depth > 0 )
-             return;
-
+    void InMemoryRecoveryUnit::abortUnitOfWork() {
          try {
              for (Changes::reverse_iterator it = _changes.rbegin(), end = _changes.rend();
                      it != end; ++it) {
diff --git a/src/mongo/db/storage/in_memory/in_memory_recovery_unit.h b/src/mongo/db/storage/in_memory/in_memory_recovery_unit.h
index f7edc336f44..343e9abf595 100644
--- a/src/mongo/db/storage/in_memory/in_memory_recovery_unit.h
+++ b/src/mongo/db/storage/in_memory/in_memory_recovery_unit.h
@@ -42,12 +42,9 @@ namespace mongo {
 
     class InMemoryRecoveryUnit : public RecoveryUnit {
     public:
-        InMemoryRecoveryUnit() : _depth(0) {}
-        virtual ~InMemoryRecoveryUnit();
-
-        virtual void beginUnitOfWork(OperationContext* opCtx);
-        virtual void commitUnitOfWork();
-        virtual void endUnitOfWork();
+        void beginUnitOfWork(OperationContext* opCtx) final { };
+        void commitUnitOfWork() final;
+        void abortUnitOfWork() final;
 
         virtual bool waitUntilDurable() {
             return true;
@@ -71,8 +68,7 @@ namespace mongo {
         typedef boost::shared_ptr<Change> ChangePtr;
         typedef std::vector<ChangePtr> Changes;
 
-        int _depth;
         Changes _changes;
     };
 
-}
+} // namespace mongo
diff --git a/src/mongo/db/storage/mmap_v1/dur_recovery_unit.cpp b/src/mongo/db/storage/mmap_v1/dur_recovery_unit.cpp
index cfe3d0613a0..e826277e7ff 100644
--- a/src/mongo/db/storage/mmap_v1/dur_recovery_unit.cpp
+++ b/src/mongo/db/storage/mmap_v1/dur_recovery_unit.cpp
@@ -46,64 +46,38 @@
 namespace mongo {
 
     DurRecoveryUnit::DurRecoveryUnit()
-        : _writeCount(0), _writeBytes(0), _mustRollback(false), _rollbackWritesDisabled(false) {
+        : _writeCount(0), _writeBytes(0), _inUnitOfWork(false), _rollbackWritesDisabled(false) {
     }
 
     void DurRecoveryUnit::beginUnitOfWork(OperationContext* opCtx) {
-        _startOfUncommittedChangesForLevel.push_back(Indexes(_changes.size(), _writeCount));
+        invariant(!_inUnitOfWork);
+        _inUnitOfWork = true;
     }
 
     void DurRecoveryUnit::commitUnitOfWork() {
-        invariant(inAUnitOfWork());
-        invariant(!_mustRollback);
-
-        if (!inOutermostUnitOfWork()) {
-            // If we are nested, make all changes for this level part of the containing UnitOfWork.
-            // They will be added to the global damages list once the outermost UnitOfWork commits,
-            // which it must now do.
-            if (haveUncommittedChangesAtCurrentLevel()) {
-                _startOfUncommittedChangesForLevel.back() =
-                    Indexes(_changes.size(), _writeCount);
-            }
-            return;
-        }
+        invariant(_inUnitOfWork);
 
         commitChanges();
 
         // global journal flush opportunity
         getDur().commitIfNeeded();
-    }
-
-    void DurRecoveryUnit::endUnitOfWork() {
-        invariant(inAUnitOfWork());
 
-        if (haveUncommittedChangesAtCurrentLevel()) {
-            _mustRollback = true;
-        }
+        resetChanges();
+    }
 
-        // Reset back to default if this is the last unwind of the recovery unit. That way, it can
-        // be reused for new operations.
-        if (inOutermostUnitOfWork()) {
-            rollbackChanges();
-            _rollbackWritesDisabled = false;
-            dassert(_changes.empty() && _initialWrites.empty() && _mergedWrites.empty());
-            dassert( _preimageBuffer.empty() && !_writeCount && !_writeBytes && !_mustRollback);
-        }
+    void DurRecoveryUnit::abortUnitOfWork() {
+        invariant(_inUnitOfWork);
 
-        _startOfUncommittedChangesForLevel.pop_back();
+        rollbackChanges();
+        resetChanges();
     }
 
     void DurRecoveryUnit::abandonSnapshot() {
-        invariant( !inAUnitOfWork() );
+        invariant(!_inUnitOfWork);
         // no-op since we have no transaction
     }
 
     void DurRecoveryUnit::commitChanges() {
-        invariant(!_mustRollback);
-        invariant(inOutermostUnitOfWork());
-        invariant(_startOfUncommittedChangesForLevel.front().changeIndex == 0);
-        invariant(_startOfUncommittedChangesForLevel.front().writeCount == 0);
-
         if (getDur().isDurable())
             markWritesForJournaling();
 
@@ -112,8 +86,6 @@ namespace mongo {
                     it != end; ++it) {
                 (*it)->commit();
             }
-
-            resetChanges();
         }
         catch (...) {
             std::terminate();
@@ -193,13 +165,11 @@ namespace mongo {
         _mergedWrites.clear();
         _changes.clear();
         _preimageBuffer.clear();
+        _rollbackWritesDisabled = false;
+        _inUnitOfWork = false;
     }
 
     void DurRecoveryUnit::rollbackChanges() {
-        invariant(inOutermostUnitOfWork());
-        invariant(!_startOfUncommittedChangesForLevel.back().changeIndex);
-        invariant(!_startOfUncommittedChangesForLevel.back().writeCount);
-
         // First rollback disk writes, then Changes. This matches behavior in other storage engines
         // that either rollback a transaction or don't write a writebatch.
 
@@ -232,9 +202,6 @@ namespace mongo {
                 LOG(2) << "CUSTOM ROLLBACK " << demangleName(typeid(*_changes[i]));
                 _changes[i]->rollback();
             }
-
-            resetChanges();
-            _mustRollback = false;
         }
         catch (...) {
             std::terminate();
@@ -242,7 +209,7 @@ namespace mongo {
     }
 
     bool DurRecoveryUnit::waitUntilDurable() {
-        invariant(!inAUnitOfWork());
+        invariant(!_inUnitOfWork);
         return getDur().waitUntilDurable();
     }
 
@@ -303,7 +270,7 @@ namespace mongo {
     }
 
     void* DurRecoveryUnit::writingPtr(void* addr, size_t len) {
-        invariant(inAUnitOfWork());
+        invariant(_inUnitOfWork);
 
         if (len == 0) {
             return addr; // Don't need to do anything for empty ranges.
@@ -349,12 +316,12 @@ namespace mongo {
     }
 
     void DurRecoveryUnit::setRollbackWritesDisabled() {
-        invariant(inOutermostUnitOfWork());
+        invariant(_inUnitOfWork);
         _rollbackWritesDisabled = true;
     }
 
     void DurRecoveryUnit::registerChange(Change* change) {
-        invariant(inAUnitOfWork());
+        invariant(_inUnitOfWork);
         _changes.push_back(change);
     }
 
diff --git a/src/mongo/db/storage/mmap_v1/dur_recovery_unit.h b/src/mongo/db/storage/mmap_v1/dur_recovery_unit.h
index 9c18f82edbc..d26032e8f26 100644
--- a/src/mongo/db/storage/mmap_v1/dur_recovery_unit.h
+++ b/src/mongo/db/storage/mmap_v1/dur_recovery_unit.h
@@ -46,11 +46,9 @@ namespace mongo {
     public:
         DurRecoveryUnit();
 
-        virtual ~DurRecoveryUnit() { }
-
-        virtual void beginUnitOfWork(OperationContext* opCtx);
-        virtual void commitUnitOfWork();
-        virtual void endUnitOfWork();
+        void beginUnitOfWork(OperationContext* opCtx) final;
+        void commitUnitOfWork() final;
+        void abortUnitOfWork() final;
 
         virtual bool waitUntilDurable();
 
@@ -86,20 +84,6 @@ namespace mongo {
          */
         void rollbackChanges();
 
-        bool inAUnitOfWork() const { return !_startOfUncommittedChangesForLevel.empty(); }
-
-        bool inOutermostUnitOfWork() const {
-            return _startOfUncommittedChangesForLevel.size() == 1;
-        }
-
-        /**
-         * If true, ending a unit of work will cause rollback. Ending a (possibly nested) unit of
-         * work without committing and without making any changes will not cause rollback.
-         */
-        bool haveUncommittedChangesAtCurrentLevel() const {
-            return _writeCount > _startOfUncommittedChangesForLevel.back().writeCount
-                || _changes.size() > _startOfUncommittedChangesForLevel.back().changeIndex;
-        }
 
         /**
          * Version of writingPtr that checks existing writes for overlap and only stores those
@@ -169,25 +153,8 @@ namespace mongo {
 
         std::string _preimageBuffer;
 
-        // Index of the first uncommitted Change in _changes and number of writes for each nesting
-        // level. Store the number of writes as maintained in _writeCount rather than the sum of
-        // _initialWrites and _mergedWrites, as coalescing might otherwise result in
-        // haveUncommittedChangesAtCurrent level missing merged writes when determining if rollback
-        // is necessary. Index 0 in this vector is always the outermost transaction and back() is
-        // always the innermost. The size() is the current nesting level.
-        struct Indexes {
-            Indexes(size_t changeIndex, size_t writeCount)
-                : changeIndex(changeIndex)
-                , writeCount(writeCount)
-            {}
-            size_t changeIndex;
-            size_t writeCount;
-        };
-        std::vector<Indexes> _startOfUncommittedChangesForLevel;
+        bool _inUnitOfWork;
 
-        // If true, this RU is in a "failed" state and all changes must be rolled back. Once the
-        // outermost WUOW rolls back it reverts to false.
-        bool _mustRollback;
 
         // Default is false.
         // If true, no preimages are tracked.  If rollback is subsequently attempted, the process
diff --git a/src/mongo/db/storage/mmap_v1/heap_record_store_btree.cpp b/src/mongo/db/storage/mmap_v1/heap_record_store_btree.cpp
index 9dfd6628a4a..a8eebea6ad8 100644
--- a/src/mongo/db/storage/mmap_v1/heap_record_store_btree.cpp
+++ b/src/mongo/db/storage/mmap_v1/heap_record_store_btree.cpp
@@ -108,23 +108,12 @@ namespace mongo {
 
     // ---------------------------
 
-    HeapRecordStoreBtreeRecoveryUnit::~HeapRecordStoreBtreeRecoveryUnit() {
-        invariant( _depth == 0 );
-    }
-
-    void HeapRecordStoreBtreeRecoveryUnit::beginUnitOfWork(OperationContext* opCtx) {
-        _depth++;
-    }
-
     void HeapRecordStoreBtreeRecoveryUnit::commitUnitOfWork() {
-        invariant( _depth == 1 );
         _insertions.clear();
         _mods.clear();
     }
 
-    void HeapRecordStoreBtreeRecoveryUnit::endUnitOfWork() {
-        invariant( _depth-- == 1 );
-
+    void HeapRecordStoreBtreeRecoveryUnit::abortUnitOfWork() {
         // reverse in case we write same area twice
         for ( size_t i = _mods.size(); i > 0; i-- ) {
             ModEntry& e = _mods[i-1];
diff --git a/src/mongo/db/storage/mmap_v1/heap_record_store_btree.h b/src/mongo/db/storage/mmap_v1/heap_record_store_btree.h
index 3dd7c4e2fe9..b6b0c03e0a3 100644
--- a/src/mongo/db/storage/mmap_v1/heap_record_store_btree.h
+++ b/src/mongo/db/storage/mmap_v1/heap_record_store_btree.h
@@ -177,15 +177,9 @@ namespace mongo {
      */
     class HeapRecordStoreBtreeRecoveryUnit : public RecoveryUnit {
     public:
-        HeapRecordStoreBtreeRecoveryUnit() {
-            _depth = 0;
-        }
-
-        virtual ~HeapRecordStoreBtreeRecoveryUnit();
-
-        virtual void beginUnitOfWork(OperationContext* opCtx);
-        virtual void commitUnitOfWork();
-        virtual void endUnitOfWork();
+        void beginUnitOfWork(OperationContext* opCtx) final { };
+        void commitUnitOfWork() final;
+        void abortUnitOfWork() final;
 
         virtual bool waitUntilDurable() { return true; }
 
@@ -209,7 +203,6 @@ namespace mongo {
                                   HeapRecordStoreBtree* rs, const RecordId& loc );
 
     private:
-        int _depth;
         struct InsertEntry {
             HeapRecordStoreBtree* rs;
             RecordId loc;
diff --git a/src/mongo/db/storage/mmap_v1/mmap_v1_database_catalog_entry.h b/src/mongo/db/storage/mmap_v1/mmap_v1_database_catalog_entry.h
index 53827f3e594..92dea158286 100644
--- a/src/mongo/db/storage/mmap_v1/mmap_v1_database_catalog_entry.h
+++ b/src/mongo/db/storage/mmap_v1/mmap_v1_database_catalog_entry.h
@@ -114,8 +114,9 @@ namespace mongo {
                                                 StringData ns ) const;
 
         /**
-         * Creates a CollectionCatalogEntry in the for an index rather than a collection. MMAPv1
-         * puts both indexes and collections into CCEs. A namespace named 'name' must not exist.
+         * Creates a CollectionCatalogEntry in the form of an index rather than a collection.
+         * MMAPv1 puts both indexes and collections into CCEs. A namespace named 'name' must not
+         * exist.
          */
         void createNamespaceForIndex(OperationContext* txn, StringData name);
 
diff --git a/src/mongo/db/storage/recovery_unit.h b/src/mongo/db/storage/recovery_unit.h
index 01182fec2ca..62cb0c7241d 100644
--- a/src/mongo/db/storage/recovery_unit.h
+++ b/src/mongo/db/storage/recovery_unit.h
@@ -57,26 +57,14 @@ namespace mongo {
         /**
          * These should be called through WriteUnitOfWork rather than directly.
          *
-         * begin and end mark the begining and end of a unit of work. Each call to begin must be
-         * matched with exactly one call to end. commit can be called any number of times between
-         * begin and end but must not be called outside. When end() is called, all changes since the
-         * last commit (if any) will be rolled back.
-         *
-         * If UnitsOfWork nest (ie begin is called twice before a call to end), the prior paragraph
-         * describes the behavior of the outermost UnitOfWork. Inner UnitsOfWork neither commit nor
-         * rollback on their own but rely on the outermost to do it. If an inner UnitOfWork commits
-         * any changes, it is illegal for an outer unit to rollback. If an inner UnitOfWork
-         * rollsback any changes, it is illegal for an outer UnitOfWork to do anything other than
-         * rollback.
-         *
-         * The goal is not to fully support nested transaction, instead we want to allow delaying
-         * commit on a unit if it is part of a larger atomic unit.
-         *
-         * TODO see if we can get rid of nested UnitsOfWork.
+         * A call to 'beginUnitOfWork' marks the beginning of a unit of work. Each call to
+         * 'beginUnitOfWork' must be matched with exactly one call to either 'commitUnitOfWork' or
+         * 'abortUnitOfWork'. When 'abortUnitOfWork' is called, all changes made since the begin
+         * of the unit of work will be rolled back.
          */
         virtual void beginUnitOfWork(OperationContext* opCtx) = 0;
         virtual void commitUnitOfWork() = 0;
-        virtual void endUnitOfWork() = 0;
+        virtual void abortUnitOfWork() = 0;
 
         // WARNING: "commit" in functions below refers to a global journal flush which implicitly
         // commits the current UnitOfWork as well. They are actually stronger than commitUnitOfWork
diff --git a/src/mongo/db/storage/recovery_unit_noop.h b/src/mongo/db/storage/recovery_unit_noop.h
index a424a6c476b..911320bed86 100644
--- a/src/mongo/db/storage/recovery_unit_noop.h
+++ b/src/mongo/db/storage/recovery_unit_noop.h
@@ -37,9 +37,9 @@ namespace mongo {
     class RecoveryUnitNoop : public RecoveryUnit {
     public:
         // TODO implement rollback
-        virtual void beginUnitOfWork(OperationContext* opCtx) {}
-        virtual void commitUnitOfWork() {}
-        virtual void endUnitOfWork() {}
+        void beginUnitOfWork(OperationContext* opCtx) final {}
+        void commitUnitOfWork() final {}
+        void abortUnitOfWork() final {}
 
         virtual void abandonSnapshot() {}
 
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
index 0fb31799cb9..5739df724ce 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_record_store.cpp
@@ -447,7 +447,9 @@ namespace {
             checked_cast<WiredTigerRecoveryUnit*>( txn->releaseRecoveryUnit() );
         invariant( realRecoveryUnit );
         WiredTigerSessionCache* sc = realRecoveryUnit->getSessionCache();
-        txn->setRecoveryUnit( new WiredTigerRecoveryUnit( sc ) );
+        OperationContext::RecoveryUnitState const realRUstate =
+            txn->setRecoveryUnit(new WiredTigerRecoveryUnit(sc),
+                                 OperationContext::kNotInUnitOfWork);
 
         WiredTigerRecoveryUnit::get(txn)->markNoTicketRequired(); // realRecoveryUnit already has
         WT_SESSION* session = WiredTigerRecoveryUnit::get(txn)->getSession(txn)->getSession();
@@ -531,18 +533,18 @@ namespace {
         }
         catch ( const WriteConflictException& wce ) {
             delete txn->releaseRecoveryUnit();
-            txn->setRecoveryUnit( realRecoveryUnit );
+            txn->setRecoveryUnit(realRecoveryUnit, realRUstate);
             log() << "got conflict truncating capped, ignoring";
             return 0;
         }
         catch ( ... ) {
             delete txn->releaseRecoveryUnit();
-            txn->setRecoveryUnit( realRecoveryUnit );
+            txn->setRecoveryUnit(realRecoveryUnit, realRUstate);
             throw;
         }
 
         delete txn->releaseRecoveryUnit();
-        txn->setRecoveryUnit( realRecoveryUnit );
+        txn->setRecoveryUnit(realRecoveryUnit, realRUstate);
         return docsRemoved;
     }
 
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.cpp
index b3330713553..aeb220ebfbb 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.cpp
@@ -82,7 +82,7 @@ namespace mongo {
     WiredTigerRecoveryUnit::WiredTigerRecoveryUnit(WiredTigerSessionCache* sc) :
         _sessionCache( sc ),
         _session( NULL ),
-        _depth(0),
+        _inUnitOfWork(false),
         _active( false ),
         _myTransactionCount( 1 ),
         _everStartedWrite( false ),
@@ -92,7 +92,7 @@ namespace mongo {
     }
 
     WiredTigerRecoveryUnit::~WiredTigerRecoveryUnit() {
-        invariant( _depth == 0 );
+        invariant(!_inUnitOfWork);
         _abort();
         if ( _session ) {
             _sessionCache->releaseSession( _session );
@@ -101,7 +101,7 @@ namespace mongo {
     }
 
     void WiredTigerRecoveryUnit::reportState( BSONObjBuilder* b ) const {
-        b->append("wt_depth", _depth);
+        b->append("wt_inUnitOfWork", _inUnitOfWork);
         b->append("wt_active", _active);
         b->append("wt_everStartedWrite", _everStartedWrite);
         b->append("wt_hasTicket", _ticket.hasTicket());
@@ -151,23 +151,23 @@ namespace mongo {
     }
 
     void WiredTigerRecoveryUnit::beginUnitOfWork(OperationContext* opCtx) {
-        invariant( !_currentlySquirreled );
-        _depth++;
+        invariant(!_inUnitOfWork);
+        invariant(!_currentlySquirreled);
+        _inUnitOfWork = true;
         _everStartedWrite = true;
         _getTicket(opCtx);
     }
 
     void WiredTigerRecoveryUnit::commitUnitOfWork() {
-        if (_depth > 1)
-            return; // only outermost WUOW gets committed.
+        invariant(_inUnitOfWork);
+        _inUnitOfWork = false;
         _commit();
     }
 
-    void WiredTigerRecoveryUnit::endUnitOfWork() {
-        _depth--;
-        if ( _depth == 0 ) {
-            _abort();
-        }
+    void WiredTigerRecoveryUnit::abortUnitOfWork() {
+        invariant(_inUnitOfWork);
+        _inUnitOfWork = false;
+        _abort();
     }
 
     void WiredTigerRecoveryUnit::goingToWaitUntilDurable() {
@@ -189,7 +189,7 @@ namespace mongo {
     }
 
     void WiredTigerRecoveryUnit::registerChange(Change* change) {
-        invariant(_depth > 0);
+        invariant(_inUnitOfWork);
         _changes.push_back(change);
     }
 
@@ -214,7 +214,7 @@ namespace mongo {
     }
 
     void WiredTigerRecoveryUnit::abandonSnapshot() {
-        invariant(_depth == 0);
+        invariant(!_inUnitOfWork);
         if (_active) {
             // Can't be in a WriteUnitOfWork, so safe to rollback
             _txnClose(false);
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.h b/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.h
index 00e5d5d1435..a11e4333d6e 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.h
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_recovery_unit.h
@@ -57,11 +57,9 @@ namespace mongo {
 
         virtual void reportState( BSONObjBuilder* b ) const;
 
-        virtual void beginUnitOfWork(OperationContext* opCtx);
-
-        virtual void commitUnitOfWork();
-
-        virtual void endUnitOfWork();
+        void beginUnitOfWork(OperationContext* opCtx) final;
+        void commitUnitOfWork() final;
+        void abortUnitOfWork() final;
 
         virtual bool waitUntilDurable();
         virtual void goingToWaitUntilDurable();
@@ -88,7 +86,6 @@ namespace mongo {
         void assertInActiveTxn() const;
 
         bool everStartedWrite() const { return _everStartedWrite; }
-        int depth() const { return _depth; }
 
         void setOplogReadTill( const RecordId& loc );
         RecordId getOplogReadTill() const { return _oplogReadTill; }
@@ -109,7 +106,7 @@ namespace mongo {
         WiredTigerSessionCache* _sessionCache; // not owned
         WiredTigerSession* _session; // owned, but from pool
         bool _defaultCommit;
-        int _depth;
+        bool _inUnitOfWork;
         bool _active;
         uint64_t _myTransactionCount;
         bool _everStartedWrite;