rdbms: avoid LoadBalancer::getConnection waste when using $groups

Add a "readIndexByGroup" field that does the same thing that the
generic read index does except for query grouped connections.

Also remove some pointless checks at the end of getReaderIndex().

Change-Id: I3bd6295be4355ee930960f89ccb07811dd8c5545

diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 3fd0189..ffb7a34 100644 (file)
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -105,6 +105,8 @@ class LoadBalancer implements ILoadBalancer {
        private $errorConnection;
        /** @var int The generic (not query grouped) replica DB index */
        private $genericReadIndex = -1;
+       /** @var int[] The group replica DB indexes keyed by group */
+       private $readIndexByGroup = [];
        /** @var bool|DBMasterPos False if not set */
        private $waitForPos;
        /** @var bool Whether the generic reader fell back to a lagged replica DB */
@@ -398,9 +400,12 @@ class LoadBalancer implements ILoadBalancer {
                if ( count( $this->servers ) == 1 ) {
                        // Skip the load balancing if there's only one server
                        return $this->getWriterIndex();
-               } elseif ( $group === false && $this->genericReadIndex >= 0 ) {
-                       // A generic reader index was already selected and "waitForPos" was handled
-                       return $this->genericReadIndex;
+               }
+
+               $index = $this->getExistingReaderIndex( $group );
+               if ( $index >= 0 ) {
+                       // A reader index was already selected and "waitForPos" was handled
+                       return $index;
                }
 
                if ( $group !== false ) {
@@ -435,11 +440,11 @@ class LoadBalancer implements ILoadBalancer {
                        $laggedReplicaMode = true;
                }
 
-               if ( $this->genericReadIndex < 0 && $this->genericLoads[$i] > 0 && $group === false ) {
-                       // Cache the generic (ungrouped) reader index for future DB_REPLICA handles
-                       $this->genericReadIndex = $i;
-                       // Record if the generic reader index is in "lagged replica DB" mode
-                       $this->laggedReplicaMode = ( $laggedReplicaMode || $this->laggedReplicaMode );
+               // Cache the reader index for future DB_REPLICA handles
+               $this->setExistingReaderIndex( $group, $i );
+               // Record whether the generic reader index is in "lagged replica DB" mode
+               if ( $group === false && $laggedReplicaMode ) {
+                       $this->laggedReplicaMode = true;
                }
 
                $serverName = $this->getServerName( $i );
@@ -448,6 +453,40 @@ class LoadBalancer implements ILoadBalancer {
                return $i;
        }
 
+       /**
+        * Get the server index chosen by the load balancer for use with the given query group
+        *
+        * @param string|bool $group Query group; use false for the generic group
+        * @return int Server index or -1 if none was chosen
+        */
+       protected function getExistingReaderIndex( $group ) {
+               if ( $group === false ) {
+                       $index = $this->genericReadIndex;
+               } else {
+                       $index = $this->readIndexByGroup[$group] ?? -1;
+               }
+
+               return $index;
+       }
+
+       /**
+        * Set the server index chosen by the load balancer for use with the given query group
+        *
+        * @param string|bool $group Query group; use false for the generic group
+        * @param int $index The index of a specific server
+        */
+       private function setExistingReaderIndex( $group, $index ) {
+               if ( $index < 0 ) {
+                       throw new UnexpectedValueException( "Cannot set a negative read server index" );
+               }
+
+               if ( $group === false ) {
+                       $this->genericReadIndex = $index;
+               } else {
+                       $this->readIndexByGroup[$group] = $index;
+               }
+       }
+
        /**
         * @param array $loads List of server weights
         * @param string|bool $domain

diff --git a/tests/phpunit/includes/db/LoadBalancerTest.php b/tests/phpunit/includes/db/LoadBalancerTest.php
index 7fc070c..169e4bf 100644 (file)
--- a/tests/phpunit/includes/db/LoadBalancerTest.php
+++ b/tests/phpunit/includes/db/LoadBalancerTest.php
@@ -26,6 +26,7 @@ use Wikimedia\Rdbms\DatabaseDomain;
 use Wikimedia\Rdbms\Database;
 use Wikimedia\Rdbms\LoadBalancer;
 use Wikimedia\Rdbms\LoadMonitorNull;
+use Wikimedia\TestingAccessWrapper;
 
 /**
  * @group Database
@@ -165,7 +166,8 @@ class LoadBalancerTest extends MediaWikiTestCase {
                global $wgDBserver, $wgDBname, $wgDBuser, $wgDBpassword, $wgDBtype, $wgSQLiteDataDir;
 
                $servers = [
-                       [ // master
+                       // Master DB
+                       0 => [
                                'host' => $wgDBserver,
                                'dbname' => $wgDBname,
                                'tablePrefix' => $this->dbPrefix(),
@@ -176,7 +178,19 @@ class LoadBalancerTest extends MediaWikiTestCase {
                                'load' => 0,
                                'flags' => $flags
                        ],
-                       [ // emulated replica
+                       // Main replica DBs
+                       1 => [
+                               'host' => $wgDBserver,
+                               'dbname' => $wgDBname,
+                               'tablePrefix' => $this->dbPrefix(),
+                               'user' => $wgDBuser,
+                               'password' => $wgDBpassword,
+                               'type' => $wgDBtype,
+                               'dbDirectory' => $wgSQLiteDataDir,
+                               'load' => 100,
+                               'flags' => $flags
+                       ],
+                       2 => [
                                'host' => $wgDBserver,
                                'dbname' => $wgDBname,
                                'tablePrefix' => $this->dbPrefix(),
@@ -186,6 +200,66 @@ class LoadBalancerTest extends MediaWikiTestCase {
                                'dbDirectory' => $wgSQLiteDataDir,
                                'load' => 100,
                                'flags' => $flags
+                       ],
+                       // RC replica DBs
+                       3 => [
+                               'host' => $wgDBserver,
+                               'dbname' => $wgDBname,
+                               'tablePrefix' => $this->dbPrefix(),
+                               'user' => $wgDBuser,
+                               'password' => $wgDBpassword,
+                               'type' => $wgDBtype,
+                               'dbDirectory' => $wgSQLiteDataDir,
+                               'load' => 0,
+                               'groupLoads' => [
+                                       'recentchanges' => 100,
+                                       'watchlist' => 100
+                               ],
+                               'flags' => $flags
+                       ],
+                       // Logging replica DBs
+                       4 => [
+                               'host' => $wgDBserver,
+                               'dbname' => $wgDBname,
+                               'tablePrefix' => $this->dbPrefix(),
+                               'user' => $wgDBuser,
+                               'password' => $wgDBpassword,
+                               'type' => $wgDBtype,
+                               'dbDirectory' => $wgSQLiteDataDir,
+                               'load' => 0,
+                               'groupLoads' => [
+                                       'logging' => 100
+                               ],
+                               'flags' => $flags
+                       ],
+                       5 => [
+                               'host' => $wgDBserver,
+                               'dbname' => $wgDBname,
+                               'tablePrefix' => $this->dbPrefix(),
+                               'user' => $wgDBuser,
+                               'password' => $wgDBpassword,
+                               'type' => $wgDBtype,
+                               'dbDirectory' => $wgSQLiteDataDir,
+                               'load' => 0,
+                               'groupLoads' => [
+                                       'logging' => 100
+                               ],
+                               'flags' => $flags
+                       ],
+                       // Maintenance query replica DBs
+                       6 => [
+                               'host' => $wgDBserver,
+                               'dbname' => $wgDBname,
+                               'tablePrefix' => $this->dbPrefix(),
+                               'user' => $wgDBuser,
+                               'password' => $wgDBpassword,
+                               'type' => $wgDBtype,
+                               'dbDirectory' => $wgSQLiteDataDir,
+                               'load' => 0,
+                               'groupLoads' => [
+                                       'vslow' => 100
+                               ],
+                               'flags' => $flags
                        ]
                ];
 
@@ -488,4 +562,47 @@ class LoadBalancerTest extends MediaWikiTestCase {
 
                $rConn->insert( 'test', [ 't' => 1 ], __METHOD__ );
        }
+
+       public function testQueryGroupIndex() {
+               $lb = $this->newMultiServerLocalLoadBalancer();
+               /** @var LoadBalancer $lbWrapper */
+               $lbWrapper = TestingAccessWrapper::newFromObject( $lb );
+
+               $rGeneric = $lb->getConnectionRef( DB_REPLICA );
+               $mainIndexPicked = $rGeneric->getLBInfo( 'serverIndex' );
+
+               $this->assertEquals( $mainIndexPicked, $lbWrapper->getExistingReaderIndex( false ) );
+               $this->assertTrue( in_array( $mainIndexPicked, [ 1, 2 ] ) );
+               for ( $i = 0; $i < 300; ++$i ) {
+                       $rLog = $lb->getConnectionRef( DB_REPLICA, [] );
+                       $this->assertEquals(
+                               $mainIndexPicked,
+                               $rLog->getLBInfo( 'serverIndex' ),
+                               "Main index unchanged" );
+               }
+
+               $rRC = $lb->getConnectionRef( DB_REPLICA, [ 'recentchanges' ] );
+               $rWL = $lb->getConnectionRef( DB_REPLICA, [ 'watchlist' ] );
+
+               $this->assertEquals( 3, $rRC->getLBInfo( 'serverIndex' ) );
+               $this->assertEquals( 3, $rWL->getLBInfo( 'serverIndex' ) );
+
+               $rLog = $lb->getConnectionRef( DB_REPLICA, [ 'logging', 'watchlist' ] );
+               $logIndexPicked = $rLog->getLBInfo( 'serverIndex' );
+
+               $this->assertEquals( $logIndexPicked, $lbWrapper->getExistingReaderIndex( 'logging' ) );
+               $this->assertTrue( in_array( $logIndexPicked, [ 4, 5 ] ) );
+
+               for ( $i = 0; $i < 300; ++$i ) {
+                       $rLog = $lb->getConnectionRef( DB_REPLICA, [ 'logging', 'watchlist' ] );
+                       $this->assertEquals(
+                               $logIndexPicked, $rLog->getLBInfo( 'serverIndex' ), "Index unchanged" );
+               }
+
+               $rVslow = $lb->getConnectionRef( DB_REPLICA, [ 'vslow', 'logging' ] );
+               $vslowIndexPicked = $rVslow->getLBInfo( 'serverIndex' );
+
+               $this->assertEquals( $vslowIndexPicked, $lbWrapper->getExistingReaderIndex( 'vslow' ) );
+               $this->assertEquals( 6, $vslowIndexPicked );
+       }
 }