Handle multiple warnings correctly in ApiBase::setWarning(). Calling this function...
[lhc/web/wiklou.git] / includes / LoadBalancer.php
index 69b5cd3..4280713 100644 (file)
@@ -1,81 +1,66 @@
 <?php
 /**
- *
- */
-
-/**
- * Depends on the database object
+ * @file
+ * @ingroup Database
  */
-require_once( 'Database.php' );
-
-# Valid database indexes
-# Operation-based indexes
-define( 'DB_SLAVE', -1 );     # Read from the slave (or only server)
-define( 'DB_MASTER', -2 );    # Write to master (or only server)
-define( 'DB_LAST', -3 );     # Whatever database was used last
-
-# Obsolete aliases
-define( 'DB_READ', -1 );
-define( 'DB_WRITE', -2 );
-
-# Task-based indexes
-# ***NOT USED YET, EXPERIMENTAL***
-# These may be defined in $wgDBservers. If they aren't, the default reader or writer will be used
-# Even numbers are always readers, odd numbers are writers
-define( 'DB_TASK_FIRST', 1000 );  # First in list
-define( 'DB_SEARCH_R', 1000 );    # Search read
-define( 'DB_SEARCH_W', 1001 );    # Search write
-define( 'DB_ASKSQL_R', 1002 );    # Special:Asksql read
-define( 'DB_WATCHLIST_R', 1004 ); # Watchlist read
-define( 'DB_TASK_LAST', 1004) ;   # Last in list
-
-define( 'MASTER_WAIT_TIMEOUT', 15 ); # Time to wait for a slave to synchronise
 
 /**
  * Database load balancing object
+ *
  * @todo document
+ * @ingroup Database
  */
 class LoadBalancer {
-       /* private */ var $mServers, $mConnections, $mLoads;
-       /* private */ var $mFailFunction;
-       /* private */ var $mForce, $mReadIndex, $mLastIndex;
-       /* private */ var $mWaitForFile, $mWaitForPos;
-
-       function LoadBalancer()
-       {
-               $this->mServers = array();
-               $this->mConnections = array();
-               $this->mFailFunction = false;
-               $this->mReadIndex = -1;
-               $this->mForce = -1;
-               $this->mLastIndex = -1;
-       }
-
-       function newFromParams( $servers, $failFunction = false )
-       {
-               $lb = new LoadBalancer;
-               $lb->initialise( $servers, $failFunction = false );
-               return $lb;
-       }
-
-       function initialise( $servers, $failFunction = false )
+       /* private */ var $mServers, $mConns, $mLoads, $mGroupLoads;
+       /* private */ var $mFailFunction, $mErrorConnection;
+       /* private */ var $mReadIndex, $mLastIndex, $mAllowLagged;
+       /* private */ var $mWaitForPos, $mWaitTimeout;
+       /* private */ var $mLaggedSlaveMode, $mLastError = 'Unknown error';
+       /* private */ var $mParentInfo, $mLagTimes;
+
+       function __construct( $servers, $failFunction = false, $waitTimeout = 10, $unused = false )
        {
                $this->mServers = $servers;
                $this->mFailFunction = $failFunction;
                $this->mReadIndex = -1;
                $this->mWriteIndex = -1;
-               $this->mForce = -1;
-               $this->mConnections = array();
-               $this->mLastIndex = 1;
+               $this->mConns = array(
+                       'local' => array(),
+                       'foreignUsed' => array(),
+                       'foreignFree' => array() );
+               $this->mLastIndex = -1;
                $this->mLoads = array();
-               $this->mWaitForFile = false;
                $this->mWaitForPos = false;
+               $this->mWaitTimeout = $waitTimeout;
+               $this->mLaggedSlaveMode = false;
+               $this->mErrorConnection = false;
+               $this->mAllowLag = false;
 
                foreach( $servers as $i => $server ) {
                        $this->mLoads[$i] = $server['load'];
+                       if ( isset( $server['groupLoads'] ) ) {
+                               foreach ( $server['groupLoads'] as $group => $ratio ) {
+                                       if ( !isset( $this->mGroupLoads[$group] ) ) {
+                                               $this->mGroupLoads[$group] = array();
+                                       }
+                                       $this->mGroupLoads[$group][$i] = $ratio;
+                               }
+                       }
                }
        }
-       
+
+       static function newFromParams( $servers, $failFunction = false, $waitTimeout = 10 )
+       {
+               return new LoadBalancer( $servers, $failFunction, $waitTimeout );
+       }
+
+       /**
+        * Get or set arbitrary data used by the parent object, usually an LBFactory
+        */
+       function parentInfo( $x = null ) {
+               return wfSetVar( $this->mParentInfo, $x );
+       }
+
        /**
         * Given an array of non-normalised probabilities, this function will select
         * an element and return the appropriate key
@@ -86,13 +71,17 @@ class LoadBalancer {
                        return false;
                }
 
-               $sum = 0;
-               foreach ( $weights as $w ) {
-                       $sum += $w;
+               $sum = array_sum( $weights );
+               if ( $sum == 0 ) {
+                       # No loads on any of them
+                       # In previous versions, this triggered an unweighted random selection,
+                       # but this feature has been removed as of April 2006 to allow for strict
+                       # separation of query groups.
+                       return false;
                }
                $max = mt_getrandmax();
                $rand = mt_rand(0, $max) / $max * $sum;
-               
+
                $sum = 0;
                foreach ( $weights as $i => $w ) {
                        $sum += $w;
@@ -103,131 +92,309 @@ class LoadBalancer {
                return $i;
        }
 
-       function getReaderIndex()
-       {
-               $fname = 'LoadBalancer::getReaderIndex';
-               wfProfileIn( $fname );
+       function getRandomNonLagged( $loads, $wiki = false ) {
+               # Unset excessively lagged servers
+               $lags = $this->getLagTimes( $wiki );
+               foreach ( $lags as $i => $lag ) {
+                       if ( $i != 0 && isset( $this->mServers[$i]['max lag'] ) ) {
+                               if ( $lag === false ) {
+                                       wfDebug( "Server #$i is not replicating\n" );
+                                       unset( $loads[$i] );
+                               } elseif ( $lag > $this->mServers[$i]['max lag'] ) {
+                                       wfDebug( "Server #$i is excessively lagged ($lag seconds)\n" );
+                                       unset( $loads[$i] );
+                               }
+                       }
+               }
 
-               $i = false;
-               if ( $this->mForce >= 0 ) {
-                       $i = $this->mForce;
-               } else {
-                       if ( $this->mReadIndex >= 0 ) {
-                               $i = $this->mReadIndex;
+               # Find out if all the slaves with non-zero load are lagged
+               $sum = 0;
+               foreach ( $loads as $load ) {
+                       $sum += $load;
+               }
+               if ( $sum == 0 ) {
+                       # No appropriate DB servers except maybe the master and some slaves with zero load
+                       # Do NOT use the master
+                       # Instead, this function will return false, triggering read-only mode,
+                       # and a lagged slave will be used instead.
+                       return false;
+               }
+
+               if ( count( $loads ) == 0 ) {
+                       return false;
+               }
+
+               #wfDebugLog( 'connect', var_export( $loads, true ) );
+
+               # Return a random representative of the remainder
+               return $this->pickRandom( $loads );
+       }
+
+       /**
+        * Get the index of the reader connection, which may be a slave
+        * This takes into account load ratios and lag times. It should
+        * always return a consistent index during a given invocation
+        *
+        * Side effect: opens connections to databases
+        */
+       function getReaderIndex( $group = false, $wiki = false ) {
+               global $wgReadOnly, $wgDBClusterTimeout, $wgDBAvgStatusPoll, $wgDBtype;
+
+               # FIXME: For now, only go through all this for mysql databases
+               if ($wgDBtype != 'mysql') {
+                       return $this->getWriterIndex();
+               }
+
+               if ( count( $this->mServers ) == 1 )  {
+                       # Skip the load balancing if there's only one server
+                       return 0;
+               } elseif ( $group === false and $this->mReadIndex >= 0 ) {
+                       # Shortcut if generic reader exists already
+                       return $this->mReadIndex;
+               }
+
+               wfProfileIn( __METHOD__ );
+
+               $totalElapsed = 0;
+
+               # convert from seconds to microseconds
+               $timeout = $wgDBClusterTimeout * 1e6;
+
+               # Find the relevant load array
+               if ( $group !== false ) {
+                       if ( isset( $this->mGroupLoads[$group] ) ) {
+                               $nonErrorLoads = $this->mGroupLoads[$group];
                        } else {
-                               # $loads is $this->mLoads except with elements knocked out if they
-                               # don't work
-                               $loads = $this->mLoads;
-                               do {
-                                       $i = $this->pickRandom( $loads );
-                                       if ( $i !== false ) {
-                                               wfDebug( "Using reader #$i: {$this->mServers[$i]['host']}\n" );
-
-                                               $this->openConnection( $i );
-
-                                               if ( !$this->isOpen( $i ) ) {
-                                                       unset( $loads[$i] );
-                                               }
+                               # No loads for this group, return false and the caller can use some other group
+                               wfDebug( __METHOD__.": no loads for group $group\n" );
+                               wfProfileOut( __METHOD__ );
+                               return false;
+                       }
+               } else {
+                       $nonErrorLoads = $this->mLoads;
+               }
+
+               if ( !$nonErrorLoads ) {
+                       throw new MWException( "Empty server array given to LoadBalancer" );
+               }
+
+               $i = false;
+               $found = false;
+               $laggedSlaveMode = false;
+
+               # First try quickly looking through the available servers for a server that
+               # meets our criteria
+               do {
+                       $totalThreadsConnected = 0;
+                       $overloadedServers = 0;
+                       $currentLoads = $nonErrorLoads;
+                       while ( count( $currentLoads ) ) {
+                               if ( $wgReadOnly || $this->mAllowLagged || $laggedSlaveMode ) {
+                                       $i = $this->pickRandom( $currentLoads );
+                               } else {
+                                       $i = $this->getRandomNonLagged( $currentLoads, $wiki );
+                                       if ( $i === false && count( $currentLoads ) != 0 )  {
+                                               # All slaves lagged. Switch to read-only mode
+                                               $wgReadOnly = wfMsgNoDBForContent( 'readonly_lag' );
+                                               $i = $this->pickRandom( $currentLoads );
+                                               $laggedSlaveMode = true;
                                        }
-                               } while ( $i !== false && !$this->isOpen( $i ) );
+                               }
+
+                               if ( $i === false ) {
+                                       # pickRandom() returned false
+                                       # This is permanent and means the configuration wants us to return false
+                                       wfDebugLog( 'connect', __METHOD__.": pickRandom() returned false\n" );
+                                       wfProfileOut( __METHOD__ );
+                                       return false;
+                               }
+
+                               wfDebugLog( 'connect', __METHOD__.": Using reader #$i: {$this->mServers[$i]['host']}...\n" );
+                               $conn = $this->openConnection( $i, $wiki );
+
+                               if ( !$conn ) {
+                                       wfDebugLog( 'connect', __METHOD__.": Failed connecting to $i/$wiki\n" );
+                                       unset( $nonErrorLoads[$i] );
+                                       unset( $currentLoads[$i] );
+                                       continue;
+                               }
 
-                               if ( $this->isOpen( $i ) ) {
-                                       $this->mReadIndex = $i;
+                               if ( isset( $this->mServers[$i]['max threads'] ) ) {
+                                       $status = $conn->getStatus("Thread%");
+                                       if ( $wiki !== false ) {
+                                               $this->reuseConnection( $conn );
+                                       }
+                                       if ( $status['Threads_running'] > $this->mServers[$i]['max threads'] ) {
+                                               $totalThreadsConnected += $status['Threads_connected'];
+                                               $overloadedServers++;
+                                               unset( $currentLoads[$i] );
+                                       } else {
+                                               # Max threads satisfied, return this server
+                                               break 2;
+                                       }
                                } else {
-                                       $i = false;
+                                       # No maximum, return this server
+                                       if ( $wiki !== false ) {
+                                               $this->reuseConnection( $conn );
+                                       }
+                                       $found = true;
+                                       break 2;
                                }
                        }
+
+                       # No server found yet
+                       $i = false;
+
+                       # If all servers were down, quit now
+                       if ( !count( $nonErrorLoads ) ) {
+                               wfDebugLog( 'connect', "All servers down\n" );
+                               break;
+                       }
+
+                       # Some servers must have been overloaded
+                       if ( $overloadedServers == 0 ) {
+                               throw new MWException( __METHOD__.": unexpectedly found no overloaded servers" );
+                       }
+                       # Back off for a while
+                       # Scale the sleep time by the number of connected threads, to produce a
+                       # roughly constant global poll rate
+                       $avgThreads = $totalThreadsConnected / $overloadedServers;
+                       $totalElapsed += $this->sleep( $wgDBAvgStatusPoll * $avgThreads );
+               } while ( $totalElapsed < $timeout );
+
+               if ( $totalElapsed >= $timeout ) {
+                       wfDebugLog( 'connect', "All servers busy\n" );
+                       $this->mErrorConnection = false;
+                       $this->mLastError = 'All servers busy';
+               }
+
+               if ( $i !== false ) {
+                       # Wait for the session master pos for a short time
+                       if ( $this->mWaitForPos && $i > 0 ) {
+                               if ( !$this->doWait( $i ) ) {
+                                       $this->mServers[$i]['slave pos'] = $conn->getSlavePos();
+                               }
+                       }
+                       if ( $this->mReadIndex <=0 && $this->mLoads[$i]>0 && $i !== false ) {
+                               $this->mReadIndex = $i;
+                       }
                }
-               wfProfileOut( $fname );
+               wfProfileOut( __METHOD__ );
                return $i;
        }
 
+       /**
+        * Wait for a specified number of microseconds, and return the period waited
+        */
+       function sleep( $t ) {
+               wfProfileIn( __METHOD__ );
+               wfDebug( __METHOD__.": waiting $t us\n" );
+               usleep( $t );
+               wfProfileOut( __METHOD__ );
+               return $t;
+       }
+
+       /**
+        * Get a random server to use in a query group
+        * @deprecated use getReaderIndex
+        */
+       function getGroupIndex( $group ) {
+               return $this->getReaderIndex( $group );
+       }
+
        /**
         * Set the master wait position
         * If a DB_SLAVE connection has been opened already, waits
         * Otherwise sets a variable telling it to wait if such a connection is opened
         */
-       function waitFor( $file, $pos ) {
-               $fname = 'LoadBalancer::waitFor';
-               wfProfileIn( $fname );
-
-               wfDebug( "User master pos: $file $pos\n" );
-               $this->mWaitForFile = false;
-               $this->mWaitForPos = false;
-
-               if ( count( $this->mServers ) > 1 ) {
-                       $this->mWaitForFile = $file;
-                       $this->mWaitForPos = $pos;
+       public function waitFor( $pos ) {
+               wfProfileIn( __METHOD__ );
+               $this->mWaitForPos = $pos;
+               $i = $this->mReadIndex;
+
+               if ( $i > 0 ) {
+                       if ( !$this->doWait( $i ) ) {
+                               $this->mServers[$i]['slave pos'] = $this->getAnyOpenConnection( $i )->getSlavePos();
+                               $this->mLaggedSlaveMode = true;
+                       }
+               }
+               wfProfileOut( __METHOD__ );
+       }
 
-                       if ( $this->mReadIndex > 0 ) {
-                               if ( !$this->doWait( $this->mReadIndex ) ) {
-                                       # Use master instead
-                                       $this->mReadIndex = 0;
-                               }
-                       } 
+       /**
+        * Get any open connection to a given server index, local or foreign
+        * Returns false if there is no connection open
+        */
+       function getAnyOpenConnection( $i ) {
+               foreach ( $this->mConns as $type => $conns ) {
+                       if ( !empty( $conns[$i] ) ) {
+                               return reset( $conns[$i] );
+                       }
                }
-               wfProfileOut( $fname );
+               return false;
        }
 
        /**
         * Wait for a given slave to catch up to the master pos stored in $this
         */
        function doWait( $index ) {
-               global $wgMemc;
-               
-               $retVal = false;
-
-               $key = 'masterpos:' . $index;
-               $memcPos = $wgMemc->get( $key );
-               if ( $memcPos ) {
-                       list( $file, $pos ) = explode( ' ', $memcPos );
-                       # If the saved position is later than the requested position, return now
-                       if ( $file == $this->mWaitForFile && $this->mWaitForPos <= $pos ) {
-                               $retVal = true;
-                       }
+               # Find a connection to wait on
+               $conn = $this->getAnyOpenConnection( $index );
+               if ( !$conn ) {
+                       wfDebug( __METHOD__ . ": no connection open\n" );
+                       return false;
                }
 
-               if ( !$retVal && $this->isOpen( $index ) ) {
-                       $conn =& $this->mConnections( $index );
-                       wfDebug( "Waiting for slave #$index to catch up...\n" );
-                       $result = $conn->masterPosWait( $this->mWaitForFile, $this->mWaitForPos, MASTER_WAIT_TIMEOUT );
+               wfDebug( __METHOD__.": Waiting for slave #$index to catch up...\n" );
+               $result = $conn->masterPosWait( $this->mWaitForPos, $this->mWaitTimeout );
 
-                       if ( $result == -1 || is_null( $result ) ) {
-                               # Timed out waiting for slave, use master instead
-                               wfDebug( "Timed out waiting for slave #$index pos {$this->mWaitForFile} {$this->mWaitForPos}\n" );
-                               $retVal = false;
-                       } else {
-                               $retVal = true;
-                               wfDebug( "Done\n" );
-                       }
+               if ( $result == -1 || is_null( $result ) ) {
+                       # Timed out waiting for slave, use master instead
+                       wfDebug( __METHOD__.": Timed out waiting for slave #$index pos {$this->mWaitForPos}\n" );
+                       return false;
+               } else {
+                       wfDebug( __METHOD__.": Done\n" );
+                       return true;
                }
-               return $retVal;
-       }               
+       }
 
        /**
         * Get a connection by index
+        * This is the main entry point for this class.
         */
-       function &getConnection( $i, $fail = true )
-       {
-               $fname = 'LoadBalancer::getConnection';
-               wfProfileIn( $fname );
-               /*
-               # Task-based index
-               if ( $i >= DB_TASK_FIRST && $i < DB_TASK_LAST ) {
-                       if ( $i % 2 ) {
-                               # Odd index use writer
-                               $i = DB_MASTER;
-                       } else {
-                               # Even index use reader
-                               $i = DB_SLAVE;
+       public function &getConnection( $i, $groups = array(), $wiki = false ) {
+               global $wgDBtype;
+               wfProfileIn( __METHOD__ );
+
+               if ( $wiki === wfWikiID() ) {
+                       $wiki = false;
+               }
+
+               # Query groups
+               if ( $i == DB_MASTER ) {
+                       $i = $this->getWriterIndex();
+               } elseif ( !is_array( $groups ) ) {
+                       $groupIndex = $this->getReaderIndex( $groups, $wiki );
+                       if ( $groupIndex !== false ) {
+                               $serverName = $this->getServerName( $groupIndex );
+                               wfDebug( __METHOD__.": using server $serverName for group $groups\n" );
+                               $i = $groupIndex;
+                       }
+               } else {
+                       foreach ( $groups as $group ) {
+                               $groupIndex = $this->getReaderIndex( $group, $wiki );
+                               if ( $groupIndex !== false ) {
+                                       $serverName = $this->getServerName( $groupIndex );
+                                       wfDebug( __METHOD__.": using server $serverName for group $group\n" );
+                                       $i = $groupIndex;
+                                       break;
+                               }
                        }
-               }*/
+               }
 
                # Operation-based index
-               if ( $i == DB_SLAVE ) { 
-                       $i = $this->getReaderIndex();
-               } elseif ( $i == DB_MASTER ) {
-                       $i = $this->getWriterIndex();
+               if ( $i == DB_SLAVE ) {
+                       $i = $this->getReaderIndex( false, $wiki );
                } elseif ( $i == DB_LAST ) {
                        # Just use $this->mLastIndex, which should already be set
                        $i = $this->mLastIndex;
@@ -237,116 +404,270 @@ class LoadBalancer {
                                $i = $this->getWriterIndex();
                        }
                }
+               # Couldn't find a working server in getReaderIndex()?
+               if ( $i === false ) {
+                       $this->reportConnectionError( $this->mErrorConnection );
+               }
+
                # Now we have an explicit index into the servers array
-               $this->openConnection( $i, $fail );
-               wfProfileOut( $fname );
-               return $this->mConnections[$i];
+               $conn = $this->openConnection( $i, $wiki );
+               if ( !$conn ) {
+                       $this->reportConnectionError( $this->mErrorConnection );
+               }
+
+               wfProfileOut( __METHOD__ );
+               return $conn;
+       }
+
+       /**
+        * Mark a foreign connection as being available for reuse under a different
+        * DB name or prefix. This mechanism is reference-counted, and must be called
+        * the same number of times as getConnection() to work.
+        */
+       public function reuseConnection( $conn ) {
+               $serverIndex = $conn->getLBInfo('serverIndex');
+               $refCount = $conn->getLBInfo('foreignPoolRefCount');
+               $dbName = $conn->getDBname();
+               $prefix = $conn->tablePrefix();
+               if ( strval( $prefix ) !== '' ) {
+                       $wiki = "$dbName-$prefix";
+               } else {
+                       $wiki = $dbName;
+               }
+               if ( $serverIndex === null || $refCount === null ) {
+                       wfDebug( __METHOD__.": this connection was not opened as a foreign connection\n" );
+                       /**
+                        * This can happen in code like:
+                        *   foreach ( $dbs as $db ) {
+                        *     $conn = $lb->getConnection( DB_SLAVE, array(), $db );
+                        *     ...
+                        *     $lb->reuseConnection( $conn );
+                        *   }
+                        * When a connection to the local DB is opened in this way, reuseConnection()
+                        * should be ignored
+                        */
+                       return;
+               }
+               if ( $this->mConns['foreignUsed'][$serverIndex][$wiki] !== $conn ) {
+                       throw new MWException( __METHOD__.": connection not found, has the connection been freed already?" );
+               }
+               $conn->setLBInfo( 'foreignPoolRefCount', --$refCount );
+               if ( $refCount <= 0 ) {
+                       $this->mConns['foreignFree'][$serverIndex][$wiki] = $conn;
+                       unset( $this->mConns['foreignUsed'][$serverIndex][$wiki] );
+                       wfDebug( __METHOD__.": freed connection $serverIndex/$wiki\n" );
+               } else {
+                       wfDebug( __METHOD__.": reference count for $serverIndex/$wiki reduced to $refCount\n" );
+               }
        }
 
        /**
         * Open a connection to the server given by the specified index
-        * Index must be an actual index into the array
-        * @private
+        * Index must be an actual index into the array.
+        * If the server is already open, returns it.
+        *
+        * On error, returns false, and the connection which caused the
+        * error will be available via $this->mErrorConnection.
+        *
+        * @param integer $i Server index
+        * @param string $wiki Wiki ID to open
+        * @return Database
+        *
+        * @access private
         */
-       function openConnection( $i, $fail = false ) {
-               $fname = 'LoadBalancer::openConnection';
-               wfProfileIn( $fname );
-
-               if ( !$this->isOpen( $i ) ) {
-                       $this->mConnections[$i] = $this->reallyOpenConnection( $this->mServers[$i] );
-                       
-                       if ( $i != 0 && $this->mWaitForFile ) {
-                               if ( !$this->doWait( $i ) ) {
-                                       # Error waiting for this slave, use master instead
-                                       $this->mReadIndex = 0;
-                                       $i = 0;
-                                       if ( !$this->isOpen( 0 ) ) {
-                                               $this->mConnections[0] = $this->reallyOpenConnection( $this->mServers[0] );
-                                       }
-                                       wfDebug( "Failed over to {$this->mConnections[0]->mServer}\n" );
-                               }
-                       }
+       function openConnection( $i, $wiki = false ) {
+               wfProfileIn( __METHOD__ );
+               if ( $wiki !== false ) {
+                       $conn = $this->openForeignConnection( $i, $wiki );
+                       wfProfileOut( __METHOD__);
+                       return $conn;
                }
-               if ( !$this->isOpen( $i ) ) {
-                       wfDebug( "Failed to connect to database $i at {$this->mServers[$i]['host']}\n" );
-                       if ( $fail ) {
-                               $this->reportConnectionError( $this->mConnections[$i] );
+               if ( isset( $this->mConns['local'][$i][0] ) ) {
+                       $conn = $this->mConns['local'][$i][0];
+               } else {
+                       $server = $this->mServers[$i];
+                       $server['serverIndex'] = $i;
+                       $conn = $this->reallyOpenConnection( $server );
+                       if ( $conn->isOpen() ) {
+                               $this->mConns['local'][$i][0] = $conn;
+                       } else {
+                               wfDebug( "Failed to connect to database $i at {$this->mServers[$i]['host']}\n" );
+                               $this->mErrorConnection = $conn;
+                               $conn = false;
                        }
-                       $this->mConnections[$i] = false;
                }
                $this->mLastIndex = $i;
-               wfProfileOut( $fname );
+               wfProfileOut( __METHOD__ );
+               return $conn;
+       }
+
+       /**
+        * Open a connection to a foreign DB, or return one if it is already open.
+        *
+        * Increments a reference count on the returned connection which locks the
+        * connection to the requested wiki. This reference count can be
+        * decremented by calling reuseConnection().
+        *
+        * If a connection is open to the appropriate server already, but with the wrong
+        * database, it will be switched to the right database and returned, as long as
+        * it has been freed first with reuseConnection().
+        *
+        * On error, returns false, and the connection which caused the
+        * error will be available via $this->mErrorConnection.
+        *
+        * @param integer $i Server index
+        * @param string $wiki Wiki ID to open
+        * @return Database
+        */
+       function openForeignConnection( $i, $wiki ) {
+               wfProfileIn(__METHOD__);
+               list( $dbName, $prefix ) = wfSplitWikiID( $wiki );
+               if ( isset( $this->mConns['foreignUsed'][$i][$wiki] ) ) {
+                       // Reuse an already-used connection
+                       $conn = $this->mConns['foreignUsed'][$i][$wiki];
+                       wfDebug( __METHOD__.": reusing connection $i/$wiki\n" );
+               } elseif ( isset( $this->mConns['foreignFree'][$i][$wiki] ) ) {
+                       // Reuse a free connection for the same wiki
+                       $conn = $this->mConns['foreignFree'][$i][$wiki];
+                       unset( $this->mConns['foreignFree'][$i][$wiki] );
+                       $this->mConns['foreignUsed'][$i][$wiki] = $conn;
+                       wfDebug( __METHOD__.": reusing free connection $i/$wiki\n" );
+               } elseif ( !empty( $this->mConns['foreignFree'][$i] ) ) {
+                       // Reuse a connection from another wiki
+                       $conn = reset( $this->mConns['foreignFree'][$i] );
+                       $oldWiki = key( $this->mConns['foreignFree'][$i] );
+
+                       if ( !$conn->selectDB( $dbName ) ) {
+                               global $wguname;
+                               $this->mLastError = "Error selecting database $dbName on server " .
+                                       $conn->getServer() . " from client host {$wguname['nodename']}\n";
+                               $this->mErrorConnection = $conn;
+                               $conn = false;
+                       } else {
+                               $conn->tablePrefix( $prefix );
+                               unset( $this->mConns['foreignFree'][$i][$oldWiki] );
+                               $this->mConns['foreignUsed'][$i][$wiki] = $conn;
+                               wfDebug( __METHOD__.": reusing free connection from $oldWiki for $wiki\n" );
+                       }
+               } else {
+                       // Open a new connection
+                       $server = $this->mServers[$i];
+                       $server['serverIndex'] = $i;
+                       $server['foreignPoolRefCount'] = 0;
+                       $conn = $this->reallyOpenConnection( $server, $dbName );
+                       if ( !$conn->isOpen() ) {
+                               wfDebug( __METHOD__.": error opening connection for $i/$wiki\n" );
+                               $this->mErrorConnection = $conn;
+                               $conn = false;
+                       } else {
+                               $this->mConns['foreignUsed'][$i][$wiki] = $conn;
+                               wfDebug( __METHOD__.": opened new connection for $i/$wiki\n" );
+                       }
+               }
+
+               // Increment reference count
+               if ( $conn ) {
+                       $refCount = $conn->getLBInfo( 'foreignPoolRefCount' );
+                       $conn->setLBInfo( 'foreignPoolRefCount', $refCount + 1 );
+               }
+               wfProfileOut(__METHOD__);
+               return $conn;
        }
 
        /**
         * Test if the specified index represents an open connection
-        * @private
+        * @access private
         */
        function isOpen( $index ) {
                if( !is_integer( $index ) ) {
                        return false;
                }
-               if ( array_key_exists( $index, $this->mConnections ) && is_object( $this->mConnections[$index] ) && 
-                 $this->mConnections[$index]->isOpen() ) 
-               {
-                       return true;
-               } else {
-                       return false;
-               }
+               return (bool)$this->getAnyOpenConnection( $index );
        }
-       
+
        /**
-        * Really opens a connection
-        * @private
+        * Really opens a connection. Uncached.
+        * Returns a Database object whether or not the connection was successful.
+        * @access private
         */
-       function reallyOpenConnection( &$server ) {
-                       extract( $server );
-                       # Get class for this database type
-                       $class = 'Database' . ucfirst( $type );
-                       if ( !class_exists( $class ) ) {
-                               require_once( "$class.php" );
-                       }
+       function reallyOpenConnection( $server, $dbNameOverride = false ) {
+               if( !is_array( $server ) ) {
+                       throw new MWException( 'You must update your load-balancing configuration. See DefaultSettings.php entry for $wgDBservers.' );
+               }
+
+               extract( $server );
+               if ( $dbNameOverride !== false ) {
+                       $dbname = $dbNameOverride;
+               }
+
+               # Get class for this database type
+               $class = 'Database' . ucfirst( $type );
 
-                       # Create object
-                       return new $class( $host, $user, $password, $dbname, 1, $flags );
+               # Create object
+               wfDebug( "Connecting to $host $dbname...\n" );
+               $db = new $class( $host, $user, $password, $dbname, 1, $flags );
+               if ( $db->isOpen() ) {
+                       wfDebug( "Connected\n" );
+               } else {
+                       wfDebug( "Failed\n" );
+               }
+               $db->setLBInfo( $server );
+               if ( isset( $server['fakeSlaveLag'] ) ) {
+                       $db->setFakeSlaveLag( $server['fakeSlaveLag'] );
+               }
+               if ( isset( $server['fakeMaster'] ) ) {
+                       $db->setFakeMaster( true );
+               }
+               return $db;
        }
-       
-       function reportConnectionError( &$conn )
-       {
-               $fname = 'LoadBalancer::reportConnectionError';
-               wfProfileIn( $fname );
+
+       function reportConnectionError( &$conn ) {
+               wfProfileIn( __METHOD__ );
                # Prevent infinite recursion
-               
+
                static $reporting = false;
                if ( !$reporting ) {
                        $reporting = true;
                        if ( !is_object( $conn ) ) {
+                               // No last connection, probably due to all servers being too busy
                                $conn = new Database;
-                       }
-                       if ( $this->mFailFunction ) {
-                               $conn->failFunction( $this->mFailFunction );
+                               if ( $this->mFailFunction ) {
+                                       $conn->failFunction( $this->mFailFunction );
+                                       $conn->reportConnectionError( $this->mLastError );
+                               } else {
+                                       // If all servers were busy, mLastError will contain something sensible
+                                       throw new DBConnectionError( $conn, $this->mLastError );
+                               }
                        } else {
-                               $conn->failFunction( 'wfEmergencyAbort' );
+                               if ( $this->mFailFunction ) {
+                                       $conn->failFunction( $this->mFailFunction );
+                               } else {
+                                       $conn->failFunction( false );
+                               }
+                               $server = $conn->getProperty( 'mServer' );
+                               $conn->reportConnectionError( "{$this->mLastError} ({$server})" );
                        }
-                       $conn->reportConnectionError();
                        $reporting = false;
                }
-               wfProfileOut( $fname );
+               wfProfileOut( __METHOD__ );
        }
-       
-       function getWriterIndex()
-       {
+
+       function getWriterIndex() {
                return 0;
        }
 
-       function force( $i )
-       {
-               $this->mForce = $i;
+       /**
+        * Returns true if the specified index is a valid server index
+        */
+       function haveIndex( $i ) {
+               return array_key_exists( $i, $this->mServers );
        }
 
-       function haveIndex( $i )
-       {
-               return array_key_exists( $i, $this->mServers );
+       /**
+        * Returns true if the specified index is valid and has non-zero load
+        */
+       function isNonZeroLoad( $i ) {
+               return array_key_exists( $i, $this->mServers ) && $this->mLoads[$i] != 0;
        }
 
        /**
@@ -357,54 +678,217 @@ class LoadBalancer {
        }
 
        /**
-        * Save master pos to the session and to memcached, if the session exists
+        * Get the host name or IP address of the server with the specified index
         */
-       function saveMasterPos() {
-               global $wgSessionStarted;
-               if ( $wgSessionStarted && count( $this->mServers ) > 1 ) {
-                       # If this entire request was served from a slave without opening a connection to the 
-                       # master (however unlikely that may be), then we can fetch the position from the slave.
-                       if ( empty( $this->mConnections[0] ) ) {
-                               $conn =& $this->getConnection( DB_SLAVE );
-                               list( $file, $pos ) = $conn->getSlavePos();
-                               wfDebug( "Saving master pos fetched from slave: $file $pos\n" );
-                       } else {
-                               $conn =& $this->getConnection( 0 );
-                               list( $file, $pos ) = $conn->getMasterPos();
-                               wfDebug( "Saving master pos: $file $pos\n" );
-                       }
-                       if ( $file !== false ) {
-                               $_SESSION['master_log_file'] = $file;
-                               $_SESSION['master_pos'] = $pos;
-                       }
+       function getServerName( $i ) {
+               if ( isset( $this->mServers[$i]['hostName'] ) ) {
+                       return $this->mServers[$i]['hostName'];
+               } elseif ( isset( $this->mServers[$i]['host'] ) ) {
+                       return $this->mServers[$i]['host'];
+               } else {
+                       return '';
                }
        }
 
        /**
-        * Loads the master pos from the session, waits for it if necessary
+        * Get the current master position for chronology control purposes
+        * @return mixed
         */
-       function loadMasterPos() {
-               if ( isset( $_SESSION['master_log_file'] ) && isset( $_SESSION['master_pos'] ) ) {
-                       $this->waitFor( $_SESSION['master_log_file'], $_SESSION['master_pos'] );
+       function getMasterPos() {
+               # If this entire request was served from a slave without opening a connection to the
+               # master (however unlikely that may be), then we can fetch the position from the slave.
+               $masterConn = $this->getAnyOpenConnection( 0 );
+               if ( !$masterConn ) {
+                       for ( $i = 1; $i < count( $this->mServers ); $i++ ) {
+                               $conn = $this->getAnyOpenConnection( $i );
+                               if ( $conn ) {
+                                       wfDebug( "Master pos fetched from slave\n" );
+                                       return $conn->getSlavePos();
+                               }
+                       }
+               } else {
+                       wfDebug( "Master pos fetched from master\n" );
+                       return $masterConn->getMasterPos();
                }
+               return false;
        }
 
        /**
         * Close all open connections
         */
        function closeAll() {
-               foreach( $this->mConnections as $i => $conn ) {
-                       if ( $this->isOpen( $i ) ) {
-                               $conn->close();
+               foreach ( $this->mConns as $conns2 ) {
+                       foreach  ( $conns2 as $conns3 ) {
+                               foreach ( $conns3 as $conn ) {
+                                       $conn->close();
+                               }
                        }
                }
+               $this->mConns = array(
+                       'local' => array(),
+                       'foreignFree' => array(),
+                       'foreignUsed' => array(),
+               );
        }
 
+       /**
+        * Close a connection
+        * Using this function makes sure the LoadBalancer knows the connection is closed.
+        * If you use $conn->close() directly, the load balancer won't update its state.
+        */
+       function closeConnecton( $conn ) {
+               $done = false;
+               foreach ( $this->mConns as $i1 => $conns2 ) {
+                       foreach ( $conns2 as $i2 => $conns3 ) {
+                               foreach ( $conns3 as $i3 => $candidateConn ) {
+                                       if ( $conn === $candidateConn ) {
+                                               $conn->close();
+                                               unset( $this->mConns[$i1][$i2][$i3] );
+                                               $done = true;
+                                               break;
+                                       }
+                               }
+                       }
+               }
+               if ( !$done ) {
+                       $conn->close();
+               }
+       }
+
+       /**
+        * Commit transactions on all open connections
+        */
        function commitAll() {
-               foreach( $this->mConnections as $i => $conn ) {
-                       if ( $this->isOpen( $i ) ) {
-                               $conn->immediateCommit();
+               foreach ( $this->mConns as $conns2 ) {
+                       foreach ( $conns2 as $conns3 ) {
+                               foreach ( $conns3 as $conn ) {
+                                       $conn->immediateCommit();
+                               }
                        }
                }
        }
+
+       /* Issue COMMIT only on master, only if queries were done on connection */
+       function commitMasterChanges() {
+               // Always 0, but who knows.. :)
+               $masterIndex = $this->getWriterIndex();
+               foreach ( $this->mConns as $type => $conns2 ) {
+                       if ( empty( $conns2[$masterIndex] ) ) {
+                               continue;
+                       }
+                       foreach ( $conns2[$masterIndex] as $conn ) {
+                               if ( $conn->lastQuery() != '' ) {
+                                       $conn->commit();
+                               }
+                       }
+               }
+       }
+
+       function waitTimeout( $value = NULL ) {
+               return wfSetVar( $this->mWaitTimeout, $value );
+       }
+
+       function getLaggedSlaveMode() {
+               return $this->mLaggedSlaveMode;
+       }
+
+       /* Disables/enables lag checks */
+       function allowLagged($mode=null) {
+               if ($mode===null)
+                       return $this->mAllowLagged;
+               $this->mAllowLagged=$mode;
+       }
+
+       function pingAll() {
+               $success = true;
+               foreach ( $this->mConns as $conns2 ) {
+                       foreach ( $conns2 as $conns3 ) {
+                               foreach ( $conns3 as $conn ) {
+                                       if ( !$conn->ping() ) {
+                                               $success = false;
+                                       }
+                               }
+                       }
+               }
+               return $success;
+       }
+
+       /**
+        * Get the hostname and lag time of the most-lagged slave.
+        * This is useful for maintenance scripts that need to throttle their updates.
+        * May attempt to open connections to slaves on the default DB.
+        */
+       function getMaxLag() {
+               $maxLag = -1;
+               $host = '';
+               foreach ( $this->mServers as $i => $conn ) {
+                       $conn = $this->getAnyOpenConnection( $i );
+                       if ( !$conn ) {
+                               $conn = $this->openConnection( $i );
+                       }
+                       if ( !$conn ) {
+                               continue;
+                       }
+                       $lag = $conn->getLag();
+                       if ( $lag > $maxLag ) {
+                               $maxLag = $lag;
+                               $host = $this->mServers[$i]['host'];
+                       }
+               }
+               return array( $host, $maxLag );
+       }
+
+       /**
+        * Get lag time for each server
+        * Results are cached for a short time in memcached, and indefinitely in the process cache
+        */
+       function getLagTimes( $wiki = false ) {
+               wfProfileIn( __METHOD__ );
+
+               if ( !isset( $this->mLagTimes ) ) {
+                       $expiry = 5;
+                       $requestRate = 10;
+
+                       global $wgMemc;
+                       $masterName = $this->getServerName( 0 );
+                       $memcKey = wfMemcKey( 'lag_times', $masterName );
+                       $times = $wgMemc->get( $memcKey );
+                       if ( $times ) {
+                               # Randomly recache with probability rising over $expiry
+                               $elapsed = time() - $times['timestamp'];
+                               $chance = max( 0, ( $expiry - $elapsed ) * $requestRate );
+                               if ( mt_rand( 0, $chance ) != 0 ) {
+                                       unset( $times['timestamp'] );
+                                       wfProfileOut( __METHOD__ );
+                                       return $times;
+                               }
+                               wfIncrStats( 'lag_cache_miss_expired' );
+                       } else {
+                               wfIncrStats( 'lag_cache_miss_absent' );
+                       }
+
+                       # Cache key missing or expired
+
+                       $times = array();
+                       foreach ( $this->mServers as $i => $conn ) {
+                               if ($i == 0) { # Master
+                                       $times[$i] = 0;
+                               } elseif ( false !== ( $conn = $this->getAnyOpenConnection( $i ) ) ) {
+                                       $times[$i] = $conn->getLag();
+                               } elseif ( false !== ( $conn = $this->openConnection( $i, $wiki ) ) ) {
+                                       $times[$i] = $conn->getLag();
+                               }
+                       }
+
+                       # Add a timestamp key so we know when it was cached
+                       $times['timestamp'] = time();
+                       $wgMemc->set( $memcKey, $times, $expiry );
+
+                       # But don't give the timestamp to the caller
+                       unset($times['timestamp']);
+                       $this->mLagTimes = $times;
+               }
+               wfProfileOut( __METHOD__ );
+               return $this->mLagTimes;
+       }
 }