SERVER-52694 Rely on Collection copy-on-write instead of restoring Collection in rollback handler

CollectionIndexUsageTrackerDecoration is separated from CollectionQueryInfo as they no longer have the same semantics. CollectionIndexUsageTrackerDecoration is shared between Collection instances and CollectionQueryInfo is not.
Fix UnmanagedCommitManagedRollback mode so it works when write conflict exception is thrown by managing it for rollback

1 files changed, 35 insertions, 21 deletions

diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index 92e0b9c4f96..034aff9d8af 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -166,17 +166,24 @@ Collection* AutoGetCollection::getWritableCollection(CollectionCatalog::Lifetime
 
         // Makes the internal CollectionPtr Yieldable and resets the writable Collection when the
         // write unit of work finishes so we re-fetches and re-clones the Collection if a new write
-        // unit of work is opened.
+        // unit of work is opened. The manage commit flag is used to control if this logic should
+        // happen on commit or not. This depends on the lifetime mode that was used to request the
+        // writable collection.
         class WritableCollectionReset : public RecoveryUnit::Change {
         public:
             WritableCollectionReset(AutoGetCollection& autoColl,
-                                    const Collection* originalCollection)
-                : _autoColl(autoColl), _originalCollection(originalCollection) {}
+                                    const Collection* originalCollection,
+                                    bool manageCommit)
+                : _autoColl(autoColl),
+                  _originalCollection(originalCollection),
+                  _manageCommit(manageCommit) {}
             void commit(boost::optional<Timestamp> commitTime) final {
-                _autoColl._coll = CollectionPtr(_autoColl.getOperationContext(),
-                                                _autoColl._coll.get(),
-                                                LookupCollectionForYieldRestore());
-                _autoColl._writableColl = nullptr;
+                if (_manageCommit) {
+                    _autoColl._coll = CollectionPtr(_autoColl.getOperationContext(),
+                                                    _autoColl._coll.get(),
+                                                    LookupCollectionForYieldRestore());
+                    _autoColl._writableColl = nullptr;
+                }
             }
             void rollback() final {
                 _autoColl._coll = CollectionPtr(_autoColl.getOperationContext(),
@@ -189,14 +196,16 @@ Collection* AutoGetCollection::getWritableCollection(CollectionCatalog::Lifetime
             AutoGetCollection& _autoColl;
             // Used to be able to restore to the original pointer in case of a rollback
             const Collection* _originalCollection;
+            bool _manageCommit;
         };
 
         auto catalog = CollectionCatalog::get(_opCtx);
         _writableColl =
             catalog->lookupCollectionByNamespaceForMetadataWrite(_opCtx, mode, _resolvedNss);
-        if (mode == CollectionCatalog::LifetimeMode::kManagedInWriteUnitOfWork) {
+        if (mode != CollectionCatalog::LifetimeMode::kInplace) {
+            bool peformCommit = mode == CollectionCatalog::LifetimeMode::kManagedInWriteUnitOfWork;
             _opCtx->recoveryUnit()->registerChange(
-                std::make_unique<WritableCollectionReset>(*this, _coll.get()));
+                std::make_unique<WritableCollectionReset>(*this, _coll.get(), peformCommit));
         }
 
         // Set to writable collection. We are no longer yieldable.
@@ -312,10 +321,6 @@ CollectionWriter::~CollectionWriter() {
     if (_sharedImpl) {
         _sharedImpl->_parent = nullptr;
     }
-
-    if (_mode == CollectionCatalog::LifetimeMode::kUnmanagedClone && _writableCollection) {
-        CollectionCatalog::discardUnmanagedClone(_opCtx, _writableCollection);
-    }
 }
 
 Collection* CollectionWriter::getWritableCollection() {
@@ -326,14 +331,19 @@ Collection* CollectionWriter::getWritableCollection() {
         // Resets the writable Collection when the write unit of work finishes so we re-fetch and
         // re-clone the Collection if a new write unit of work is opened. Holds the back pointer to
         // the CollectionWriter via a shared_ptr so we can detect if the instance is already
-        // destroyed.
+        // destroyed. The manage commit flag is used to control if this logic should
+        // happen on commit or not. This depends on the lifetime mode that was used to request the
+        // writable collection.
         class WritableCollectionReset : public RecoveryUnit::Change {
         public:
             WritableCollectionReset(std::shared_ptr<SharedImpl> shared,
-                                    CollectionPtr rollbackCollection)
-                : _shared(std::move(shared)), _rollbackCollection(std::move(rollbackCollection)) {}
+                                    CollectionPtr rollbackCollection,
+                                    bool manageCommit)
+                : _shared(std::move(shared)),
+                  _rollbackCollection(std::move(rollbackCollection)),
+                  _manageCommit(manageCommit) {}
             void commit(boost::optional<Timestamp> commitTime) final {
-                if (_shared->_parent)
+                if (_manageCommit && _shared->_parent)
                     _shared->_parent->_writableCollection = nullptr;
             }
             void rollback() final {
@@ -346,15 +356,19 @@ Collection* CollectionWriter::getWritableCollection() {
         private:
             std::shared_ptr<SharedImpl> _shared;
             CollectionPtr _rollbackCollection;
+            bool _manageCommit;
         };
 
         // If we are using our stored Collection then we are not managed by an AutoGetCollection and
         // we need to manage lifetime here.
-        if (_mode == CollectionCatalog::LifetimeMode::kManagedInWriteUnitOfWork) {
+        if (_mode != CollectionCatalog::LifetimeMode::kInplace) {
             bool usingStoredCollection = *_collection == _storedCollection;
+            bool performCommit =
+                _mode == CollectionCatalog::LifetimeMode::kManagedInWriteUnitOfWork;
             _opCtx->recoveryUnit()->registerChange(std::make_unique<WritableCollectionReset>(
                 _sharedImpl,
-                usingStoredCollection ? std::move(_storedCollection) : CollectionPtr()));
+                usingStoredCollection ? std::move(_storedCollection) : CollectionPtr(),
+                performCommit));
             if (usingStoredCollection) {
                 _storedCollection = _writableCollection;
             }
@@ -364,8 +378,8 @@ Collection* CollectionWriter::getWritableCollection() {
 }
 
 void CollectionWriter::commitToCatalog() {
-    dassert(_mode == CollectionCatalog::LifetimeMode::kUnmanagedClone);
-    dassert(_writableCollection);
+    invariant(_mode == CollectionCatalog::LifetimeMode::kUnmanagedCommitManagedRollback);
+    invariant(_writableCollection);
     CollectionCatalog::commitUnmanagedClone(_opCtx, _writableCollection);
     _writableCollection = nullptr;
 }