Refactor LoadBalancer::setDomainPrefix() to use forEachOpen* method

This makes the code more consistent with the other class methods
that iterate over connection handles and less fragile.

Change-Id: Ie2d8efb80a66a5c87b5b3edd64cb6bc4268fa911

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 677a69f..72217da 100644 (file)
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -1625,13 +1625,19 @@ class LoadBalancer implements ILoadBalancer {
        }
 
        public function setDomainPrefix( $prefix ) {
-               if ( $this->mConns[self::KEY_FOREIGN_INUSE] ) {
-                       // Do not switch connections to explicit foreign domains unless marked as free
-                       $domains = [];
-                       foreach ( $this->mConns[self::KEY_FOREIGN_INUSE] as $i => $connsByDomain ) {
-                               $domains = array_merge( $domains, array_keys( $connsByDomain ) );
+               // Find connections to explicit foreign domains still marked as in-use...
+               $domainsInUse = [];
+               $this->forEachOpenConnection( function ( IDatabase $conn ) use ( &$domainsInUse ) {
+                       // Once reuseConnection() is called on a handle, its reference count goes from 1 to 0.
+                       // Until then, it is still in use by the caller (explicitly or via DBConnRef scope).
+                       if ( $conn->getLBInfo( 'foreignPoolRefCount' ) > 0 ) {
+                               $domainsInUse[] = $conn->getDomainID();
                        }
-                       $domains = implode( ', ', $domains );
+               } );
+
+               // Do not switch connections to explicit foreign domains unless marked as safe
+               if ( $domainsInUse ) {
+                       $domains = implode( ', ', $domainsInUse );
                        throw new DBUnexpectedError( null,
                                "Foreign domain connections are still in use ($domains)." );
                }