JobQueueRedis.php

Go to the documentation of this file.

<?php
class JobQueueRedis extends JobQueue {
    protected $redisPool;

    protected $server;
    protected $compression;

    public function __construct( array $params ) {
        parent::__construct( $params );
        $params['redisConfig']['serializer'] = 'none'; // make it easy to use Lua
        $this->server = $params['redisServer'];
        $this->compression = isset( $params['compression'] ) ? $params['compression'] : 'none';
        $this->redisPool = RedisConnectionPool::singleton( $params['redisConfig'] );
        if ( empty( $params['daemonized'] ) ) {
            throw new InvalidArgumentException(
                "Non-daemonized mode is no longer supported. Please install the " .
                "mediawiki/services/jobrunner service and update \$wgJobTypeConf as needed." );
        }
    }

    protected function supportedOrders() {
        return [ 'timestamp', 'fifo' ];
    }

    protected function optimalOrder() {
        return 'fifo';
    }

    protected function supportsDelayedJobs() {
        return true;
    }

    protected function doIsEmpty() {
        return $this->doGetSize() == 0;
    }

    protected function doGetSize() {
        $conn = $this->getConnection();
        try {
            return $conn->lSize( $this->getQueueKey( 'l-unclaimed' ) );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function doGetAcquiredCount() {
        $conn = $this->getConnection();
        try {
            $conn->multi( Redis::PIPELINE );
            $conn->zSize( $this->getQueueKey( 'z-claimed' ) );
            $conn->zSize( $this->getQueueKey( 'z-abandoned' ) );

            return array_sum( $conn->exec() );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function doGetDelayedCount() {
        $conn = $this->getConnection();
        try {
            return $conn->zSize( $this->getQueueKey( 'z-delayed' ) );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function doGetAbandonedCount() {
        $conn = $this->getConnection();
        try {
            return $conn->zSize( $this->getQueueKey( 'z-abandoned' ) );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function doBatchPush( array $jobs, $flags ) {
        // Convert the jobs into field maps (de-duplicated against each other)
        $items = []; // (job ID => job fields map)
        foreach ( $jobs as $job ) {
            $item = $this->getNewJobFields( $job );
            if ( strlen( $item['sha1'] ) ) { // hash identifier => de-duplicate
                $items[$item['sha1']] = $item;
            } else {
                $items[$item['uuid']] = $item;
            }
        }

        if ( !count( $items ) ) {
            return; // nothing to do
        }

        $conn = $this->getConnection();
        try {
            // Actually push the non-duplicate jobs into the queue...
            if ( $flags & self::QOS_ATOMIC ) {
                $batches = [ $items ]; // all or nothing
            } else {
                $batches = array_chunk( $items, 100 ); // avoid tying up the server
            }
            $failed = 0;
            $pushed = 0;
            foreach ( $batches as $itemBatch ) {
                $added = $this->pushBlobs( $conn, $itemBatch );
                if ( is_int( $added ) ) {
                    $pushed += $added;
                } else {
                    $failed += count( $itemBatch );
                }
            }
            JobQueue::incrStats( 'inserts', $this->type, count( $items ) );
            JobQueue::incrStats( 'inserts_actual', $this->type, $pushed );
            JobQueue::incrStats( 'dupe_inserts', $this->type,
                count( $items ) - $failed - $pushed );
            if ( $failed > 0 ) {
                $err = "Could not insert {$failed} {$this->type} job(s).";
                wfDebugLog( 'JobQueueRedis', $err );
                throw new RedisException( $err );
            }
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function pushBlobs( RedisConnRef $conn, array $items ) {
        $args = [ $this->encodeQueueName() ];
        // Next args come in 4s ([id, sha1, rtime, blob [, id, sha1, rtime, blob ... ] ] )
        foreach ( $items as $item ) {
            $args[] = (string)$item['uuid'];
            $args[] = (string)$item['sha1'];
            $args[] = (string)$item['rtimestamp'];
            $args[] = (string)$this->serialize( $item );
        }
        static $script =
<<<LUA
        local kUnclaimed, kSha1ById, kIdBySha1, kDelayed, kData, kQwJobs = unpack(KEYS)
        -- First argument is the queue ID
        local queueId = ARGV[1]
        -- Next arguments all come in 4s (one per job)
        local variadicArgCount = #ARGV - 1
        if variadicArgCount % 4 ~= 0 then
            return redis.error_reply('Unmatched arguments')
        end
        -- Insert each job into this queue as needed
        local pushed = 0
        for i = 2,#ARGV,4 do
            local id,sha1,rtimestamp,blob = ARGV[i],ARGV[i+1],ARGV[i+2],ARGV[i+3]
            if sha1 == '' or redis.call('hExists',kIdBySha1,sha1) == 0 then
                if 1*rtimestamp > 0 then
                    -- Insert into delayed queue (release time as score)
                    redis.call('zAdd',kDelayed,rtimestamp,id)
                else
                    -- Insert into unclaimed queue
                    redis.call('lPush',kUnclaimed,id)
                end
                if sha1 ~= '' then
                    redis.call('hSet',kSha1ById,id,sha1)
                    redis.call('hSet',kIdBySha1,sha1,id)
                end
                redis.call('hSet',kData,id,blob)
                pushed = pushed + 1
            end
        end
        -- Mark this queue as having jobs
        redis.call('sAdd',kQwJobs,queueId)
        return pushed
LUA;
        return $conn->luaEval( $script,
            array_merge(
                [
                    $this->getQueueKey( 'l-unclaimed' ), # KEYS[1]
                    $this->getQueueKey( 'h-sha1ById' ), # KEYS[2]
                    $this->getQueueKey( 'h-idBySha1' ), # KEYS[3]
                    $this->getQueueKey( 'z-delayed' ), # KEYS[4]
                    $this->getQueueKey( 'h-data' ), # KEYS[5]
                    $this->getGlobalKey( 's-queuesWithJobs' ), # KEYS[6]
                ],
                $args
            ),
            6 # number of first argument(s) that are keys
        );
    }

    protected function doPop() {
        $job = false;

        $conn = $this->getConnection();
        try {
            do {
                $blob = $this->popAndAcquireBlob( $conn );
                if ( !is_string( $blob ) ) {
                    break; // no jobs; nothing to do
                }

                JobQueue::incrStats( 'pops', $this->type );
                $item = $this->unserialize( $blob );
                if ( $item === false ) {
                    wfDebugLog( 'JobQueueRedis', "Could not unserialize {$this->type} job." );
                    continue;
                }

                // If $item is invalid, the runner loop recyling will cleanup as needed
                $job = $this->getJobFromFields( $item ); // may be false
            } while ( !$job ); // job may be false if invalid
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $job;
    }

    protected function popAndAcquireBlob( RedisConnRef $conn ) {
        static $script =
<<<LUA
        local kUnclaimed, kSha1ById, kIdBySha1, kClaimed, kAttempts, kData = unpack(KEYS)
        local rTime = unpack(ARGV)
        -- Pop an item off the queue
        local id = redis.call('rPop',kUnclaimed)
        if not id then
            return false
        end
        -- Allow new duplicates of this job
        local sha1 = redis.call('hGet',kSha1ById,id)
        if sha1 then redis.call('hDel',kIdBySha1,sha1) end
        redis.call('hDel',kSha1ById,id)
        -- Mark the jobs as claimed and return it
        redis.call('zAdd',kClaimed,rTime,id)
        redis.call('hIncrBy',kAttempts,id,1)
        return redis.call('hGet',kData,id)
LUA;
        return $conn->luaEval( $script,
            [
                $this->getQueueKey( 'l-unclaimed' ), # KEYS[1]
                $this->getQueueKey( 'h-sha1ById' ), # KEYS[2]
                $this->getQueueKey( 'h-idBySha1' ), # KEYS[3]
                $this->getQueueKey( 'z-claimed' ), # KEYS[4]
                $this->getQueueKey( 'h-attempts' ), # KEYS[5]
                $this->getQueueKey( 'h-data' ), # KEYS[6]
                time(), # ARGV[1] (injected to be replication-safe)
            ],
            6 # number of first argument(s) that are keys
        );
    }

    protected function doAck( Job $job ) {
        if ( !isset( $job->metadata['uuid'] ) ) {
            throw new UnexpectedValueException( "Job of type '{$job->getType()}' has no UUID." );
        }

        $uuid = $job->metadata['uuid'];
        $conn = $this->getConnection();
        try {
            static $script =
<<<LUA
            local kClaimed, kAttempts, kData = unpack(KEYS)
            local id = unpack(ARGV)
            -- Unmark the job as claimed
            local removed = redis.call('zRem',kClaimed,id)
            -- Check if the job was recycled
            if removed == 0 then
                return 0
            end
            -- Delete the retry data
            redis.call('hDel',kAttempts,id)
            -- Delete the job data itself
            return redis.call('hDel',kData,id)
LUA;
            $res = $conn->luaEval( $script,
                [
                    $this->getQueueKey( 'z-claimed' ), # KEYS[1]
                    $this->getQueueKey( 'h-attempts' ), # KEYS[2]
                    $this->getQueueKey( 'h-data' ), # KEYS[3]
                    $uuid # ARGV[1]
                ],
                3 # number of first argument(s) that are keys
            );

            if ( !$res ) {
                wfDebugLog( 'JobQueueRedis', "Could not acknowledge {$this->type} job $uuid." );

                return false;
            }

            JobQueue::incrStats( 'acks', $this->type );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return true;
    }

    protected function doDeduplicateRootJob( IJobSpecification $job ) {
        if ( !$job->hasRootJobParams() ) {
            throw new LogicException( "Cannot register root job; missing parameters." );
        }
        $params = $job->getRootJobParams();

        $key = $this->getRootJobCacheKey( $params['rootJobSignature'] );

        $conn = $this->getConnection();
        try {
            $timestamp = $conn->get( $key ); // current last timestamp of this job
            if ( $timestamp && $timestamp >= $params['rootJobTimestamp'] ) {
                return true; // a newer version of this root job was enqueued
            }

            // Update the timestamp of the last root job started at the location...
            return $conn->set( $key, $params['rootJobTimestamp'], self::ROOTJOB_TTL ); // 2 weeks
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    protected function doIsRootJobOldDuplicate( Job $job ) {
        if ( !$job->hasRootJobParams() ) {
            return false; // job has no de-deplication info
        }
        $params = $job->getRootJobParams();

        $conn = $this->getConnection();
        try {
            // Get the last time this root job was enqueued
            $timestamp = $conn->get( $this->getRootJobCacheKey( $params['rootJobSignature'] ) );
        } catch ( RedisException $e ) {
            $timestamp = false;
            $this->throwRedisException( $conn, $e );
        }

        // Check if a new root job was started at the location after this one's...
        return ( $timestamp && $timestamp > $params['rootJobTimestamp'] );
    }

    protected function doDelete() {
        static $props = [ 'l-unclaimed', 'z-claimed', 'z-abandoned',
            'z-delayed', 'h-idBySha1', 'h-sha1ById', 'h-attempts', 'h-data' ];

        $conn = $this->getConnection();
        try {
            $keys = [];
            foreach ( $props as $prop ) {
                $keys[] = $this->getQueueKey( $prop );
            }

            $ok = ( $conn->delete( $keys ) !== false );
            $conn->sRem( $this->getGlobalKey( 's-queuesWithJobs' ), $this->encodeQueueName() );

            return $ok;
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    public function getAllQueuedJobs() {
        $conn = $this->getConnection();
        try {
            $uids = $conn->lRange( $this->getQueueKey( 'l-unclaimed' ), 0, -1 );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $this->getJobIterator( $conn, $uids );
    }

    public function getAllDelayedJobs() {
        $conn = $this->getConnection();
        try {
            $uids = $conn->zRange( $this->getQueueKey( 'z-delayed' ), 0, -1 );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $this->getJobIterator( $conn, $uids );
    }

    public function getAllAcquiredJobs() {
        $conn = $this->getConnection();
        try {
            $uids = $conn->zRange( $this->getQueueKey( 'z-claimed' ), 0, -1 );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $this->getJobIterator( $conn, $uids );
    }

    public function getAllAbandonedJobs() {
        $conn = $this->getConnection();
        try {
            $uids = $conn->zRange( $this->getQueueKey( 'z-abandoned' ), 0, -1 );
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $this->getJobIterator( $conn, $uids );
    }

    protected function getJobIterator( RedisConnRef $conn, array $uids ) {
        return new MappedIterator(
            $uids,
            function ( $uid ) use ( $conn ) {
                return $this->getJobFromUidInternal( $uid, $conn );
            },
            [ 'accept' => function ( $job ) {
                return is_object( $job );
            } ]
        );
    }

    public function getCoalesceLocationInternal() {
        return "RedisServer:" . $this->server;
    }

    protected function doGetSiblingQueuesWithJobs( array $types ) {
        return array_keys( array_filter( $this->doGetSiblingQueueSizes( $types ) ) );
    }

    protected function doGetSiblingQueueSizes( array $types ) {
        $sizes = []; // (type => size)
        $types = array_values( $types ); // reindex
        $conn = $this->getConnection();
        try {
            $conn->multi( Redis::PIPELINE );
            foreach ( $types as $type ) {
                $conn->lSize( $this->getQueueKey( 'l-unclaimed', $type ) );
            }
            $res = $conn->exec();
            if ( is_array( $res ) ) {
                foreach ( $res as $i => $size ) {
                    $sizes[$types[$i]] = $size;
                }
            }
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $sizes;
    }

    public function getJobFromUidInternal( $uid, RedisConnRef $conn ) {
        try {
            $data = $conn->hGet( $this->getQueueKey( 'h-data' ), $uid );
            if ( $data === false ) {
                return false; // not found
            }
            $item = $this->unserialize( $data );
            if ( !is_array( $item ) ) { // this shouldn't happen
                throw new UnexpectedValueException( "Could not find job with ID '$uid'." );
            }
            $title = Title::makeTitle( $item['namespace'], $item['title'] );
            $job = Job::factory( $item['type'], $title, $item['params'] );
            $job->metadata['uuid'] = $item['uuid'];
            $job->metadata['timestamp'] = $item['timestamp'];
            // Add in attempt count for debugging at showJobs.php
            $job->metadata['attempts'] = $conn->hGet( $this->getQueueKey( 'h-attempts' ), $uid );

            return $job;
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }
    }

    public function getServerQueuesWithJobs() {
        $queues = [];

        $conn = $this->getConnection();
        try {
            $set = $conn->sMembers( $this->getGlobalKey( 's-queuesWithJobs' ) );
            foreach ( $set as $queue ) {
                $queues[] = $this->decodeQueueName( $queue );
            }
        } catch ( RedisException $e ) {
            $this->throwRedisException( $conn, $e );
        }

        return $queues;
    }

    protected function getNewJobFields( IJobSpecification $job ) {
        return [
            // Fields that describe the nature of the job
            'type' => $job->getType(),
            'namespace' => $job->getTitle()->getNamespace(),
            'title' => $job->getTitle()->getDBkey(),
            'params' => $job->getParams(),
            // Some jobs cannot run until a "release timestamp"
            'rtimestamp' => $job->getReleaseTimestamp() ?: 0,
            // Additional job metadata
            'uuid' => UIDGenerator::newRawUUIDv4( UIDGenerator::QUICK_RAND ),
            'sha1' => $job->ignoreDuplicates()
                ? Wikimedia\base_convert( sha1( serialize( $job->getDeduplicationInfo() ) ), 16, 36, 31 )
                : '',
            'timestamp' => time() // UNIX timestamp
        ];
    }

    protected function getJobFromFields( array $fields ) {
        $title = Title::makeTitle( $fields['namespace'], $fields['title'] );
        $job = Job::factory( $fields['type'], $title, $fields['params'] );
        $job->metadata['uuid'] = $fields['uuid'];
        $job->metadata['timestamp'] = $fields['timestamp'];

        return $job;
    }

    protected function serialize( array $fields ) {
        $blob = serialize( $fields );
        if ( $this->compression === 'gzip'
            && strlen( $blob ) >= 1024
            && function_exists( 'gzdeflate' )
        ) {
            $object = (object)[ 'blob' => gzdeflate( $blob ), 'enc' => 'gzip' ];
            $blobz = serialize( $object );

            return ( strlen( $blobz ) < strlen( $blob ) ) ? $blobz : $blob;
        } else {
            return $blob;
        }
    }

    protected function unserialize( $blob ) {
        $fields = unserialize( $blob );
        if ( is_object( $fields ) ) {
            if ( $fields->enc === 'gzip' && function_exists( 'gzinflate' ) ) {
                $fields = unserialize( gzinflate( $fields->blob ) );
            } else {
                $fields = false;
            }
        }

        return is_array( $fields ) ? $fields : false;
    }

    protected function getConnection() {
        $conn = $this->redisPool->getConnection( $this->server );
        if ( !$conn ) {
            throw new JobQueueConnectionError(
                "Unable to connect to redis server {$this->server}." );
        }

        return $conn;
    }

    protected function throwRedisException( RedisConnRef $conn, $e ) {
        $this->redisPool->handleError( $conn, $e );
        throw new JobQueueError( "Redis server error: {$e->getMessage()}\n" );
    }

    private function encodeQueueName() {
        return json_encode( [ $this->type, $this->wiki ] );
    }

    private function decodeQueueName( $name ) {
        return json_decode( $name );
    }

    private function getGlobalKey( $name ) {
        $parts = [ 'global', 'jobqueue', $name ];
        foreach ( $parts as $part ) {
            if ( !preg_match( '/[a-zA-Z0-9_-]+/', $part ) ) {
                throw new InvalidArgumentException( "Key part characters are out of range." );
            }
        }

        return implode( ':', $parts );
    }

    private function getQueueKey( $prop, $type = null ) {
        $type = is_string( $type ) ? $type : $this->type;
        list( $db, $prefix ) = wfSplitWikiID( $this->wiki );
        $keyspace = $prefix ? "$db-$prefix" : $db;

        $parts = [ $keyspace, 'jobqueue', $type, $prop ];

        // Parts are typically ASCII, but encode for sanity to escape ":"
        return implode( ':', array_map( 'rawurlencode', $parts ) );
    }
}