tests/phpunit/includes/db/LoadBalancerTest.php

   1 <?php
   2 /**
   3  * This program is free software; you can redistribute it and/or modify
   4  * it under the terms of the GNU General Public License as published by
   5  * the Free Software Foundation; either version 2 of the License, or
   6  * (at your option) any later version.
   7  *
   8  * This program is distributed in the hope that it will be useful,
   9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  11  * GNU General Public License for more details.
  12  *
  13  * You should have received a copy of the GNU General Public License along
  14  * with this program; if not, write to the Free Software Foundation, Inc.,
  15  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16  * http://www.gnu.org/copyleft/gpl.html
  17  *
  18  * @file
  19  */
  20
  21 use Wikimedia\Rdbms\ChronologyProtector;
  22 use Wikimedia\Rdbms\Database;
  23 use Wikimedia\Rdbms\DatabaseDomain;
  24 use Wikimedia\Rdbms\DBConnRef;
  25 use Wikimedia\Rdbms\DBReadOnlyRoleError;
  26 use Wikimedia\Rdbms\IDatabase;
  27 use Wikimedia\Rdbms\IMaintainableDatabase;
  28 use Wikimedia\Rdbms\LoadBalancer;
  29 use Wikimedia\Rdbms\LoadMonitorNull;
  30 use Wikimedia\Rdbms\ServerInfo;
  31 use Wikimedia\Rdbms\TransactionManager;
  32 use Wikimedia\TestingAccessWrapper;
  33
  34 /**
  35  * @group Database
  36  * @group medium
  37  * @covers \Wikimedia\Rdbms\LoadBalancer
  38  * @covers \Wikimedia\Rdbms\ServerInfo
  39  */
  40 class LoadBalancerTest extends MediaWikiIntegrationTestCase {
  41         private function makeServerConfig( $flags = DBO_DEFAULT ) {
  42                 global $wgDBserver, $wgDBport, $wgDBname, $wgDBuser, $wgDBpassword, $wgDBtype;
  43                 global $wgSQLiteDataDir;
  44
  45                 return [
  46                         'host' => $wgDBserver,
  47                         'port' => $wgDBport,
  48                         'serverName' => 'testhost',
  49                         'dbname' => $wgDBname,
  50                         'tablePrefix' => self::dbPrefix(),
  51                         'user' => $wgDBuser,
  52                         'password' => $wgDBpassword,
  53                         'type' => $wgDBtype,
  54                         'dbDirectory' => $wgSQLiteDataDir,
  55                         'load' => 0,
  56                         'flags' => $flags
  57                 ];
  58         }
  59
  60         public function testWithoutReplica() {
  61                 global $wgDBname, $wgDBmwschema;
  62
  63                 $called = false;
  64                 $chronologyProtector = $this->createMock( ChronologyProtector::class );
  65                 $chronologyProtector->method( 'getSessionPrimaryPos' )
  66                         ->willReturnCallback(
  67                                 static function () use ( &$called ) {
  68                                         $called = true;
  69                                 }
  70                         );
  71                 $lb = new LoadBalancer( [
  72                         // Simulate web request with DBO_TRX
  73                         'servers' => [ $this->makeServerConfig( DBO_TRX ) ],
  74                         'logger' => MediaWiki\Logger\LoggerFactory::getInstance( 'rdbms' ),
  75                         'localDomain' => new DatabaseDomain( $wgDBname, $wgDBmwschema, self::dbPrefix() ),
  76                         'chronologyProtector' => $chronologyProtector,
  77                         'clusterName' => 'xyz'
  78                 ] );
  79
  80                 $this->assertSame( 1, $lb->getServerCount() );
  81                 $this->assertFalse( $lb->hasReplicaServers() );
  82                 $this->assertFalse( $lb->hasStreamingReplicaServers() );
  83                 $this->assertSame( 'xyz', $lb->getClusterName() );
  84
  85                 $ld = DatabaseDomain::newFromId( $lb->getLocalDomainID() );
  86                 $this->assertSame( $wgDBname, $ld->getDatabase(), 'local domain DB set' );
  87                 $this->assertSame( self::dbPrefix(), $ld->getTablePrefix(), 'local domain prefix set' );
  88                 $this->assertSame( 'my_test_wiki', $lb->resolveDomainID( 'my_test_wiki' ) );
  89                 $this->assertSame( $ld->getId(), $lb->resolveDomainID( false ) );
  90                 $this->assertSame( $ld->getId(), $lb->resolveDomainID( $ld ) );
  91                 $this->assertFalse( $called );
  92
  93                 $dbw = $lb->getConnection( DB_PRIMARY );
  94                 $dbw->getServerName();
  95                 $this->assertFalse( $called, "getServerName() optimized for DB_PRIMARY" );
  96
  97                 $dbw->ensureConnection();
  98                 $this->assertFalse( $called, "Session replication pos not used with single server" );
  99                 $this->assertTrue( $dbw->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on master" );
 100                 $this->assertWriteAllowed( $dbw );
 101
 102                 $dbr = $lb->getConnection( DB_REPLICA );
 103                 $this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on replica" );
 104
 105                 if ( !$lb->getServerAttributes( ServerInfo::WRITER_INDEX )[Database::ATTR_DB_LEVEL_LOCKING] ) {
 106                         $dbwAC1 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 107                         $this->assertFalse(
 108                                 $dbwAC1->getFlag( $dbw::DBO_TRX ),
 109                                 "No DBO_TRX with CONN_TRX_AUTOCOMMIT"
 110                         );
 111                         $this->assertTrue( $dbw->getFlag( $dbw::DBO_TRX ), "DBO_TRX still set on master" );
 112                         $this->assertUnsharedHandle( $dbw, $dbwAC1, "CONN_TRX_AUTOCOMMIT separate connection" );
 113
 114                         $dbrAC1 = $lb->getConnection( DB_REPLICA, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 115                         $this->assertFalse(
 116                                 $dbrAC1->getFlag( $dbw::DBO_TRX ),
 117                                 "No DBO_TRX with CONN_TRX_AUTOCOMMIT"
 118                         );
 119                         $this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX still set on replica" );
 120                         $this->assertUnsharedHandle( $dbr, $dbrAC1, "CONN_TRX_AUTOCOMMIT separate connection" );
 121
 122                         $dbwAC2 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 123                         $dbwAC2->ensureConnection();
 124                         $this->assertSharedHandle( $dbwAC2, $dbwAC1, "CONN_TRX_AUTOCOMMIT reuses connections" );
 125
 126                         $dbrAC2 = $lb->getConnection( DB_REPLICA, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 127                         $dbrAC2->ensureConnection();
 128                         $this->assertSharedHandle( $dbrAC2, $dbrAC1, "CONN_TRX_AUTOCOMMIT reuses connections" );
 129                 }
 130
 131                 $lb->closeAll( __METHOD__ );
 132         }
 133
 134         public function testWithReplica() {
 135                 // Simulate web request with DBO_TRX
 136                 $lb = $this->newMultiServerLocalLoadBalancer( [], [ 'flags' => DBO_TRX ] );
 137
 138                 $this->assertSame( 8, $lb->getServerCount() );
 139                 $this->assertTrue( $lb->hasReplicaServers() );
 140                 $this->assertTrue( $lb->hasStreamingReplicaServers() );
 141                 $this->assertSame( 'main-test-cluster', $lb->getClusterName() );
 142
 143                 for ( $i = 0; $i < $lb->getServerCount(); ++$i ) {
 144                         $this->assertIsString( $lb->getServerName( $i ) );
 145                         $this->assertIsArray( $lb->getServerInfo( $i ) );
 146                         $this->assertIsString( $lb->getServerType( $i ) );
 147                         $this->assertIsArray( $lb->getServerAttributes( $i ) );
 148                 }
 149
 150                 $dbw = $lb->getConnection( DB_PRIMARY );
 151                 $dbw->ensureConnection();
 152                 $wConn = TestingAccessWrapper::newFromObject( $dbw )->conn;
 153                 $wConnWrap = TestingAccessWrapper::newFromObject( $wConn );
 154
 155                 $this->assertTrue( $dbw->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on primary" );
 156                 $this->assertWriteAllowed( $dbw );
 157
 158                 $dbr = $lb->getConnection( DB_REPLICA );
 159                 $dbr->ensureConnection();
 160                 $rConn = TestingAccessWrapper::newFromObject( $dbr )->conn;
 161                 $rConnWrap = TestingAccessWrapper::newFromObject( $rConn );
 162
 163                 $this->assertTrue( $dbr->isReadOnly(), 'replica shows as replica' );
 164                 $this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on replica" );
 165                 $this->assertSame( $dbr->getLBInfo( 'serverIndex' ), $lb->getReaderIndex() );
 166
 167                 if ( !$lb->getServerAttributes( ServerInfo::WRITER_INDEX )[Database::ATTR_DB_LEVEL_LOCKING] ) {
 168                         $dbwAC1 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 169                         $this->assertFalse(
 170                                 $dbwAC1->getFlag( $dbw::DBO_TRX ),
 171                                 "No DBO_TRX with CONN_TRX_AUTOCOMMIT"
 172                         );
 173                         $this->assertTrue( $dbw->getFlag( $dbw::DBO_TRX ), "DBO_TRX still set on master" );
 174                         $this->assertUnsharedHandle( $dbw, $dbwAC1, "CONN_TRX_AUTOCOMMIT separate connection" );
 175
 176                         $dbrAC1 = $lb->getConnection( DB_REPLICA, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 177                         $this->assertFalse(
 178                                 $dbrAC1->getFlag( $dbw::DBO_TRX ),
 179                                 "No DBO_TRX with CONN_TRX_AUTOCOMMIT"
 180                         );
 181                         $this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX still set on replica" );
 182                         $this->assertUnsharedHandle( $dbr, $dbrAC1, "CONN_TRX_AUTOCOMMIT separate connection" );
 183
 184                         $dbwAC2 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 185                         $dbwAC2->ensureConnection();
 186                         $this->assertSharedHandle( $dbwAC2, $dbwAC1, "CONN_TRX_AUTOCOMMIT reuses connections" );
 187
 188                         $dbrAC2 = $lb->getConnection( DB_REPLICA, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 189                         $dbrAC2->ensureConnection();
 190                         $this->assertSharedHandle( $dbrAC2, $dbrAC1, "CONN_TRX_AUTOCOMMIT reuses connections" );
 191                 }
 192
 193                 $lb->closeAll( __METHOD__ );
 194         }
 195
 196         private function assertSharedHandle( DBConnRef $connRef1, DBConnRef $connRef2, $msg ) {
 197                 $connRef1Wrap = TestingAccessWrapper::newFromObject( $connRef1 );
 198                 $connRef2Wrap = TestingAccessWrapper::newFromObject( $connRef2 );
 199                 $this->assertSame( $connRef1Wrap->conn, $connRef2Wrap->conn, $msg );
 200         }
 201
 202         private function assertUnsharedHandle( DBConnRef $connRef1, DBConnRef $connRef2, $msg ) {
 203                 $connRef1Wrap = TestingAccessWrapper::newFromObject( $connRef1 );
 204                 $connRef2Wrap = TestingAccessWrapper::newFromObject( $connRef2 );
 205                 $this->assertNotSame( $connRef1Wrap->conn, $connRef2Wrap->conn, $msg );
 206         }
 207
 208         private function newSingleServerLocalLoadBalancer() {
 209                 global $wgDBname;
 210
 211                 return new LoadBalancer( [
 212                         'servers' => [ $this->makeServerConfig() ],
 213                         'localDomain' => new DatabaseDomain( $wgDBname, null, self::dbPrefix() ),
 214                         'cliMode' => false
 215                 ] );
 216         }
 217
 218         private function newMultiServerLocalLoadBalancer(
 219                 $lbExtra = [], $srvExtra = [], $masterOnly = false
 220         ) {
 221                 global $wgDBserver, $wgDBport, $wgDBuser, $wgDBpassword, $wgDBtype;
 222                 global $wgDBname, $wgDBmwschema;
 223                 global $wgSQLiteDataDir;
 224
 225                 $servers = [
 226                         // Primary DB
 227                         0 => $srvExtra + [
 228                                         'serverName' => 'db0',
 229                                         'host' => $wgDBserver,
 230                                         'port' => $wgDBport,
 231                                         'dbname' => $wgDBname,
 232                                         'tablePrefix' => self::dbPrefix(),
 233                                         'user' => $wgDBuser,
 234                                         'password' => $wgDBpassword,
 235                                         'type' => $wgDBtype,
 236                                         'dbDirectory' => $wgSQLiteDataDir,
 237                                         'load' => $masterOnly ? 100 : 0,
 238                                 ],
 239                         // Main replica DBs
 240                         1 => $srvExtra + [
 241                                         'serverName' => 'db1',
 242                                         'host' => $wgDBserver,
 243                                         'port' => $wgDBport,
 244                                         'dbname' => $wgDBname,
 245                                         'tablePrefix' => self::dbPrefix(),
 246                                         'user' => $wgDBuser,
 247                                         'password' => $wgDBpassword,
 248                                         'type' => $wgDBtype,
 249                                         'dbDirectory' => $wgSQLiteDataDir,
 250                                         'load' => $masterOnly ? 0 : 100,
 251                                 ],
 252                         2 => $srvExtra + [
 253                                         'serverName' => 'db2',
 254                                         'host' => $wgDBserver,
 255                                         'port' => $wgDBport,
 256                                         'dbname' => $wgDBname,
 257                                         'tablePrefix' => self::dbPrefix(),
 258                                         'user' => $wgDBuser,
 259                                         'password' => $wgDBpassword,
 260                                         'type' => $wgDBtype,
 261                                         'dbDirectory' => $wgSQLiteDataDir,
 262                                         'load' => $masterOnly ? 0 : 100,
 263                                 ],
 264                         // RC replica DBs
 265                         3 => $srvExtra + [
 266                                         'serverName' => 'db3',
 267                                         'host' => $wgDBserver,
 268                                         'port' => $wgDBport,
 269                                         'dbname' => $wgDBname,
 270                                         'tablePrefix' => self::dbPrefix(),
 271                                         'user' => $wgDBuser,
 272                                         'password' => $wgDBpassword,
 273                                         'type' => $wgDBtype,
 274                                         'dbDirectory' => $wgSQLiteDataDir,
 275                                         'load' => 0,
 276                                         'groupLoads' => [
 277                                                 'foo' => 100,
 278                                                 'bar' => 100
 279                                         ],
 280                                 ],
 281                         // Logging replica DBs
 282                         4 => $srvExtra + [
 283                                         'serverName' => 'db4',
 284                                         'host' => $wgDBserver,
 285                                         'port' => $wgDBport,
 286                                         'dbname' => $wgDBname,
 287                                         'tablePrefix' => self::dbPrefix(),
 288                                         'user' => $wgDBuser,
 289                                         'password' => $wgDBpassword,
 290                                         'type' => $wgDBtype,
 291                                         'dbDirectory' => $wgSQLiteDataDir,
 292                                         'load' => 0,
 293                                         'groupLoads' => [
 294                                                 'baz' => 100
 295                                         ],
 296                                 ],
 297                         5 => $srvExtra + [
 298                                         'serverName' => 'db5',
 299                                         'host' => $wgDBserver,
 300                                         'port' => $wgDBport,
 301                                         'dbname' => $wgDBname,
 302                                         'tablePrefix' => self::dbPrefix(),
 303                                         'user' => $wgDBuser,
 304                                         'password' => $wgDBpassword,
 305                                         'type' => $wgDBtype,
 306                                         'dbDirectory' => $wgSQLiteDataDir,
 307                                         'load' => 0,
 308                                         'groupLoads' => [
 309                                                 'baz' => 100
 310                                         ],
 311                                 ],
 312                         // Maintenance query replica DBs
 313                         6 => $srvExtra + [
 314                                         'serverName' => 'db6',
 315                                         'host' => $wgDBserver,
 316                                         'port' => $wgDBport,
 317                                         'dbname' => $wgDBname,
 318                                         'tablePrefix' => self::dbPrefix(),
 319                                         'user' => $wgDBuser,
 320                                         'password' => $wgDBpassword,
 321                                         'type' => $wgDBtype,
 322                                         'dbDirectory' => $wgSQLiteDataDir,
 323                                         'load' => 0,
 324                                         'groupLoads' => [
 325                                                 'vslow' => 100
 326                                         ],
 327                                 ],
 328                         // Replica DB that only has a copy of some static tables
 329                         7 => $srvExtra + [
 330                                         'serverName' => 'db7',
 331                                         'host' => $wgDBserver,
 332                                         'port' => $wgDBport,
 333                                         'dbname' => $wgDBname,
 334                                         'tablePrefix' => self::dbPrefix(),
 335                                         'user' => $wgDBuser,
 336                                         'password' => $wgDBpassword,
 337                                         'type' => $wgDBtype,
 338                                         'dbDirectory' => $wgSQLiteDataDir,
 339                                         'load' => 0,
 340                                         'groupLoads' => [
 341                                                 'archive' => 100
 342                                         ],
 343                                         'is static' => true
 344                                 ]
 345                 ];
 346
 347                 return new LoadBalancer( $lbExtra + [
 348                         'servers' => $servers,
 349                         'localDomain' => new DatabaseDomain( $wgDBname, $wgDBmwschema, self::dbPrefix() ),
 350                         'logger' => MediaWiki\Logger\LoggerFactory::getInstance( 'rdbms' ),
 351                         'loadMonitor' => [ 'class' => LoadMonitorNull::class ],
 352                         'clusterName' => 'main-test-cluster'
 353                 ] );
 354         }
 355
 356         private function assertWriteAllowed( IMaintainableDatabase $db ) {
 357                 $table = $db->tableName( 'some_table' );
 358                 // Trigger a transaction so that rollback() will remove all the tables.
 359                 // Don't do this for MySQL as it auto-commits transactions for DDL
 360                 // statements such as CREATE TABLE.
 361                 $useAtomicSection = in_array( $db->getType(), [ 'sqlite', 'postgres' ], true );
 362                 /** @var Database $db */
 363                 try {
 364                         $db->dropTable( 'some_table' );
 365                         $this->assertNotEquals( TransactionManager::STATUS_TRX_ERROR, $db->trxStatus() );
 366
 367                         if ( $useAtomicSection ) {
 368                                 $db->startAtomic( __METHOD__ );
 369                         }
 370                         // Use only basic SQL and trivial types for these queries for compatibility
 371                         $this->assertNotFalse(
 372                                 $db->query( "CREATE TABLE $table (id INT, time INT)", __METHOD__ ),
 373                                 "table created"
 374                         );
 375                         $this->assertNotEquals( TransactionManager::STATUS_TRX_ERROR, $db->trxStatus() );
 376                         $this->assertNotFalse(
 377                                 $db->query( "DELETE FROM $table WHERE id=57634126", __METHOD__ ),
 378                                 "delete query"
 379                         );
 380                         $this->assertNotEquals( TransactionManager::STATUS_TRX_ERROR, $db->trxStatus() );
 381                 } finally {
 382                         if ( !$useAtomicSection ) {
 383                                 // Drop the table to clean up, ignoring any error.
 384                                 $db->dropTable( 'some_table' );
 385                         }
 386                         // Rollback the atomic section for sqlite's benefit.
 387                         $db->rollback( __METHOD__, 'flush' );
 388                         $this->assertNotEquals( TransactionManager::STATUS_TRX_ERROR, $db->trxStatus() );
 389                 }
 390         }
 391
 392         public function testServerAttributes() {
 393                 $servers = [
 394                         [ // master
 395                                 'dbname' => 'my_unittest_wiki',
 396                                 'tablePrefix' => self::DB_PREFIX,
 397                                 'type' => 'sqlite',
 398                                 'dbDirectory' => "some_directory",
 399                                 'load' => 0
 400                         ]
 401                 ];
 402
 403                 $lb = new LoadBalancer( [
 404                         'servers' => $servers,
 405                         'localDomain' => new DatabaseDomain( 'my_unittest_wiki', null, self::DB_PREFIX ),
 406                         'loadMonitor' => [ 'class' => LoadMonitorNull::class ]
 407                 ] );
 408
 409                 $this->assertTrue( $lb->getServerAttributes( 0 )[Database::ATTR_DB_LEVEL_LOCKING] );
 410
 411                 $servers = [
 412                         [ // master
 413                                 'serverName' => 'db1',
 414                                 'host' => 'db1001',
 415                                 'user' => 'wikiuser',
 416                                 'password' => 'none',
 417                                 'dbname' => 'my_unittest_wiki',
 418                                 'tablePrefix' => self::DB_PREFIX,
 419                                 'type' => 'mysql',
 420                                 'load' => 100
 421                         ],
 422                         [ // emulated replica
 423                                 'serverName' => 'db2',
 424                                 'host' => 'db1002',
 425                                 'user' => 'wikiuser',
 426                                 'password' => 'none',
 427                                 'dbname' => 'my_unittest_wiki',
 428                                 'tablePrefix' => self::DB_PREFIX,
 429                                 'type' => 'mysql',
 430                                 'load' => 100
 431                         ]
 432                 ];
 433
 434                 $lb = new LoadBalancer( [
 435                         'servers' => $servers,
 436                         'localDomain' => new DatabaseDomain( 'my_unittest_wiki', null, self::DB_PREFIX ),
 437                         'loadMonitor' => [ 'class' => LoadMonitorNull::class ]
 438                 ] );
 439
 440                 $this->assertFalse( $lb->getServerAttributes( 1 )[Database::ATTR_DB_LEVEL_LOCKING] );
 441         }
 442
 443         public function testOpenConnection() {
 444                 $lb = $this->newSingleServerLocalLoadBalancer();
 445                 $i = ServerInfo::WRITER_INDEX;
 446
 447                 $this->assertFalse( $lb->getAnyOpenConnection( $i ) );
 448                 $this->assertFalse( $lb->getAnyOpenConnection( $i, $lb::CONN_TRX_AUTOCOMMIT ) );
 449
 450                 // Get two live round-aware handles
 451                 $raConnRef1 = $lb->getConnection( $i );
 452                 $raConnRef1->ensureConnection();
 453                 $raConnRef1Wrapper = TestingAccessWrapper::newFromObject( $raConnRef1 );
 454                 $raConnRef2 = $lb->getConnection( $i );
 455                 $raConnRef2->ensureConnection();
 456                 $raConnRef2Wrapper = TestingAccessWrapper::newFromObject( $raConnRef2 );
 457
 458                 $this->assertNotNull( $raConnRef1Wrapper->conn );
 459                 $this->assertSame( $raConnRef1Wrapper->conn, $raConnRef2Wrapper->conn );
 460                 $this->assertTrue( $raConnRef1Wrapper->conn->getFlag( DBO_TRX ) );
 461
 462                 // Get two live autocommit handles
 463                 $acConnRef1 = $lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 464                 $acConnRef1->ensureConnection();
 465                 $acConnRef1Wrapper = TestingAccessWrapper::newFromObject( $acConnRef1 );
 466                 $acConnRef2 = $lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 467                 $acConnRef2->ensureConnection();
 468                 $acConnRef2Wrapper = TestingAccessWrapper::newFromObject( $acConnRef2 );
 469
 470                 $this->assertNotNull( $acConnRef1Wrapper->conn );
 471                 $this->assertSame( $acConnRef1Wrapper->conn, $acConnRef2Wrapper->conn );
 472
 473                 $this->assertNotFalse( $lb->getAnyOpenConnection( $i ) );
 474
 475                 $lb->closeAll( __METHOD__ );
 476         }
 477
 478         public function testReconfigure() {
 479                 $serverA = $this->makeServerConfig();
 480                 $serverA['serverName'] = 'test_one';
 481
 482                 $serverB = $this->makeServerConfig();
 483                 $serverB['serverName'] = 'test_two';
 484                 $conf = [
 485                         'servers' => [ $serverA, $serverB ],
 486                         'clusterName' => 'A',
 487                         'localDomain' => $this->getDb()->getDomainID()
 488                 ];
 489
 490                 $lb = new LoadBalancer( $conf );
 491                 $this->assertSame( 2, $lb->getServerCount() );
 492
 493                 $con = $lb->getConnectionInternal( DB_PRIMARY );
 494                 $ref = $lb->getConnection( DB_PRIMARY );
 495
 496                 $this->assertTrue( $con->isOpen() );
 497                 $this->assertTrue( $ref->isOpen() );
 498
 499                 // Depool the second server
 500                 $conf['servers'] = [ $serverA ];
 501                 $lb->reconfigure( $conf );
 502                 $this->assertSame( 1, $lb->getServerCount() );
 503
 504                 // Reconfiguring should not close connections immediately.
 505                 $this->assertTrue( $con->isOpen() );
 506
 507                 // Connection refs should detect the config change, close the old connection,
 508                 // and get a new connection.
 509                 $this->assertTrue( $ref->isOpen() );
 510
 511                 // The old connection should have been called by DBConnRef.
 512                 $this->assertFalse( $con->isOpen() );
 513         }
 514
 515         public function testTransactionCallbackChains() {
 516                 global $wgDBserver, $wgDBport, $wgDBname, $wgDBuser, $wgDBpassword, $wgDBtype;
 517                 global $wgSQLiteDataDir;
 518
 519                 $servers = [
 520                         [
 521                                 'host' => $wgDBserver,
 522                                 'port' => $wgDBport,
 523                                 'dbname' => $wgDBname,
 524                                 'tablePrefix' => self::dbPrefix(),
 525                                 'user' => $wgDBuser,
 526                                 'password' => $wgDBpassword,
 527                                 'type' => $wgDBtype,
 528                                 'dbDirectory' => $wgSQLiteDataDir,
 529                                 'load' => 0,
 530                                 'flags' => DBO_TRX // simulate a web request with DBO_TRX
 531                         ],
 532                 ];
 533
 534                 $lb = new LoadBalancer( [
 535                         'servers' => $servers,
 536                         'localDomain' => new DatabaseDomain( $wgDBname, null, self::dbPrefix() )
 537                 ] );
 538                 /** @var LoadBalancer $lbWrapper */
 539                 $lbWrapper = TestingAccessWrapper::newFromObject( $lb );
 540
 541                 $conn1 = $lb->getConnection( ServerInfo::WRITER_INDEX, [], false );
 542                 $count = iterator_count( $lbWrapper->getOpenPrimaryConnections() );
 543                 $this->assertSame( 0, $count, 'Connection handle count' );
 544                 $conn1->getServerName();
 545                 $count = iterator_count( $lbWrapper->getOpenPrimaryConnections() );
 546                 $this->assertSame( 0, $count, 'Connection handle count' );
 547                 $conn1->ensureConnection();
 548
 549                 $conn2 = $lb->getConnection( ServerInfo::WRITER_INDEX, [], '' );
 550                 $count = iterator_count( $lbWrapper->getOpenPrimaryConnections() );
 551                 $this->assertSame( 1, $count, 'Connection handle count' );
 552                 $conn2->getServerName();
 553                 $count = iterator_count( $lbWrapper->getOpenPrimaryConnections() );
 554                 $this->assertSame( 1, $count, 'Connection handle count' );
 555                 $conn2->ensureConnection();
 556
 557                 $count = iterator_count( $lbWrapper->getOpenPrimaryConnections() );
 558                 $this->assertSame( 1, $count, 'Connection handle count' );
 559
 560                 $tlCalls = 0;
 561                 $lb->setTransactionListener( 'test-listener', static function () use ( &$tlCalls ) {
 562                         ++$tlCalls;
 563                 } );
 564
 565                 $lb->beginPrimaryChanges( __METHOD__ );
 566                 $bc = array_fill_keys( [ 'a', 'b', 'c', 'd' ], 0 );
 567                 $conn1->onTransactionPreCommitOrIdle( static function () use ( &$bc, $conn1, $conn2 ) {
 568                         $bc['a'] = 1;
 569                         $conn2->onTransactionPreCommitOrIdle( static function () use ( &$bc, $conn1 ) {
 570                                 $bc['b'] = 1;
 571                                 $conn1->onTransactionPreCommitOrIdle( static function () use ( &$bc, $conn1 ) {
 572                                         $bc['c'] = 1;
 573                                         $conn1->onTransactionPreCommitOrIdle( static function () use ( &$bc ) {
 574                                                 $bc['d'] = 1;
 575                                         } );
 576                                 } );
 577                         } );
 578                 } );
 579                 $lb->finalizePrimaryChanges();
 580                 $lb->approvePrimaryChanges( 0 );
 581                 $lb->commitPrimaryChanges( __METHOD__ );
 582                 $lb->runPrimaryTransactionIdleCallbacks();
 583                 $lb->runPrimaryTransactionListenerCallbacks();
 584
 585                 $this->assertSame( array_fill_keys( [ 'a', 'b', 'c', 'd' ], 1 ), $bc );
 586                 $this->assertSame( 1, $tlCalls );
 587
 588                 $tlCalls = 0;
 589                 $lb->beginPrimaryChanges( __METHOD__ );
 590                 $ac = array_fill_keys( [ 'a', 'b', 'c', 'd' ], 0 );
 591                 $conn1->onTransactionCommitOrIdle( static function () use ( &$ac, $conn1, $conn2 ) {
 592                         $ac['a'] = 1;
 593                         $conn2->onTransactionCommitOrIdle( static function () use ( &$ac, $conn1 ) {
 594                                 $ac['b'] = 1;
 595                                 $conn1->onTransactionCommitOrIdle( static function () use ( &$ac, $conn1 ) {
 596                                         $ac['c'] = 1;
 597                                         $conn1->onTransactionCommitOrIdle( static function () use ( &$ac ) {
 598                                                 $ac['d'] = 1;
 599                                         } );
 600                                 } );
 601                         } );
 602                 } );
 603                 $lb->finalizePrimaryChanges();
 604                 $lb->approvePrimaryChanges( 0 );
 605                 $lb->commitPrimaryChanges( __METHOD__ );
 606                 $lb->runPrimaryTransactionIdleCallbacks();
 607                 $lb->runPrimaryTransactionListenerCallbacks();
 608
 609                 $this->assertSame( array_fill_keys( [ 'a', 'b', 'c', 'd' ], 1 ), $ac );
 610                 $this->assertSame( 1, $tlCalls );
 611
 612                 $conn1->lock( 'test_lock_' . mt_rand(), __METHOD__, 0 );
 613                 $lb->flushPrimarySessions( __METHOD__ );
 614                 $this->assertSame( TransactionManager::STATUS_TRX_NONE, $conn1->trxStatus() );
 615                 $this->assertSame( TransactionManager::STATUS_TRX_NONE, $conn2->trxStatus() );
 616         }
 617
 618         public function testForbiddenWritesNoRef() {
 619                 // Simulate web request with DBO_TRX
 620                 $lb = $this->newMultiServerLocalLoadBalancer( [], [ 'flags' => DBO_TRX ] );
 621
 622                 $dbr = $lb->getConnection( DB_REPLICA );
 623                 $this->assertTrue( $dbr->isReadOnly(), 'replica shows as replica' );
 624                 $this->expectException( DBReadOnlyRoleError::class );
 625                 $dbr->newDeleteQueryBuilder()
 626                         ->deleteFrom( 'some_table' )
 627                         ->where( [ 'id' => 57634126 ] )
 628                         ->caller( __METHOD__ )
 629                         ->execute();
 630
 631                 // FIXME: not needed?
 632                 $lb->closeAll( __METHOD__ );
 633         }
 634
 635         public function testDBConnRefReadsMasterAndReplicaRoles() {
 636                 $lb = $this->newSingleServerLocalLoadBalancer();
 637
 638                 $rConn = $lb->getConnection( DB_REPLICA );
 639                 $wConn = $lb->getConnection( DB_PRIMARY );
 640                 $wConn2 = $lb->getConnection( 0 );
 641
 642                 $v = [ 'value' => '1', '1' ];
 643                 $sql = 'SELECT MAX(1) AS value';
 644                 foreach ( [ $rConn, $wConn, $wConn2 ] as $conn ) {
 645                         $conn->clearFlag( $conn::DBO_TRX );
 646
 647                         $res = $conn->query( $sql, __METHOD__ );
 648                         $this->assertEquals( $v, $res->fetchRow() );
 649
 650                         $res = $conn->query( $sql, __METHOD__, $conn::QUERY_REPLICA_ROLE );
 651                         $this->assertEquals( $v, $res->fetchRow() );
 652                 }
 653
 654                 $wConn->getScopedLockAndFlush( 'key', __METHOD__, 1 );
 655                 $wConn2->getScopedLockAndFlush( 'key2', __METHOD__, 1 );
 656         }
 657
 658         public function testDBConnRefWritesReplicaRole() {
 659                 $lb = $this->newSingleServerLocalLoadBalancer();
 660
 661                 $rConn = $lb->getConnection( DB_REPLICA );
 662
 663                 $this->expectException( DBReadOnlyRoleError::class );
 664                 $rConn->query( 'DELETE FROM sometesttable WHERE 1=0' );
 665         }
 666
 667         public function testDBConnRefWritesReplicaRoleIndex() {
 668                 $lb = $this->newMultiServerLocalLoadBalancer();
 669
 670                 $rConn = $lb->getConnection( 1 );
 671
 672                 $this->expectException( DBReadOnlyRoleError::class );
 673                 $rConn->query( 'DELETE FROM sometesttable WHERE 1=0' );
 674         }
 675
 676         public function testDBConnRefWritesReplicaRoleInsert() {
 677                 $lb = $this->newMultiServerLocalLoadBalancer();
 678
 679                 $rConn = $lb->getConnection( DB_REPLICA );
 680
 681                 $this->expectException( DBReadOnlyRoleError::class );
 682                 $rConn->insert( 'test', [ 't' => 1 ], __METHOD__ );
 683         }
 684
 685         public function testGetConnectionRefDefaultGroup() {
 686                 $lb = $this->newMultiServerLocalLoadBalancer( [ 'defaultGroup' => 'vslow' ] );
 687                 $lbWrapper = TestingAccessWrapper::newFromObject( $lb );
 688
 689                 $rVslow = $lb->getConnection( DB_REPLICA );
 690                 $vslowIndexPicked = $rVslow->getLBInfo( 'serverIndex' );
 691
 692                 $this->assertSame( $vslowIndexPicked, $lbWrapper->getExistingReaderIndex( 'vslow' ) );
 693         }
 694
 695         public function testGetConnectionRefUnknownDefaultGroup() {
 696                 $lb = $this->newMultiServerLocalLoadBalancer( [ 'defaultGroup' => 'invalid' ] );
 697
 698                 $this->assertInstanceOf(
 699                         IDatabase::class,
 700                         $lb->getConnection( DB_REPLICA )
 701                 );
 702         }
 703
 704         public function testQueryGroupIndex() {
 705                 $lb = $this->newMultiServerLocalLoadBalancer( [ 'defaultGroup' => false ] );
 706                 /** @var LoadBalancer $lbWrapper */
 707                 $lbWrapper = TestingAccessWrapper::newFromObject( $lb );
 708
 709                 $rGeneric = $lb->getConnection( DB_REPLICA );
 710                 $mainIndexPicked = $rGeneric->getLBInfo( 'serverIndex' );
 711
 712                 $this->assertSame(
 713                         $mainIndexPicked,
 714                         $lbWrapper->getExistingReaderIndex( $lb::GROUP_GENERIC )
 715                 );
 716                 $this->assertContains( $mainIndexPicked, [ 1, 2 ] );
 717                 for ( $i = 0; $i < 300; ++$i ) {
 718                         $rLog = $lb->getConnection( DB_REPLICA, [] );
 719                         $this->assertSame(
 720                                 $mainIndexPicked,
 721                                 $rLog->getLBInfo( 'serverIndex' ),
 722                                 "Main index unchanged" );
 723                 }
 724
 725                 $rRC = $lb->getConnection( DB_REPLICA, [ 'foo' ] );
 726                 $rWL = $lb->getConnection( DB_REPLICA, [ 'bar' ] );
 727                 $rRCMaint = $lb->getMaintenanceConnectionRef( DB_REPLICA, [ 'foo' ] );
 728                 $rWLMaint = $lb->getMaintenanceConnectionRef( DB_REPLICA, [ 'bar' ] );
 729
 730                 $this->assertSame( 3, $rRC->getLBInfo( 'serverIndex' ) );
 731                 $this->assertSame( 3, $rWL->getLBInfo( 'serverIndex' ) );
 732                 $this->assertSame( 3, $rRCMaint->getLBInfo( 'serverIndex' ) );
 733                 $this->assertSame( 3, $rWLMaint->getLBInfo( 'serverIndex' ) );
 734
 735                 $rLog = $lb->getConnection( DB_REPLICA, [ 'baz', 'bar' ] );
 736                 $logIndexPicked = $rLog->getLBInfo( 'serverIndex' );
 737
 738                 $this->assertSame( $logIndexPicked, $lbWrapper->getExistingReaderIndex( 'baz' ) );
 739                 $this->assertContains( $logIndexPicked, [ 4, 5 ] );
 740
 741                 for ( $i = 0; $i < 300; ++$i ) {
 742                         $rLog = $lb->getConnection( DB_REPLICA, [ 'baz', 'bar' ] );
 743                         $this->assertSame(
 744                                 $logIndexPicked, $rLog->getLBInfo( 'serverIndex' ), "Index unchanged" );
 745                 }
 746
 747                 $rVslow = $lb->getConnection( DB_REPLICA, [ 'vslow', 'baz' ] );
 748                 $vslowIndexPicked = $rVslow->getLBInfo( 'serverIndex' );
 749
 750                 $this->assertSame( $vslowIndexPicked, $lbWrapper->getExistingReaderIndex( 'vslow' ) );
 751                 $this->assertSame( 6, $vslowIndexPicked );
 752         }
 753
 754         public function testNonZeroMasterLoad() {
 755                 $lb = $this->newMultiServerLocalLoadBalancer( [], [ 'flags' => DBO_DEFAULT ], true );
 756                 // Make sure that no infinite loop occurs (T226678)
 757                 $rGeneric = $lb->getConnection( DB_REPLICA );
 758                 $this->assertSame( ServerInfo::WRITER_INDEX, $rGeneric->getLBInfo( 'serverIndex' ) );
 759         }
 760
 761         public function testSetDomainAliases() {
 762                 $lb = $this->newMultiServerLocalLoadBalancer();
 763                 $origDomain = $lb->getLocalDomainID();
 764
 765                 $this->assertSame( $origDomain, $lb->resolveDomainID( false ) );
 766                 $this->assertSame( "db-prefix_", $lb->resolveDomainID( "db-prefix_" ) );
 767
 768                 $lb->setDomainAliases( [
 769                         'alias-db' => 'realdb',
 770                         'alias-db-prefix_' => 'realdb-realprefix_'
 771                 ] );
 772
 773                 $this->assertSame( 'realdb', $lb->resolveDomainID( 'alias-db' ) );
 774                 $this->assertSame( "realdb-realprefix_", $lb->resolveDomainID( "alias-db-prefix_" ) );
 775         }
 776
 777         public function testClusterName() {
 778                 global $wgDBname;
 779                 $chronologyProtector = $this->createMock( ChronologyProtector::class );
 780                 $lb1 = new LoadBalancer( [
 781                         'servers' => [ $this->makeServerConfig() ],
 782                         'logger' => MediaWiki\Logger\LoggerFactory::getInstance( 'rdbms' ),
 783                         'localDomain' => new DatabaseDomain( $wgDBname, null, self::dbPrefix() ),
 784                         'chronologyProtector' => $chronologyProtector,
 785                         'clusterName' => 'xx'
 786                 ] );
 787                 $this->assertSame( 'xx', $lb1->getClusterName() );
 788
 789                 $lb2 = new LoadBalancer( [
 790                         'servers' => [ $this->makeServerConfig() ],
 791                         'logger' => MediaWiki\Logger\LoggerFactory::getInstance( 'rdbms' ),
 792                         'localDomain' => new DatabaseDomain( $wgDBname, null, self::dbPrefix() ),
 793                         'chronologyProtector' => $chronologyProtector,
 794                         'clusterName' => null
 795                 ] );
 796                 $this->assertSame( 'testhost', $lb2->getClusterName() );
 797         }
 798 }