MediaWiki  REL1_23
PoolCounterRedis.php
Go to the documentation of this file.
00001 <?php
00053 class PoolCounterRedis extends PoolCounter {
00055     protected $ring;
00057     protected $pool;
00059     protected $serversByLabel;
00061     protected $keySha1;
00063     protected $lockTTL;
00064 
00066     protected $conn;
00068     protected $slot;
00070     protected $onRelease;
00072     protected $session;
00074     protected $slotTime;
00075 
00076     const AWAKE_ONE = 1; // wake-up if when a slot can be taken from an existing process
00077     const AWAKE_ALL = 2; // wake-up if an existing process finishes and wake up such others
00078 
00080     protected static $active = null;
00081 
00082     function __construct( $conf, $type, $key ) {
00083         parent::__construct( $conf, $type, $key );
00084 
00085         $this->serversByLabel = $conf['servers'];
00086         $this->ring = new HashRing( array_fill_keys( array_keys( $conf['servers'] ), 100 ) );
00087 
00088         $conf['redisConfig']['serializer'] = 'none'; // for use with Lua
00089         $this->pool = RedisConnectionPool::singleton( $conf['redisConfig'] );
00090 
00091         $this->keySha1 = sha1( $this->key );
00092         $met = ini_get( 'max_execution_time' ); // usually 0 in CLI mode
00093         $this->lockTTL = $met ? 2*$met : 3600;
00094 
00095         if ( self::$active === null ) {
00096             self::$active = array();
00097             register_shutdown_function( array( __CLASS__, 'releaseAll' ) );
00098         }
00099     }
00100 
00104     protected function getConnection() {
00105         if ( !isset( $this->conn ) ) {
00106             $conn = false;
00107             $servers = $this->ring->getLocations( $this->key, 3 );
00108             ArrayUtils::consistentHashSort( $servers, $this->key );
00109             foreach ( $servers as $server ) {
00110                 $conn = $this->pool->getConnection( $this->serversByLabel[$server] );
00111                 if ( $conn ) {
00112                     break;
00113                 }
00114             }
00115             if ( !$conn ) {
00116                 return Status::newFatal( 'pool-servererror', implode( ', ', $servers ) );
00117             }
00118             $this->conn = $conn;
00119         }
00120         return Status::newGood( $this->conn );
00121     }
00122 
00123     function acquireForMe() {
00124         $section = new ProfileSection( __METHOD__ );
00125 
00126         return $this->waitForSlotOrNotif( self::AWAKE_ONE );
00127     }
00128 
00129     function acquireForAnyone() {
00130         $section = new ProfileSection( __METHOD__ );
00131 
00132         return $this->waitForSlotOrNotif( self::AWAKE_ALL );
00133     }
00134 
00135     function release() {
00136         $section = new ProfileSection( __METHOD__ );
00137 
00138         if ( $this->slot === null ) {
00139             return Status::newGood( PoolCounter::NOT_LOCKED ); // not locked
00140         }
00141 
00142         $status = $this->getConnection();
00143         if ( !$status->isOK() ) {
00144             return $status;
00145         }
00146         $conn = $status->value;
00147 
00148         static $script =
00149 <<<LUA
00150         local kSlots,kSlotsNextRelease,kWakeup,kWaiting = unpack(KEYS)
00151         local rMaxWorkers,rExpiry,rSlot,rSlotTime,rAwakeAll,rTime = unpack(ARGV)
00152         -- Add the slots back to the list (if rSlot is "w" then it is not a slot).
00153         -- Treat the list as expired if the "next release" time sorted-set is missing.
00154         if rSlot ~= 'w' and redis.call('exists',kSlotsNextRelease) == 1 then
00155             if 1*redis.call('zScore',kSlotsNextRelease,rSlot) ~= (rSlotTime + rExpiry) then
00156                 -- Slot lock expired and was released already
00157             elseif redis.call('lLen',kSlots) >= (1*rMaxWorkers - 1) then
00158                 -- Clear list to save space; it will re-init as needed
00159                 redis.call('del',kSlots,kSlotsNextRelease)
00160             else
00161                 -- Add slot back to pool and update the "next release" time
00162                 redis.call('rPush',kSlots,rSlot)
00163                 redis.call('zAdd',kSlotsNextRelease,rTime + 30,rSlot)
00164                 -- Always keep renewing the expiry on use
00165                 redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
00166                 redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
00167             end
00168         end
00169         -- Update an ephemeral list to wake up other clients that can
00170         -- reuse any cached work from this process. Only do this if no
00171         -- slots are currently free (e.g. clients could be waiting).
00172         if 1*rAwakeAll == 1 then
00173             local count = redis.call('zCard',kWaiting)
00174             for i = 1,count do
00175                 redis.call('rPush',kWakeup,'w')
00176             end
00177             redis.call('pexpire',kWakeup,1)
00178         end
00179         return 1
00180 LUA;
00181         try {
00182             $res = $conn->luaEval( $script,
00183                 array(
00184                     $this->getSlotListKey(),
00185                     $this->getSlotRTimeSetKey(),
00186                     $this->getWakeupListKey(),
00187                     $this->getWaitSetKey(),
00188                     $this->workers,
00189                     $this->lockTTL,
00190                     $this->slot,
00191                     $this->slotTime, // used for CAS-style sanity check
00192                     ( $this->onRelease === self::AWAKE_ALL ) ? 1 : 0,
00193                     microtime( true )
00194                 ),
00195                 4 # number of first argument(s) that are keys
00196             );
00197         } catch ( RedisException $e ) {
00198             return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
00199         }
00200 
00201         $this->slot = null;
00202         $this->slotTime = null;
00203         $this->onRelease = null;
00204         unset( self::$active[$this->session] );
00205 
00206         return Status::newGood( PoolCounter::RELEASED );
00207     }
00208 
00213     protected function waitForSlotOrNotif( $doWakeup ) {
00214         if ( $this->slot !== null ) {
00215             return Status::newGood( PoolCounter::LOCK_HELD ); // already acquired
00216         }
00217 
00218         $status = $this->getConnection();
00219         if ( !$status->isOK() ) {
00220             return $status;
00221         }
00222         $conn = $status->value;
00223 
00224         $now = microtime( true );
00225         try {
00226             $slot = $this->initAndPopPoolSlotList( $conn, $now );
00227             if ( ctype_digit( $slot ) ) {
00228                 // Pool slot acquired by this process
00229                 $slotTime = $now;
00230             } elseif ( $slot === 'QUEUE_FULL' ) {
00231                 // Too many processes are waiting for pooled processes to finish
00232                 return Status::newGood( PoolCounter::QUEUE_FULL );
00233             } elseif ( $slot === 'QUEUE_WAIT' ) {
00234                 // This process is now registered as waiting
00235                 $keys = ( $doWakeup == self::AWAKE_ALL )
00236                     // Wait for an open slot or wake-up signal (preferring the later)
00237                     ? array( $this->getWakeupListKey(), $this->getSlotListKey() )
00238                     // Just wait for an actual pool slot
00239                     : array( $this->getSlotListKey() );
00240 
00241                 $res = $conn->blPop( $keys, $this->timeout );
00242                 if ( $res === array() ) {
00243                     $conn->zRem( $this->getWaitSetKey(), $this->session ); // no longer waiting
00244                     return Status::newGood( PoolCounter::TIMEOUT );
00245                 }
00246 
00247                 $slot = $res[1]; // pool slot or "w" for wake-up notifications
00248                 $slotTime = microtime( true ); // last microtime() was a few RTTs ago
00249                 // Unregister this process as waiting and bump slot "next release" time
00250                 $this->registerAcquisitionTime( $conn, $slot, $slotTime );
00251             } else {
00252                 return Status::newFatal( 'pool-error-unknown', "Server gave slot '$slot'." );
00253             }
00254         } catch ( RedisException $e ) {
00255             return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
00256         }
00257 
00258         if ( $slot !== 'w' ) {
00259             $this->slot = $slot;
00260             $this->slotTime = $slotTime;
00261             $this->onRelease = $doWakeup;
00262             self::$active[$this->session] = $this;
00263         }
00264 
00265         return Status::newGood( $slot === 'w' ? PoolCounter::DONE : PoolCounter::LOCKED );
00266     }
00267 
00273     protected function initAndPopPoolSlotList( RedisConnRef $conn, $now ) {
00274         static $script =
00275 <<<LUA
00276         local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
00277         local rMaxWorkers,rMaxQueue,rTimeout,rExpiry,rSess,rTime = unpack(ARGV)
00278         -- Initialize if the "next release" time sorted-set is empty. The slot key
00279         -- itself is empty if all slots are busy or when nothing is initialized.
00280         -- If the list is empty but the set is not, then it is the later case.
00281         -- For sanity, if the list exists but not the set, then reset everything.
00282         if redis.call('exists',kSlotsNextRelease) == 0 then
00283             redis.call('del',kSlots)
00284             for i = 1,1*rMaxWorkers do
00285                 redis.call('rPush',kSlots,i)
00286                 redis.call('zAdd',kSlotsNextRelease,-1,i)
00287             end
00288         -- Otherwise do maintenance to clean up after network partitions
00289         else
00290             -- Find stale slot locks and add free them (avoid duplicates for sanity)
00291             local staleLocks = redis.call('zRangeByScore',kSlotsNextRelease,0,rTime)
00292             for k,slot in ipairs(staleLocks) do
00293                 redis.call('lRem',kSlots,0,slot)
00294                 redis.call('rPush',kSlots,slot)
00295                 redis.call('zAdd',kSlotsNextRelease,rTime + 30,slot)
00296             end
00297             -- Find stale wait slot entries and remove them
00298             redis.call('zRemRangeByScore',kSlotWaits,0,rTime - 2*rTimeout)
00299         end
00300         local slot
00301         -- Try to acquire a slot if possible now
00302         if redis.call('lLen',kSlots) > 0 then
00303             slot = redis.call('lPop',kSlots)
00304             -- Update the slot "next release" time
00305             redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,slot)
00306         elseif redis.call('zCard',kSlotWaits) >= 1*rMaxQueue then
00307             slot = 'QUEUE_FULL'
00308         else
00309             slot = 'QUEUE_WAIT'
00310             -- Register this process as waiting
00311             redis.call('zAdd',kSlotWaits,rTime,rSess)
00312             redis.call('expireAt',kSlotWaits,math.ceil(rTime + 2*rTimeout))
00313         end
00314         -- Always keep renewing the expiry on use
00315         redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
00316         redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
00317         return slot
00318 LUA;
00319         return $conn->luaEval( $script,
00320             array(
00321                 $this->getSlotListKey(),
00322                 $this->getSlotRTimeSetKey(),
00323                 $this->getWaitSetKey(),
00324                 $this->workers,
00325                 $this->maxqueue,
00326                 $this->timeout,
00327                 $this->lockTTL,
00328                 $this->session,
00329                 $now
00330             ),
00331             3 # number of first argument(s) that are keys
00332         );
00333     }
00334 
00341     protected function registerAcquisitionTime( RedisConnRef $conn, $slot, $now ) {
00342         static $script =
00343 <<<LUA
00344         local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
00345         local rSlot,rExpiry,rSess,rTime = unpack(ARGV)
00346         -- If rSlot is 'w' then the client was told to wake up but got no slot
00347         if rSlot ~= 'w' then
00348             -- Update the slot "next release" time
00349             redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,rSlot)
00350             -- Always keep renewing the expiry on use
00351             redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
00352             redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
00353         end
00354         -- Unregister this process as waiting
00355         redis.call('zRem',kSlotWaits,rSess)
00356         return 1
00357 LUA;
00358         return $conn->luaEval( $script,
00359             array(
00360                 $this->getSlotListKey(),
00361                 $this->getSlotRTimeSetKey(),
00362                 $this->getWaitSetKey(),
00363                 $slot,
00364                 $this->lockTTL,
00365                 $this->session,
00366                 $now
00367             ),
00368             3 # number of first argument(s) that are keys
00369         );
00370     }
00371 
00375     protected function getSlotListKey() {
00376         return "poolcounter:l-slots-{$this->keySha1}-{$this->workers}";
00377     }
00378 
00382     protected function getSlotRTimeSetKey() {
00383         return "poolcounter:z-renewtime-{$this->keySha1}-{$this->workers}";
00384     }
00385 
00389     protected function getWaitSetKey() {
00390         return "poolcounter:z-wait-{$this->keySha1}-{$this->workers}";
00391     }
00392 
00396     protected function getWakeupListKey() {
00397         return "poolcounter:l-wakeup-{$this->keySha1}-{$this->workers}";
00398     }
00399 
00403     public static function releaseAll() {
00404         foreach ( self::$active as $poolCounter ) {
00405             try {
00406                 if ( $poolCounter->slot !== null ) {
00407                     $poolCounter->release();
00408                 }
00409             } catch ( Exception $e ) {}
00410         }
00411     }
00412 }