PoolCounterRedis.php

Go to the documentation of this file.

<?php
namespace MediaWiki\PoolCounter;
 
use Exception;
use MediaWiki\Status\Status;
use RedisException;
use Wikimedia\ArrayUtils\ArrayUtils;
use Wikimedia\HashRing\HashRing;
use Wikimedia\ObjectCache\RedisConnectionPool;
use Wikimedia\ObjectCache\RedisConnRef;
 
class PoolCounterRedis extends PoolCounter {
    protected $ring;
    protected $pool;
    protected $serversByLabel;
    protected $keySha1;
    protected $lockTTL;
    protected $conn;
    protected $slot;
    protected $onRelease;
    protected $session;
    protected $slotTime;
 
    private const AWAKE_ONE = 1; // wake-up if when a slot can be taken from an existing process
    private const AWAKE_ALL = 2; // wake-up if an existing process finishes and wake up such others
 
    protected static $active = null;
 
    public function __construct( array $conf, string $type, string $key ) {
        parent::__construct( $conf, $type, $key );
 
        $this->serversByLabel = $conf['servers'];
 
        $serverLabels = array_keys( $conf['servers'] );
        $this->ring = new HashRing( array_fill_keys( $serverLabels, 10 ) );
 
        $conf['redisConfig']['serializer'] = 'none'; // for use with Lua
        $this->pool = RedisConnectionPool::singleton( $conf['redisConfig'] );
 
        $this->keySha1 = sha1( $this->key );
        $met = ini_get( 'max_execution_time' ); // usually 0 in CLI mode
        $this->lockTTL = $met ? 2 * (int)$met : 3600;
 
        if ( self::$active === null ) {
            self::$active = [];
            register_shutdown_function( self::releaseAll( ... ) );
        }
    }
 
    protected function getConnection() {
        if ( !$this->conn ) {
            $conn = false;
            $servers = $this->ring->getLocations( $this->key, 3 );
            ArrayUtils::consistentHashSort( $servers, $this->key );
            foreach ( $servers as $server ) {
                $conn = $this->pool->getConnection( $this->serversByLabel[$server], $this->logger );
                if ( $conn ) {
                    break;
                }
            }
            if ( !$conn ) {
                return Status::newFatal( 'pool-servererror', implode( ', ', $servers ) );
            }
            $this->conn = $conn;
        }
        return Status::newGood( $this->conn );
    }
 
    public function acquireForMe( $timeout = null ) {
        $status = $this->precheckAcquire();
        if ( !$status->isGood() ) {
            return $status;
        }
 
        return $this->waitForSlotOrNotif( self::AWAKE_ONE, $timeout );
    }
 
    public function acquireForAnyone( $timeout = null ) {
        $status = $this->precheckAcquire();
        if ( !$status->isGood() ) {
            return $status;
        }
 
        return $this->waitForSlotOrNotif( self::AWAKE_ALL, $timeout );
    }
 
    public function release() {
        if ( $this->slot === null ) {
            return Status::newGood( PoolCounter::NOT_LOCKED ); // not locked
        }
 
        $status = $this->getConnection();
        if ( !$status->isOK() ) {
            return $status;
        }
        $conn = $status->value;
        '@phan-var RedisConnRef $conn';
 
        // phpcs:disable Generic.Files.LineLength
        static $script =
<<<LUA
        local kSlots,kSlotsNextRelease,kWakeup,kWaiting = unpack(KEYS)
        local rMaxWorkers,rExpiry,rSlot,rSlotTime,rAwakeAll,rTime = unpack(ARGV)
        -- Add the slots back to the list (if rSlot is "w" then it is not a slot).
        -- Treat the list as expired if the "next release" time sorted-set is missing.
        if rSlot ~= 'w' and redis.call('exists',kSlotsNextRelease) == 1 then
            if 1*redis.call('zScore',kSlotsNextRelease,rSlot) ~= (rSlotTime + rExpiry) then
                -- Slot lock expired and was released already
            elseif redis.call('lLen',kSlots) >= 1*rMaxWorkers then
                -- Slots somehow got out of sync; reset the list
                redis.call('del',kSlots,kSlotsNextRelease)
            elseif redis.call('lLen',kSlots) == (1*rMaxWorkers - 1) and redis.call('zCard',kWaiting) == 0 then
                -- Slot list will be made full; clear it to save space (it re-inits as needed)
                -- since nothing is waiting on being unblocked by a push to the list
                redis.call('del',kSlots,kSlotsNextRelease)
            else
                -- Add slot back to pool and update the "next release" time
                redis.call('rPush',kSlots,rSlot)
                redis.call('zAdd',kSlotsNextRelease,rTime + 30,rSlot)
                -- Always keep renewing the expiry on use
                redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
                redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
            end
        end
        -- Update an ephemeral list to wake up other clients that can
        -- reuse any cached work from this process. Only do this if no
        -- slots are currently free (e.g. clients could be waiting).
        if 1*rAwakeAll == 1 then
            local count = redis.call('zCard',kWaiting)
            for i = 1,count do
                redis.call('rPush',kWakeup,'w')
            end
            redis.call('pexpire',kWakeup,1)
        end
        return 1
LUA;
        // phpcs:enable
 
        try {
            $conn->luaEval(
                $script,
                [
                    $this->getSlotListKey(),
                    $this->getSlotRTimeSetKey(),
                    $this->getWakeupListKey(),
                    $this->getWaitSetKey(),
                    $this->workers,
                    $this->lockTTL,
                    $this->slot,
                    $this->slotTime, // used for CAS-style check
                    ( $this->onRelease === self::AWAKE_ALL ) ? 1 : 0,
                    microtime( true ),
                ],
                4 # number of first argument(s) that are keys
            );
        } catch ( RedisException $e ) {
            return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
        }
 
        $this->slot = null;
        $this->slotTime = null;
        $this->onRelease = null;
        unset( self::$active[$this->session] );
 
        $this->onRelease();
 
        return Status::newGood( PoolCounter::RELEASED );
    }
 
    protected function waitForSlotOrNotif( $doWakeup, $timeout = null ) {
        if ( $this->slot !== null ) {
            return Status::newGood( PoolCounter::LOCK_HELD ); // already acquired
        }
 
        $status = $this->getConnection();
        if ( !$status->isOK() ) {
            return $status;
        }
        $conn = $status->value;
        '@phan-var RedisConnRef $conn';
 
        $now = microtime( true );
        $timeout ??= $this->timeout;
        try {
            $slot = $this->initAndPopPoolSlotList( $conn, $now );
            if ( ctype_digit( $slot ) ) {
                // Pool slot acquired by this process
                $slotTime = $now;
            } elseif ( $slot === 'QUEUE_FULL' ) {
                // Too many processes are waiting for pooled processes to finish
                return Status::newGood( PoolCounter::QUEUE_FULL );
            } elseif ( $slot === 'QUEUE_WAIT' ) {
                // This process is now registered as waiting
                $keys =
                    ( $doWakeup == self::AWAKE_ALL ) // Wait for an open slot or wake-up signal (preferring the latter)
                        ? [ $this->getWakeupListKey(), $this->getSlotListKey() ] // Just wait for an actual pool slot
                        : [ $this->getSlotListKey() ];
 
                $res = $conn->blPop( $keys, $timeout );
                if ( $res === [] ) {
                    $conn->zRem( $this->getWaitSetKey(), $this->session ); // no longer waiting
                    return Status::newGood( PoolCounter::TIMEOUT );
                }
 
                $slot = $res[1]; // pool slot or "w" for wake-up notifications
                $slotTime = microtime( true ); // last microtime() was a few RTTs ago
                // Unregister this process as waiting and bump slot "next release" time
                $this->registerAcquisitionTime( $conn, $slot, $slotTime );
            } else {
                return Status::newFatal( 'pool-error-unknown', "Server gave slot '$slot'." );
            }
        } catch ( RedisException $e ) {
            return Status::newFatal( 'pool-error-unknown', $e->getMessage() );
        }
 
        if ( $slot !== 'w' ) {
            $this->slot = $slot;
            $this->slotTime = $slotTime;
            $this->onRelease = $doWakeup;
            self::$active[$this->session] = $this;
        }
 
        $this->onAcquire();
 
        return Status::newGood( $slot === 'w' ? PoolCounter::DONE : PoolCounter::LOCKED );
    }
 
    protected function initAndPopPoolSlotList( RedisConnRef $conn, $now ) {
        static $script =  <<<LUA
        local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
        local rMaxWorkers,rMaxQueue,rTimeout,rExpiry,rSess,rTime = unpack(ARGV)
        -- Initialize if the "next release" time sorted-set is empty. The slot key
        -- itself is empty if all slots are busy or when nothing is initialized.
        -- If the list is empty but the set is not, then it is the latter case.
        -- If the list exists but not the set, then reset everything.
        if redis.call('exists',kSlotsNextRelease) == 0 then
            redis.call('del',kSlots)
            for i = 1,1*rMaxWorkers do
                redis.call('rPush',kSlots,i)
                redis.call('zAdd',kSlotsNextRelease,-1,i)
            end
        -- Otherwise do maintenance to clean up after network partitions
        else
            -- Find stale slot locks and add free them (avoid duplicates)
            local staleLocks = redis.call('zRangeByScore',kSlotsNextRelease,0,rTime)
            for k,slot in ipairs(staleLocks) do
                redis.call('lRem',kSlots,0,slot)
                redis.call('rPush',kSlots,slot)
                redis.call('zAdd',kSlotsNextRelease,rTime + 30,slot)
            end
            -- Find stale wait slot entries and remove them
            redis.call('zRemRangeByScore',kSlotWaits,0,rTime - 2*rTimeout)
        end
        local slot
        -- Try to acquire a slot if possible now
        if redis.call('lLen',kSlots) > 0 then
            slot = redis.call('lPop',kSlots)
            -- Update the slot "next release" time
            redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,slot)
        elseif redis.call('zCard',kSlotWaits) >= 1*rMaxQueue then
            slot = 'QUEUE_FULL'
        else
            slot = 'QUEUE_WAIT'
            -- Register this process as waiting
            redis.call('zAdd',kSlotWaits,rTime,rSess)
            redis.call('expireAt',kSlotWaits,math.ceil(rTime + 2*rTimeout))
        end
        -- Always keep renewing the expiry on use
        redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
        redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
        return slot
LUA;
        return $conn->luaEval(
            $script,
            [
                $this->getSlotListKey(),
                $this->getSlotRTimeSetKey(),
                $this->getWaitSetKey(),
                $this->workers,
                $this->maxqueue,
                $this->timeout,
                $this->lockTTL,
                $this->session,
                $now,
            ],
            3 # number of first argument(s) that are keys
        );
    }
 
    protected function registerAcquisitionTime( RedisConnRef $conn, $slot, $now ) {
        static $script =  <<<LUA
        local kSlots,kSlotsNextRelease,kSlotWaits = unpack(KEYS)
        local rSlot,rExpiry,rSess,rTime = unpack(ARGV)
        -- If rSlot is 'w' then the client was told to wake up but got no slot
        if rSlot ~= 'w' then
            -- Update the slot "next release" time
            redis.call('zAdd',kSlotsNextRelease,rTime + rExpiry,rSlot)
            -- Always keep renewing the expiry on use
            redis.call('expireAt',kSlots,math.ceil(rTime + rExpiry))
            redis.call('expireAt',kSlotsNextRelease,math.ceil(rTime + rExpiry))
        end
        -- Unregister this process as waiting
        redis.call('zRem',kSlotWaits,rSess)
        return 1
LUA;
        return $conn->luaEval(
            $script,
            [
                $this->getSlotListKey(),
                $this->getSlotRTimeSetKey(),
                $this->getWaitSetKey(),
                $slot,
                $this->lockTTL,
                $this->session,
                $now,
            ],
            3 # number of first argument(s) that are keys
        );
    }
 
    protected function getSlotListKey() {
        return "poolcounter:l-slots-{$this->keySha1}-{$this->workers}";
    }
 
    protected function getSlotRTimeSetKey() {
        return "poolcounter:z-renewtime-{$this->keySha1}-{$this->workers}";
    }
 
    protected function getWaitSetKey() {
        return "poolcounter:z-wait-{$this->keySha1}-{$this->workers}";
    }
 
    protected function getWakeupListKey() {
        return "poolcounter:l-wakeup-{$this->keySha1}-{$this->workers}";
    }
 
    public static function releaseAll() {
        $e = null;
        foreach ( self::$active as $poolCounter ) {
            try {
                if ( $poolCounter->slot !== null ) {
                    $poolCounter->release();
                }
            } catch ( Exception $e ) {
            }
        }
        if ( $e ) {
            throw $e;
        }
    }
}
 
class_alias( PoolCounterRedis::class, 'PoolCounterRedis' );