PHPLIB-466: Respect pinned session during server selection

src/Client.php

@@ -170,7 +170,7 @@ class Client
             $options['typeMap'] = $this->typeMap;
         }
 
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         if (! isset($options['writeConcern']) && server_supports_feature($server, self::$wireVersionForWritableCommandWriteConcern)) {
             $options['writeConcern'] = $this->writeConcern;
@@ -246,7 +246,7 @@ class Client
     public function listDatabases(array $options = [])
     {
         $operation = new ListDatabases($options);
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         return $operation->execute($server);
     }
@@ -311,7 +311,7 @@ class Client
             $options['readPreference'] = $this->readPreference;
         }
 
-        $server = $this->manager->selectServer($options['readPreference']);
+        $server = select_server($this->manager, $options['readPreference'], extract_session_from_options($options));
 
         if (! isset($options['readConcern']) && server_supports_feature($server, self::$wireVersionForReadConcern)) {
             $options['readConcern'] = $this->readConcern;

src/Database.php

@@ -206,7 +206,7 @@ class Database
             $options['readPreference'] = new ReadPreference(ReadPreference::RP_PRIMARY);
         }
 
-        $server = $this->manager->selectServer($options['readPreference']);
+        $server = select_server($this->manager, $options['readPreference'], extract_session_from_options($options));
 
         /* MongoDB 4.2 and later supports a read concern when an $out stage is
          * being used, but earlier versions do not.
@@ -258,7 +258,7 @@ class Database
         }
 
         $operation = new DatabaseCommand($this->databaseName, $command, $options);
-        $server = $this->manager->selectServer($options['readPreference']);
+        $server = select_server($this->manager, $options['readPreference'], extract_session_from_options($options));
 
         return $operation->execute($server);
     }
@@ -280,7 +280,7 @@ class Database
             $options['typeMap'] = $this->typeMap;
         }
 
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         if (! isset($options['writeConcern']) && server_supports_feature($server, self::$wireVersionForWritableCommandWriteConcern)) {
             $options['writeConcern'] = $this->writeConcern;
@@ -307,7 +307,7 @@ class Database
             $options['typeMap'] = $this->typeMap;
         }
 
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         if (! isset($options['writeConcern']) && server_supports_feature($server, self::$wireVersionForWritableCommandWriteConcern)) {
             $options['writeConcern'] = $this->writeConcern;
@@ -335,7 +335,7 @@ class Database
             $options['typeMap'] = $this->typeMap;
         }
 
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         if (! isset($options['writeConcern']) && server_supports_feature($server, self::$wireVersionForWritableCommandWriteConcern)) {
             $options['writeConcern'] = $this->writeConcern;
@@ -420,7 +420,7 @@ class Database
     public function listCollections(array $options = [])
     {
         $operation = new ListCollections($this->databaseName, $options);
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         return $operation->execute($server);
     }
@@ -442,7 +442,7 @@ class Database
             $options['typeMap'] = $this->typeMap;
         }
 
-        $server = $this->manager->selectServer(new ReadPreference(ReadPreference::RP_PRIMARY));
+        $server = select_server($this->manager, new ReadPreference(ReadPreference::RP_PRIMARY), extract_session_from_options($options));
 
         if (! isset($options['writeConcern']) && server_supports_feature($server, self::$wireVersionForWritableCommandWriteConcern)) {
             $options['writeConcern'] = $this->writeConcern;
@@ -509,7 +509,7 @@ class Database
             $options['readPreference'] = $this->readPreference;
         }
 
-        $server = $this->manager->selectServer($options['readPreference']);
+        $server = select_server($this->manager, $options['readPreference'], extract_session_from_options($options));
 
         if (! isset($options['readConcern']) && server_supports_feature($server, self::$wireVersionForReadConcern)) {
             $options['readConcern'] = $this->readConcern;

src/Operation/Watch.php

@@ -40,6 +40,8 @@ use function is_object;
 use function is_string;
 use function MongoDB\Driver\Monitoring\addSubscriber;
 use function MongoDB\Driver\Monitoring\removeSubscriber;
+use function MongoDB\extract_session_from_options;
+use function MongoDB\select_server;
 use function MongoDB\server_supports_feature;
 
 /**
@@ -375,8 +377,11 @@ class Watch implements Executable, /* @internal */ CommandSubscriber
 
         $this->hasResumed = true;
 
-        // Select a new server using the original read preference
-        $server = $this->manager->selectServer($this->aggregateOptions['readPreference']);
+        /* Select a new server using the original read preference. While watch
+         * is not usable within transactions, we still check if there is a
+         * pinned session. This is to avoid an ambiguous error message about
+         * running a command on the wrong server. */
+        $server = select_server($this->manager, $this->aggregateOptions['readPreference'], extract_session_from_options($this->aggregateOptions));
 
         $resumeOption = isset($this->changeStreamOptions['startAfter']) && ! $hasAdvanced ? 'startAfter' : 'resumeAfter';
 

src/functions.php

@@ -19,6 +19,8 @@ namespace MongoDB;
 
 use Exception;
 use MongoDB\BSON\Serializable;
+use MongoDB\Driver\Manager;
+use MongoDB\Driver\ReadPreference;
 use MongoDB\Driver\Server;
 use MongoDB\Driver\Session;
 use MongoDB\Exception\InvalidArgumentException;
@@ -373,3 +375,32 @@ function with_transaction(Session $session, callable $callback, array $transacti
     $operation = new WithTransaction($callback, $transactionOptions);
     $operation->execute($session);
 }
+
+/**
+ * Returns the session option if it is set and valid.
+ *
+ * @internal
+ * @param array $options
+ * @return Session|null
+ */
+function extract_session_from_options(array $options)
+{
+    if (! isset($options['session']) || ! $options['session'] instanceof Session) {
+        return null;
+    }
+
+    return $options['session'];
+}
+
+/**
+ * Performs server selection, respecting the server a session may be pinned to
+ *
+ * @internal
+ * @return Server
+ */
+function select_server(Manager $manager, ReadPreference $readPreference = null, Session $session = null)
+{
+    $server = $session !== null ? $session->getServer() : null;
+
+    return $server ?: $manager->selectServer($readPreference);
+}