PHPLIB-1400: Work around MigrationConflict errors on sharded clusters

tests/UnifiedSpecTests/CollectionData.php

@@ -39,22 +39,22 @@ public function __construct(stdClass $o)
         $this->documents = $o->documents;
     }
 
-    public function prepareInitialData(Client $client): void
+    public function prepareInitialData(Client $client, ?Session $session = null): void
     {
         $database = $client->selectDatabase(
             $this->databaseName,
             ['writeConcern' => new WriteConcern(WriteConcern::MAJORITY)],
         );
 
-        $database->dropCollection($this->collectionName);
+        $database->dropCollection($this->collectionName, ['session' => $session]);
 
         if (empty($this->documents)) {
-            $database->createCollection($this->collectionName);
+            $database->createCollection($this->collectionName, ['session' => $session]);
 
             return;
         }
 
-        $database->selectCollection($this->collectionName)->insertMany($this->documents);
+        $database->selectCollection($this->collectionName)->insertMany($this->documents, ['session' => $session]);
     }
 
     public function assertOutcome(Client $client): void

tests/UnifiedSpecTests/Context.php

@@ -58,6 +58,8 @@ final class Context
 
     private string $multiMongosUri;
 
+    private ?object $advanceClusterTime = null;
+
     public function __construct(Client $internalClient, string $uri)
     {
         $this->entityMap = new EntityMap();
@@ -150,6 +152,17 @@ public function setActiveClient(?string $clientId = null): void
         $this->activeClient = $clientId;
     }
 
+    /**
+     * Set a cluster time to use for advancing newly created session entities.
+     *
+     * This is used to ensure causal consistency with initialData collections
+     * in sharded environments (see: DRIVERS-2816).
+     */
+    public function setAdvanceClusterTime(?object $clusterTime): void
+    {
+        $this->advanceClusterTime = $clusterTime;
+    }
+
     public function isInLoop(): bool
     {
         return $this->inLoop;
@@ -463,7 +476,13 @@ private function createSession(string $id, stdClass $o): void
             $options = self::prepareSessionOptions((array) $o->sessionOptions);
         }
 
-        $this->entityMap->set($id, $client->startSession($options), $clientId);
+        $session = $client->startSession($options);
+
+        if ($this->advanceClusterTime !== null) {
+            $session->advanceClusterTime($this->advanceClusterTime);
+        }
+
+        $this->entityMap->set($id, $session, $clientId);
     }
 
     private function createBucket(string $id, stdClass $o): void

tests/UnifiedSpecTests/UnifiedTestRunner.php

@@ -183,12 +183,14 @@ private function doTestCase(stdClass $test, string $schemaVersion, ?array $runOn
             $this->checkRunOnRequirements($test->runOnRequirements);
         }
 
-        if (isset($initialData)) {
-            $this->prepareInitialData($initialData);
-        }
+        assertIsArray($test->operations);
 
         $context = $this->createContext();
 
+        if (isset($initialData)) {
+            $this->prepareInitialData($initialData, $context, $this->isAdvanceClusterTimeNeeded($test->operations));
+        }
+
         /* If an EntityMap observer has been configured, assign the Context's
          * EntityMap to a class property so it can later be accessed from run(),
          * irrespective of whether this test succeeds or fails. */
@@ -200,7 +202,6 @@ private function doTestCase(stdClass $test, string $schemaVersion, ?array $runOn
             $context->createEntities($createEntities);
         }
 
-        assertIsArray($test->operations);
         $this->preventStaleDbVersionError($test->operations, $context);
 
         $context->startEventObservers();
@@ -423,17 +424,48 @@ private function assertOutcome(array $outcome): void
         }
     }
 
-    private function prepareInitialData(array $initialData): void
+    private function prepareInitialData(array $initialData, Context $context, bool $isAdvanceClusterTimeNeeded): void
     {
         assertNotEmpty($initialData);
         assertContainsOnly('object', $initialData);
 
+        /* In order to avoid MigrationConflict errors on sharded clusters, use the cluster time obtained from creating
+         * collections to advance session entities. This is necessary because initialData uses an internal MongoClient,
+         * which will not share/gossip its cluster time via the test entities. */
+        if ($isAdvanceClusterTimeNeeded) {
+            $session = $this->internalClient->startSession();
+        }
+
         foreach ($initialData as $data) {
             $collectionData = new CollectionData($data);
-            $collectionData->prepareInitialData($this->internalClient);
+            $collectionData->prepareInitialData($this->internalClient, $session ?? null);
+        }
+
+        if (isset($session)) {
+            $context->setAdvanceClusterTime($session->getClusterTime());
         }
     }
 
+    /**
+     * Work around potential MigrationConflict errors on sharded clusters.
+     */
+    private function isAdvanceClusterTimeNeeded(array $operations): bool
+    {
+        if (! in_array($this->getPrimaryServer()->getType(), [Server::TYPE_MONGOS, Server::TYPE_LOAD_BALANCER], true)) {
+            return false;
+        }
+
+        foreach ($operations as $operation) {
+            switch ($operation->name) {
+                case 'startTransaction':
+                case 'withTransaction':
+                    return true;
+            }
+        }
+
+        return false;
+    }
+
     /**
      * Work around potential error executing distinct on sharded clusters.
      *