Preserve collection changes until transaction has succeeded

lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php

@@ -1512,14 +1512,11 @@ private function handleCollections(object $document, array $options): void
             $collections[] = $coll;
         }
 
-        if (! empty($collections)) {
-            $this->cp->update($document, $collections, $options);
+        if (empty($collections)) {
+            return;
         }
 
-        // Take new snapshots from visited collections
-        foreach ($this->uow->getVisitedCollections($document) as $coll) {
-            $coll->takeSnapshot();
-        }
+        $this->cp->update($document, $collections, $options);
     }
 
     /**

lib/Doctrine/ODM/MongoDB/UnitOfWork.php

@@ -477,6 +477,12 @@ function (Session $session) use ($options): void {
             $this->evm->dispatchEvent(Events::postFlush, new Event\PostFlushEventArgs($this->dm));
 
             // Clear up
+            foreach ($this->visitedCollections as $collections) {
+                foreach ($collections as $coll) {
+                    $coll->takeSnapshot();
+                }
+            }
+
             $this->scheduledDocumentInsertions  =
             $this->scheduledDocumentUpserts     =
             $this->scheduledDocumentUpdates     =

tests/Doctrine/ODM/MongoDB/Tests/Functional/CommitImprovementTest.php

@@ -9,7 +9,6 @@
 use Doctrine\ODM\MongoDB\Event\LifecycleEventArgs;
 use Doctrine\ODM\MongoDB\Events;
 use Doctrine\ODM\MongoDB\LockException;
-use Doctrine\ODM\MongoDB\PersistentCollection\PersistentCollectionInterface;
 use Doctrine\ODM\MongoDB\Tests\BaseTestCase;
 use Documents\Phonebook;
 use Documents\Phonenumber;
@@ -87,41 +86,6 @@ public function testCollectionsAreUpdatedJustAfterOwningDocument(): void
         self::assertEquals('87654321', $phonenumbers[1]->getPhonenumber());
     }
 
-    /**
-     * This test checks few things:
-     *  - if collections were updated after post* events, our changes would be saved
-     *  - if collection snapshot would be taken after post* events, collection
-     *    wouldn't be dirty and wouldn't be updated in next flush
-     */
-    public function testChangingCollectionInPostEventsHasNoIllEffects(): void
-    {
-        $this->dm->getEventManager()->addEventSubscriber(new PhonenumberMachine());
-
-        $user = new VersionedUser();
-        $user->setUsername('malarzm');
-        $this->dm->persist($user);
-        $this->dm->flush();
-
-        $phoneNumbers = $user->getPhonenumbers();
-        self::assertCount(1, $phoneNumbers); // so we got a number on postPersist
-        self::assertInstanceOf(PersistentCollectionInterface::class, $phoneNumbers); // so we got a number on postPersist
-        self::assertTrue($phoneNumbers->isDirty()); // but they should be dirty
-
-        $collection = $this->dm->getDocumentCollection($user::class);
-        $inDb       = $collection->findOne();
-        self::assertArrayNotHasKey('phonenumbers', $inDb, 'Collection modified in postPersist should not be in database without recomputing change set');
-
-        $this->dm->flush();
-
-        $phoneNumbers = $user->getPhonenumbers();
-        self::assertInstanceOf(PersistentCollectionInterface::class, $phoneNumbers);
-        self::assertCount(2, $phoneNumbers); // so we got a number on postUpdate
-        self::assertTrue($phoneNumbers->isDirty()); // but they should be dirty
-
-        $inDb = $collection->findOne();
-        self::assertCount(1, $inDb['phonenumbers'], 'Collection changes from postUpdate should not be in database');
-    }
-
     public function testSchedulingCollectionDeletionAfterSchedulingForUpdate(): void
     {
         $user = new User();

tests/Doctrine/ODM/MongoDB/Tests/UnitOfWorkTransactionalCommitConsistencyTest.php

@@ -9,6 +9,7 @@
 use Documents\Address;
 use Documents\ForumUser;
 use Documents\FriendUser;
+use Documents\Phonenumber;
 use Documents\User;
 use MongoDB\BSON\ObjectId;
 use MongoDB\Client;
@@ -561,6 +562,36 @@ public function testTransientDeleteErrorWithEmbeddedDocument(): void
         self::assertFalse($this->uow->isScheduledForDelete($user));
     }
 
+    public function testTransientErrorPreservesCollectionChangesets(): void
+    {
+        // Create a dummy user so that we later have a user to remove
+        $fooUser = new User();
+        $fooUser->setUsername('foo');
+        $this->uow->persist($fooUser);
+        $this->uow->commit();
+
+        // Create a new user with a collection update
+        $alcaeus = new User();
+        // Set an identifier to force an upsert, which causes separate queries
+        // for collection updates
+        $alcaeus->setId(new ObjectId());
+        $alcaeus->setUsername('alcaeus');
+        $alcaeus->getPhonenumbers()->add(new Phonenumber('12345'));
+        $this->uow->persist($alcaeus);
+
+        // Remove fooUser and create a transient failpoint to force the deletion
+        // to fail. This exposes the issue with collections
+        $this->uow->remove($fooUser);
+        $this->createTransientFailPoint('delete');
+
+        $this->uow->commit();
+
+        $this->dm->clear();
+        $check = $this->dm->find(User::class, $alcaeus->getId());
+        self::assertNotNull($check);
+        self::assertCount(1, $check->getPhonenumbers());
+    }
+
     /** Create a document manager with a single host to ensure failpoints target the correct server */
     protected static function createTestDocumentManager(): DocumentManager
     {