includes/poolcounter/PoolCounterRedis.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 use Psr\Log\LoggerInterface;
  21
  22 /**
  23  * Version of PoolCounter that uses Redis
  24  *
  25  * There are four main redis keys used to track each pool counter key:
  26  *   - poolcounter:l-slots-*     : A list of available slot IDs for a pool.
  27  *   - poolcounter:z-renewtime-* : A sorted set of (slot ID, UNIX timestamp as score)
  28  *                                 used for tracking the next time a slot should be
  29  *                                 released. This is -1 when a slot is created, and is
  30  *                                 set when released (expired), locked, and unlocked.
  31  *   - poolcounter:z-wait-*      : A sorted set of (slot ID, UNIX timestamp as score)
  32  *                                 used for tracking waiting processes (and wait time).
  33  *   - poolcounter:l-wakeup-*    : A list pushed to for the sake of waking up processes
  34  *                                 when a any process in the pool finishes (lasts for 1ms).
  35  * For a given pool key, all the redis keys start off non-existing and are deleted if not
  36  * used for a while to prevent garbage from building up on the server. They are atomically
  37  * re-initialized as needed. The "z-renewtime" key is used for detecting sessions which got
  38  * slots but then disappeared. Stale entries from there have their timestamp updated and the
  39  * corresponding slots freed up. The "z-wait" key is used for detecting processes registered
  40  * as waiting but that disappeared. Stale entries from there are deleted and the corresponding
  41  * slots are freed up. The worker count is included in all the redis key names as it does not
  42  * vary within each $wgPoolCounterConf type and doing so handles configuration changes.
  43  *
  44  * This class requires Redis 2.6 as it makes use Lua scripts for fast atomic operations.
  45  * Also this should be on a server plenty of RAM for the working set to avoid evictions.
  46  * Evictions could temporarily allow wait queues to double in size or temporarily cause
  47  * pools to appear as full when they are not. Using volatile-ttl and bumping memory-samples
  48  * in redis.conf can be helpful otherwise.
  49  *
  50  * @ingroup Redis
  51  * @since 1.23
  52  */
  53 class PoolCounterRedis extends PoolCounter {
  54         /** @var HashRing */
  55         protected $ring;
  56         /** @var RedisConnectionPool */
  57         protected $pool;
  58         /** @var LoggerInterface */
  59         protected $logger;
  60         /** @var array (server label => host) map */
  61         protected $serversByLabel;
  62         /** @var string SHA-1 of the key */
  63         protected $keySha1;
  64         /** @var int TTL for locks to expire (work should finish in this time) */
  65         protected $lockTTL;
  66
  67         /** @var RedisConnRef */
  68         protected $conn;
  69         /** @var string Pool slot value */
  70         protected $slot;
  71         /** @var int AWAKE_* constant */
  72         protected $onRelease;
  73         /** @var string Unique string to identify this process */
  74         protected $session;
  75         /** @var int UNIX timestamp */
  76         protected $slotTime;
  77
  78         const AWAKE_ONE = 1; // wake-up if when a slot can be taken from an existing process
  79         const AWAKE_ALL = 2; // wake-up if an existing process finishes and wake up such others
  80
  81         /** @var PoolCounterRedis[] List of active PoolCounterRedis objects in this script */
  82         protected static $active = null;
  83
  84         function __construct( $conf, $type, $key ) {
  85                 parent::__construct( $conf, $type, $key );
  86
  87                 $this->serversByLabel = $conf['servers'];
  88
  89                 $serverLabels = array_keys( $conf['servers'] );
  90                 $this->ring = new HashRing( array_fill_keys( $serverLabels, 10 ) );
  91
  92                 $conf['redisConfig']['serializer'] = 'none'; // for use with Lua
  93                 $this->pool = RedisConnectionPool::singleton( $conf['redisConfig'] );
  94                 $this->logger = \MediaWiki\Logger\LoggerFactory::getInstance( 'redis' );
  95
  96                 $this->keySha1 = sha1( $this->key );
  97                 $met = ini_get( 'max_execution_time' ); // usually 0 in CLI mode
  98                 $this->lockTTL = $met ? 2 * $met : 3600;
  99
 100                 if ( self::$active === null ) {
 101                         self::$active = [];
 102                         register_shutdown_function( [ __CLASS__, 'releaseAll' ] );
 103                 }
 104         }
 105
 106         /**
 107          * @return Status Uses RediConnRef as value on success
 108          */
 109         protected function getConnection() {
 110                 if ( !isset( $this->conn ) ) {
 111                         $conn = false;
 112                         $servers = $this->ring->getLocations( $this->key, 3 );
 113                         ArrayUtils::consistentHashSort( $servers, $this->key );
 114                         foreach ( $servers as $server ) {
 115                                 $conn = $this->pool->getConnection( $this->serversByLabel[$server], $this->logger );
 116                                 if ( $conn ) {
 117                                         break;
 118                                 }
 119                         }
 120                         if ( !$conn ) {
 121                                 return Status::newFatal( 'pool-servererror', implode( ', ', $servers ) );
 122                         }
 123                         $this->conn = $conn;
 124                 }
 125                 return Status::newGood( $this->conn );
 126         }
 127
 128         function acquireForMe() {
 129                 $status = $this->precheckAcquire();
 130                 if ( !$status->isGood() ) {
 131                         return $status;
 132                 }
 133
 134                 return $this->waitForSlotOrNotif( self::AWAKE_ONE );
 135         }
 136
 137         function acquireForAnyone() {
 138                 $status = $this->precheckAcquire();
 139                 if ( !$status->isGood() ) {
 140                         return $status;
 141                 }
 142
 143                 return $this->waitForSlotOrNotif( self::AWAKE_ALL );
 144         }
 145
 146         function release() {
 147                 if ( $this->slot === null ) {
 148                         return Status::newGood( PoolCounter::NOT_LOCKED ); // not locked
 149                 }
 150
 151                 $status = $this->getConnection();
 152                 if ( !$status->isOK() ) {
 153                         return $status;
 154                 }
 155                 /** @var RedisConnRef $conn */
 156                 $conn = $status->value;
 157                 '@phan-var RedisConnRef $conn';
 158
 159                 // phpcs:disable Generic.Files.LineLength
 160                 static $script =
 161                 /** @lang Lua */
 162 <<<LUA
 163                 local kSlots,kSlotsNextRelease,kWakeup,kWaiting = unpack(KEYS)
 164                 local rMaxWorkers,rExpiry,rSlot,rSlotTime,rAwakeAll,rTime = unpack(ARGV)
 165                 -- Add the slots back to the list (if rSlot is "w" then it is not a slot).
 166                 -- Treat the list as expired if the "next release" time sorted-set is missing.
 167                 if rSlot ~= 'w' and redis.call('exists',kSlotsNextRelease) == 1 then
 168                         if 1*redis.call('zScore',kSlotsNextRelease,rSlot) ~= (rSlotTime + rExpiry) then
 169                                 -- Slot lock expired and was released already
 170                         elseif redis.call('lLen',kSlots) >= 1*rMaxWorkers then
 171                                 -- Slots somehow got out of sync; reset the list for sanity
 172                                 redis.call('del',kSlots,kSlotsNextRelease)
 173                         elseif redis.call('lLen',kSlots) == (1*rMaxWorkers - 1) and redis.call('zCard',kWaiting) == 0 then
 174                                 -- Slot list will be made full; clear it to save space (it re-inits as needed)
 175                                 -- since nothing is waiting on being unblocked by a push to the list
 176                                 redis.call('del',kSlots,kSlotsNextRelease)
 177                         else
 178                                 -- Add slot back to pool and update the "next release" time
 179                                 redis.call('rPush',kSlots,rSlot)
 180                                 redis.call('zAdd',kSlotsNextRelease,rTime + 30,rSlot)
 181                                 -- Always keep renewing the expiry on use
 182                                 redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
 183                                 redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
 184                         end
 185                 end
 186                 -- Update an ephemeral list to wake up other clients that can
 187                 -- reuse any cached work from this process. Only do this if no
 188                 -- slots are currently free (e.g. clients could be waiting).
 189                 if 1*rAwakeAll == 1 then
 190                         local count = redis.call('zCard',kWaiting)
 191                         for i = 1,count do
 192                                 redis.call('rPush',kWakeup,'w')
 193                         end
 194                         redis.call('pexpire',kWakeup,1)
 195                 end
 196                 return 1
 197 LUA;
 198                 // phpcs:enable
 199
 200                 try {
 201                         $conn->luaEval( $script,
 202                                 [
 203                                         $this->getSlotListKey(),
 204                                         $this->getSlotRTimeSetKey(),
 205                                         $this->getWakeupListKey(),
 206                                         $this->getWaitSetKey(),
 207                                         $this->workers,
 208                                         $this->lockTTL,
 209                                         $this->slot,
 210                                         $this->slotTime, // used for CAS-style sanity check
 211                                         ( $this->onRelease === self::AWAKE_ALL ) ? 1 : 0,
 212                                         microtime( true )
 213                                 ],
 214                                 4 # number of first argument(s) that are keys
 215                         );
 216                 } catch ( RedisException $e ) {
 217                         return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
 218                 }
 219
 220                 $this->slot = null;
 221                 $this->slotTime = null;
 222                 $this->onRelease = null;
 223                 unset( self::$active[$this->session] );
 224
 225                 $this->onRelease();
 226
 227                 return Status::newGood( PoolCounter::RELEASED );
 228         }
 229
 230         /**
 231          * @param int $doWakeup AWAKE_* constant
 232          * @return Status
 233          */
 234         protected function waitForSlotOrNotif( $doWakeup ) {
 235                 if ( $this->slot !== null ) {
 236                         return Status::newGood( PoolCounter::LOCK_HELD ); // already acquired
 237                 }
 238
 239                 $status = $this->getConnection();
 240                 if ( !$status->isOK() ) {
 241                         return $status;
 242                 }
 243                 /** @var RedisConnRef $conn */
 244                 $conn = $status->value;
 245                 '@phan-var RedisConnRef $conn';
 246
 247                 $now = microtime( true );
 248                 try {
 249                         $slot = $this->initAndPopPoolSlotList( $conn, $now );
 250                         if ( ctype_digit( $slot ) ) {
 251                                 // Pool slot acquired by this process
 252                                 $slotTime = $now;
 253                         } elseif ( $slot === 'QUEUE_FULL' ) {
 254                                 // Too many processes are waiting for pooled processes to finish
 255                                 return Status::newGood( PoolCounter::QUEUE_FULL );
 256                         } elseif ( $slot === 'QUEUE_WAIT' ) {
 257                                 // This process is now registered as waiting
 258                                 $keys = ( $doWakeup == self::AWAKE_ALL )
 259                                         // Wait for an open slot or wake-up signal (preferring the latter)
 260                                         ? [ $this->getWakeupListKey(), $this->getSlotListKey() ]
 261                                         // Just wait for an actual pool slot
 262                                         : [ $this->getSlotListKey() ];
 263
 264                                 $res = $conn->blPop( $keys, $this->timeout );
 265                                 if ( $res === [] ) {
 266                                         $conn->zRem( $this->getWaitSetKey(), $this->session ); // no longer waiting
 267                                         return Status::newGood( PoolCounter::TIMEOUT );
 268                                 }
 269
 270                                 $slot = $res[1]; // pool slot or "w" for wake-up notifications
 271                                 $slotTime = microtime( true ); // last microtime() was a few RTTs ago
 272                                 // Unregister this process as waiting and bump slot "next release" time
 273                                 $this->registerAcquisitionTime( $conn, $slot, $slotTime );
 274                         } else {
 275                                 return Status::newFatal( 'pool-error-unknown', "Server gave slot '$slot'." );
 276                         }
 277                 } catch ( RedisException $e ) {
 278                         return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
 279                 }
 280
 281                 if ( $slot !== 'w' ) {
 282                         $this->slot = $slot;
 283                         $this->slotTime = $slotTime;
 284                         $this->onRelease = $doWakeup;
 285                         self::$active[$this->session] = $this;
 286                 }
 287
 288                 $this->onAcquire();
 289
 290                 return Status::newGood( $slot === 'w' ? PoolCounter::DONE : PoolCounter::LOCKED );
 291         }
 292
 293         /**
 294          * @param RedisConnRef $conn
 295          * @param float $now UNIX timestamp
 296          * @return string|bool False on failure
 297          */
 298         protected function initAndPopPoolSlotList( RedisConnRef $conn, $now ) {
 299                 static $script =
 300                 /** @lang Lua */
 301 <<<LUA
 302                 local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
 303                 local rMaxWorkers,rMaxQueue,rTimeout,rExpiry,rSess,rTime = unpack(ARGV)
 304                 -- Initialize if the "next release" time sorted-set is empty. The slot key
 305                 -- itself is empty if all slots are busy or when nothing is initialized.
 306                 -- If the list is empty but the set is not, then it is the latter case.
 307                 -- For sanity, if the list exists but not the set, then reset everything.
 308                 if redis.call('exists',kSlotsNextRelease) == 0 then
 309                         redis.call('del',kSlots)
 310                         for i = 1,1*rMaxWorkers do
 311                                 redis.call('rPush',kSlots,i)
 312                                 redis.call('zAdd',kSlotsNextRelease,-1,i)
 313                         end
 314                 -- Otherwise do maintenance to clean up after network partitions
 315                 else
 316                         -- Find stale slot locks and add free them (avoid duplicates for sanity)
 317                         local staleLocks = redis.call('zRangeByScore',kSlotsNextRelease,0,rTime)
 318                         for k,slot in ipairs(staleLocks) do
 319                                 redis.call('lRem',kSlots,0,slot)
 320                                 redis.call('rPush',kSlots,slot)
 321                                 redis.call('zAdd',kSlotsNextRelease,rTime + 30,slot)
 322                         end
 323                         -- Find stale wait slot entries and remove them
 324                         redis.call('zRemRangeByScore',kSlotWaits,0,rTime - 2*rTimeout)
 325                 end
 326                 local slot
 327                 -- Try to acquire a slot if possible now
 328                 if redis.call('lLen',kSlots) > 0 then
 329                         slot = redis.call('lPop',kSlots)
 330                         -- Update the slot "next release" time
 331                         redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,slot)
 332                 elseif redis.call('zCard',kSlotWaits) >= 1*rMaxQueue then
 333                         slot = 'QUEUE_FULL'
 334                 else
 335                         slot = 'QUEUE_WAIT'
 336                         -- Register this process as waiting
 337                         redis.call('zAdd',kSlotWaits,rTime,rSess)
 338                         redis.call('expireAt',kSlotWaits,math.ceil(rTime + 2*rTimeout))
 339                 end
 340                 -- Always keep renewing the expiry on use
 341                 redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
 342                 redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
 343                 return slot
 344 LUA;
 345                 return $conn->luaEval( $script,
 346                         [
 347                                 $this->getSlotListKey(),
 348                                 $this->getSlotRTimeSetKey(),
 349                                 $this->getWaitSetKey(),
 350                                 $this->workers,
 351                                 $this->maxqueue,
 352                                 $this->timeout,
 353                                 $this->lockTTL,
 354                                 $this->session,
 355                                 $now
 356                         ],
 357                         3 # number of first argument(s) that are keys
 358                 );
 359         }
 360
 361         /**
 362          * @param RedisConnRef $conn
 363          * @param string $slot
 364          * @param float $now
 365          * @return int|bool False on failure
 366          */
 367         protected function registerAcquisitionTime( RedisConnRef $conn, $slot, $now ) {
 368                 static $script =
 369                 /** @lang Lua */
 370 <<<LUA
 371                 local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
 372                 local rSlot,rExpiry,rSess,rTime = unpack(ARGV)
 373                 -- If rSlot is 'w' then the client was told to wake up but got no slot
 374                 if rSlot ~= 'w' then
 375                         -- Update the slot "next release" time
 376                         redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,rSlot)
 377                         -- Always keep renewing the expiry on use
 378                         redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
 379                         redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
 380                 end
 381                 -- Unregister this process as waiting
 382                 redis.call('zRem',kSlotWaits,rSess)
 383                 return 1
 384 LUA;
 385                 return $conn->luaEval( $script,
 386                         [
 387                                 $this->getSlotListKey(),
 388                                 $this->getSlotRTimeSetKey(),
 389                                 $this->getWaitSetKey(),
 390                                 $slot,
 391                                 $this->lockTTL,
 392                                 $this->session,
 393                                 $now
 394                         ],
 395                         3 # number of first argument(s) that are keys
 396                 );
 397         }
 398
 399         /**
 400          * @return string
 401          */
 402         protected function getSlotListKey() {
 403                 return "poolcounter:l-slots-{$this->keySha1}-{$this->workers}";
 404         }
 405
 406         /**
 407          * @return string
 408          */
 409         protected function getSlotRTimeSetKey() {
 410                 return "poolcounter:z-renewtime-{$this->keySha1}-{$this->workers}";
 411         }
 412
 413         /**
 414          * @return string
 415          */
 416         protected function getWaitSetKey() {
 417                 return "poolcounter:z-wait-{$this->keySha1}-{$this->workers}";
 418         }
 419
 420         /**
 421          * @return string
 422          */
 423         protected function getWakeupListKey() {
 424                 return "poolcounter:l-wakeup-{$this->keySha1}-{$this->workers}";
 425         }
 426
 427         /**
 428          * Try to make sure that locks get released (even with exceptions and fatals)
 429          */
 430         public static function releaseAll() {
 431                 foreach ( self::$active as $poolCounter ) {
 432                         try {
 433                                 if ( $poolCounter->slot !== null ) {
 434                                         $poolCounter->release();
 435                                 }
 436                         } catch ( Exception $e ) {
 437                         }
 438                 }
 439         }
 440 }