BagOStuff.php

Go to the documentation of this file.

<?php
use Psr\Log\LoggerAwareInterface;
use Psr\Log\LoggerInterface;
use Psr\Log\NullLogger;
use Wikimedia\ScopedCallback;
use Wikimedia\WaitConditionLoop;

abstract class BagOStuff implements IExpiringStore, LoggerAwareInterface {
    protected $locks = [];
    protected $lastError = self::ERR_NONE;
    protected $keyspace = 'local';
    protected $logger;
    protected $asyncHandler;
    protected $syncTimeout;
    protected $segmentationSize;
    protected $segmentedValueMaxSize;

    private $debugMode = false;
    private $duplicateKeyLookups = [];
    private $reportDupes = false;
    private $dupeTrackScheduled = false;

    protected $busyCallbacks = [];

    private $wallClockOverride;

    protected $attrMap = [];

    const READ_LATEST = 1; // use latest data for replicated stores
    const READ_VERIFIED = 2; // promise that caller can tell when keys are stale
    const WRITE_SYNC = 4; // synchronously write to all locations for replicated stores
    const WRITE_CACHE_ONLY = 8; // Only change state of the in-memory cache
    const WRITE_ALLOW_SEGMENTS = 16; // Allow partitioning of the value if it is large
    const WRITE_PRUNE_SEGMENTS = 32; // Delete all partition segments of the value

    const SEGMENT_COMPONENT = 'segment';

    public function __construct( array $params = [] ) {
        $this->setLogger( $params['logger'] ?? new NullLogger() );

        if ( isset( $params['keyspace'] ) ) {
            $this->keyspace = $params['keyspace'];
        }

        $this->asyncHandler = $params['asyncHandler'] ?? null;

        if ( !empty( $params['reportDupes'] ) && is_callable( $this->asyncHandler ) ) {
            $this->reportDupes = true;
        }

        $this->syncTimeout = $params['syncTimeout'] ?? 3;
        $this->segmentationSize = $params['segmentationSize'] ?? 8388608; // 8MiB
        $this->segmentedValueMaxSize = $params['segmentedValueMaxSize'] ?? 67108864; // 64MiB
    }

    public function setLogger( LoggerInterface $logger ) {
        $this->logger = $logger;
    }

    public function setDebug( $bool ) {
        $this->debugMode = $bool;
    }

    final public function getWithSetCallback( $key, $ttl, $callback, $flags = 0 ) {
        $value = $this->get( $key, $flags );

        if ( $value === false ) {
            if ( !is_callable( $callback ) ) {
                throw new InvalidArgumentException( "Invalid cache miss callback provided." );
            }
            $value = call_user_func( $callback );
            if ( $value !== false ) {
                $this->set( $key, $value, $ttl, $flags );
            }
        }

        return $value;
    }

    public function get( $key, $flags = 0 ) {
        $this->trackDuplicateKeys( $key );

        return $this->resolveSegments( $key, $this->doGet( $key, $flags ) );
    }

    private function trackDuplicateKeys( $key ) {
        if ( !$this->reportDupes ) {
            return;
        }

        if ( !isset( $this->duplicateKeyLookups[$key] ) ) {
            // Track that we have seen this key. This N-1 counting style allows
            // easy filtering with array_filter() later.
            $this->duplicateKeyLookups[$key] = 0;
        } else {
            $this->duplicateKeyLookups[$key] += 1;

            if ( $this->dupeTrackScheduled === false ) {
                $this->dupeTrackScheduled = true;
                // Schedule a callback that logs keys processed more than once by get().
                call_user_func( $this->asyncHandler, function () {
                    $dups = array_filter( $this->duplicateKeyLookups );
                    foreach ( $dups as $key => $count ) {
                        $this->logger->warning(
                            'Duplicate get(): "{key}" fetched {count} times',
                            // Count is N-1 of the actual lookup count
                            [ 'key' => $key, 'count' => $count + 1, ]
                        );
                    }
                } );
            }
        }
    }

    abstract protected function doGet( $key, $flags = 0, &$casToken = null );

    public function set( $key, $value, $exptime = 0, $flags = 0 ) {
        if (
            ( $flags & self::WRITE_ALLOW_SEGMENTS ) != self::WRITE_ALLOW_SEGMENTS ||
            is_infinite( $this->segmentationSize )
        ) {
            return $this->doSet( $key, $value, $exptime, $flags );
        }

        $serialized = $this->serialize( $value );
        $segmentSize = $this->getSegmentationSize();
        $maxTotalSize = $this->getSegmentedValueMaxSize();

        $size = strlen( $serialized );
        if ( $size <= $segmentSize ) {
            // Since the work of serializing it was already done, just use it inline
            return $this->doSet(
                $key,
                SerializedValueContainer::newUnified( $serialized ),
                $exptime,
                $flags
            );
        } elseif ( $size > $maxTotalSize ) {
            $this->setLastError( "Key $key exceeded $maxTotalSize bytes." );

            return false;
        }

        $chunksByKey = [];
        $segmentHashes = [];
        $count = intdiv( $size, $segmentSize ) + ( ( $size % $segmentSize ) ? 1 : 0 );
        for ( $i = 0; $i < $count; ++$i ) {
            $segment = substr( $serialized, $i * $segmentSize, $segmentSize );
            $hash = sha1( $segment );
            $chunkKey = $this->makeGlobalKey( self::SEGMENT_COMPONENT, $key, $hash );
            $chunksByKey[$chunkKey] = $segment;
            $segmentHashes[] = $hash;
        }

        $ok = $this->setMulti( $chunksByKey, $exptime, $flags );
        if ( $ok ) {
            // Only when all segments are stored should the main key be changed
            $ok = $this->doSet(
                $key,
                SerializedValueContainer::newSegmented( $segmentHashes ),
                $exptime,
                $flags
            );
        }

        return $ok;
    }

    abstract protected function doSet( $key, $value, $exptime = 0, $flags = 0 );

    public function delete( $key, $flags = 0 ) {
        if ( ( $flags & self::WRITE_PRUNE_SEGMENTS ) != self::WRITE_PRUNE_SEGMENTS ) {
            return $this->doDelete( $key, $flags );
        }

        $mainValue = $this->doGet( $key, self::READ_LATEST );
        if ( !$this->doDelete( $key, $flags ) ) {
            return false;
        }

        if ( !SerializedValueContainer::isSegmented( $mainValue ) ) {
            return true; // no segments to delete
        }

        $orderedKeys = array_map(
            function ( $segmentHash ) use ( $key ) {
                return $this->makeGlobalKey( self::SEGMENT_COMPONENT, $key, $segmentHash );
            },
            $mainValue->{SerializedValueContainer::SEGMENTED_HASHES}
        );

        return $this->deleteMulti( $orderedKeys, $flags );
    }

    abstract protected function doDelete( $key, $flags = 0 );

    abstract public function add( $key, $value, $exptime = 0, $flags = 0 );

    public function merge( $key, callable $callback, $exptime = 0, $attempts = 10, $flags = 0 ) {
        return $this->mergeViaCas( $key, $callback, $exptime, $attempts, $flags );
    }

    protected function mergeViaCas( $key, $callback, $exptime = 0, $attempts = 10, $flags = 0 ) {
        do {
            $casToken = null; // passed by reference
            // Get the old value and CAS token from cache
            $this->clearLastError();
            $currentValue = $this->resolveSegments(
                $key,
                $this->doGet( $key, self::READ_LATEST, $casToken )
            );
            if ( $this->getLastError() ) {
                $this->logger->warning(
                    __METHOD__ . ' failed due to I/O error on get() for {key}.',
                    [ 'key' => $key ]
                );

                return false; // don't spam retries (retry only on races)
            }

            // Derive the new value from the old value
            $value = call_user_func( $callback, $this, $key, $currentValue, $exptime );
            $hadNoCurrentValue = ( $currentValue === false );
            unset( $currentValue ); // free RAM in case the value is large

            $this->clearLastError();
            if ( $value === false ) {
                $success = true; // do nothing
            } elseif ( $hadNoCurrentValue ) {
                // Try to create the key, failing if it gets created in the meantime
                $success = $this->add( $key, $value, $exptime, $flags );
            } else {
                // Try to update the key, failing if it gets changed in the meantime
                $success = $this->cas( $casToken, $key, $value, $exptime, $flags );
            }
            if ( $this->getLastError() ) {
                $this->logger->warning(
                    __METHOD__ . ' failed due to I/O error for {key}.',
                    [ 'key' => $key ]
                );

                return false; // IO error; don't spam retries
            }

        } while ( !$success && --$attempts );

        return $success;
    }

    protected function cas( $casToken, $key, $value, $exptime = 0, $flags = 0 ) {
        if ( !$this->lock( $key, 0 ) ) {
            return false; // non-blocking
        }

        $curCasToken = null; // passed by reference
        $this->doGet( $key, self::READ_LATEST, $curCasToken );
        if ( $casToken === $curCasToken ) {
            $success = $this->set( $key, $value, $exptime, $flags );
        } else {
            $this->logger->info(
                __METHOD__ . ' failed due to race condition for {key}.',
                [ 'key' => $key ]
            );

            $success = false; // mismatched or failed
        }

        $this->unlock( $key );

        return $success;
    }

    public function changeTTL( $key, $exptime = 0, $flags = 0 ) {
        $expiry = $this->convertToExpiry( $exptime );
        $delete = ( $expiry != 0 && $expiry < $this->getCurrentTime() );

        if ( !$this->lock( $key, 0 ) ) {
            return false;
        }
        // Use doGet() to avoid having to trigger resolveSegments()
        $blob = $this->doGet( $key, self::READ_LATEST );
        if ( $blob ) {
            if ( $delete ) {
                $ok = $this->doDelete( $key, $flags );
            } else {
                $ok = $this->doSet( $key, $blob, $exptime, $flags );
            }
        } else {
            $ok = false;
        }

        $this->unlock( $key );

        return $ok;
    }

    public function lock( $key, $timeout = 6, $expiry = 6, $rclass = '' ) {
        // Avoid deadlocks and allow lock reentry if specified
        if ( isset( $this->locks[$key] ) ) {
            if ( $rclass != '' && $this->locks[$key]['class'] === $rclass ) {
                ++$this->locks[$key]['depth'];
                return true;
            } else {
                return false;
            }
        }

        $fname = __METHOD__;
        $expiry = min( $expiry ?: INF, self::TTL_DAY );
        $loop = new WaitConditionLoop(
            function () use ( $key, $expiry, $fname ) {
                $this->clearLastError();
                if ( $this->add( "{$key}:lock", 1, $expiry ) ) {
                    return WaitConditionLoop::CONDITION_REACHED; // locked!
                } elseif ( $this->getLastError() ) {
                    $this->logger->warning(
                        $fname . ' failed due to I/O error for {key}.',
                        [ 'key' => $key ]
                    );

                    return WaitConditionLoop::CONDITION_ABORTED; // network partition?
                }

                return WaitConditionLoop::CONDITION_CONTINUE;
            },
            $timeout
        );

        $code = $loop->invoke();
        $locked = ( $code === $loop::CONDITION_REACHED );
        if ( $locked ) {
            $this->locks[$key] = [ 'class' => $rclass, 'depth' => 1 ];
        } elseif ( $code === $loop::CONDITION_TIMED_OUT ) {
            $this->logger->warning(
                "$fname failed due to timeout for {key}.",
                [ 'key' => $key, 'timeout' => $timeout ]
            );
        }

        return $locked;
    }

    public function unlock( $key ) {
        if ( isset( $this->locks[$key] ) && --$this->locks[$key]['depth'] <= 0 ) {
            unset( $this->locks[$key] );

            $ok = $this->doDelete( "{$key}:lock" );
            if ( !$ok ) {
                $this->logger->warning(
                    __METHOD__ . ' failed to release lock for {key}.',
                    [ 'key' => $key ]
                );
            }

            return $ok;
        }

        return true;
    }

    final public function getScopedLock( $key, $timeout = 6, $expiry = 30, $rclass = '' ) {
        $expiry = min( $expiry ?: INF, self::TTL_DAY );

        if ( !$this->lock( $key, $timeout, $expiry, $rclass ) ) {
            return null;
        }

        $lSince = $this->getCurrentTime(); // lock timestamp

        return new ScopedCallback( function () use ( $key, $lSince, $expiry ) {
            $latency = 0.050; // latency skew (err towards keeping lock present)
            $age = ( $this->getCurrentTime() - $lSince + $latency );
            if ( ( $age + $latency ) >= $expiry ) {
                $this->logger->warning(
                    "Lock for {key} held too long ({age} sec).",
                    [ 'key' => $key, 'age' => $age ]
                );
                return; // expired; it's not "safe" to delete the key
            }
            $this->unlock( $key );
        } );
    }

    public function deleteObjectsExpiringBefore( $date, $progressCallback = false ) {
        // stub
        return false;
    }

    public function getMulti( array $keys, $flags = 0 ) {
        $valuesBykey = $this->doGetMulti( $keys, $flags );
        foreach ( $valuesBykey as $key => $value ) {
            // Resolve one blob at a time (avoids too much I/O at once)
            $valuesBykey[$key] = $this->resolveSegments( $key, $value );
        }

        return $valuesBykey;
    }

    protected function doGetMulti( array $keys, $flags = 0 ) {
        $res = [];
        foreach ( $keys as $key ) {
            $val = $this->doGet( $key, $flags );
            if ( $val !== false ) {
                $res[$key] = $val;
            }
        }

        return $res;
    }

    public function setMulti( array $data, $exptime = 0, $flags = 0 ) {
        if ( ( $flags & self::WRITE_ALLOW_SEGMENTS ) === self::WRITE_ALLOW_SEGMENTS ) {
            throw new InvalidArgumentException( __METHOD__ . ' got WRITE_ALLOW_SEGMENTS' );
        }

        $res = true;
        foreach ( $data as $key => $value ) {
            $res = $this->doSet( $key, $value, $exptime, $flags ) && $res;
        }

        return $res;
    }

    public function deleteMulti( array $keys, $flags = 0 ) {
        $res = true;
        foreach ( $keys as $key ) {
            $res = $this->doDelete( $key, $flags ) && $res;
        }

        return $res;
    }

    abstract public function incr( $key, $value = 1 );

    public function decr( $key, $value = 1 ) {
        return $this->incr( $key, - $value );
    }

    public function incrWithInit( $key, $ttl, $value = 1, $init = 1 ) {
        $this->clearLastError();
        $newValue = $this->incr( $key, $value );
        if ( $newValue === false && !$this->getLastError() ) {
            // No key set; initialize
            $newValue = $this->add( $key, (int)$init, $ttl ) ? $init : false;
            if ( $newValue === false && !$this->getLastError() ) {
                // Raced out initializing; increment
                $newValue = $this->incr( $key, $value );
            }
        }

        return $newValue;
    }

    protected function resolveSegments( $key, $mainValue ) {
        if ( SerializedValueContainer::isUnified( $mainValue ) ) {
            return $this->unserialize( $mainValue->{SerializedValueContainer::UNIFIED_DATA} );
        }

        if ( SerializedValueContainer::isSegmented( $mainValue ) ) {
            $orderedKeys = array_map(
                function ( $segmentHash ) use ( $key ) {
                    return $this->makeGlobalKey( self::SEGMENT_COMPONENT, $key, $segmentHash );
                },
                $mainValue->{SerializedValueContainer::SEGMENTED_HASHES}
            );

            $segmentsByKey = $this->doGetMulti( $orderedKeys );

            $parts = [];
            foreach ( $orderedKeys as $segmentKey ) {
                if ( isset( $segmentsByKey[$segmentKey] ) ) {
                    $parts[] = $segmentsByKey[$segmentKey];
                } else {
                    return false; // missing segment
                }
            }

            return $this->unserialize( implode( '', $parts ) );
        }

        return $mainValue;
    }

    public function getLastError() {
        return $this->lastError;
    }

    public function clearLastError() {
        $this->lastError = self::ERR_NONE;
    }

    protected function setLastError( $err ) {
        $this->lastError = $err;
    }

    public function addBusyCallback( callable $workCallback ) {
        $this->busyCallbacks[] = $workCallback;
    }

    protected function debug( $text ) {
        if ( $this->debugMode ) {
            $this->logger->debug( "{class} debug: $text", [
                'class' => static::class,
            ] );
        }
    }

    protected function expiryIsRelative( $exptime ) {
        return ( $exptime != 0 && $exptime < ( 10 * self::TTL_YEAR ) );
    }

    protected function convertToExpiry( $exptime ) {
        if ( $this->expiryIsRelative( $exptime ) ) {
            return (int)$this->getCurrentTime() + $exptime;
        } else {
            return $exptime;
        }
    }

    protected function convertToRelative( $exptime ) {
        if ( $exptime >= ( 10 * self::TTL_YEAR ) ) {
            $exptime -= (int)$this->getCurrentTime();
            if ( $exptime <= 0 ) {
                $exptime = 1;
            }
            return $exptime;
        } else {
            return $exptime;
        }
    }

    protected function isInteger( $value ) {
        if ( is_int( $value ) ) {
            return true;
        } elseif ( !is_string( $value ) ) {
            return false;
        }

        $integer = (int)$value;

        return ( $value === (string)$integer );
    }

    public function makeKeyInternal( $keyspace, $args ) {
        $key = $keyspace;
        foreach ( $args as $arg ) {
            $key .= ':' . str_replace( ':', '%3A', $arg );
        }
        return strtr( $key, ' ', '_' );
    }

    public function makeGlobalKey( $class, $component = null ) {
        return $this->makeKeyInternal( 'global', func_get_args() );
    }

    public function makeKey( $class, $component = null ) {
        return $this->makeKeyInternal( $this->keyspace, func_get_args() );
    }

    public function getQoS( $flag ) {
        return $this->attrMap[$flag] ?? self::QOS_UNKNOWN;
    }

    public function getSegmentationSize() {
        return $this->segmentationSize;
    }

    public function getSegmentedValueMaxSize() {
        return $this->segmentedValueMaxSize;
    }

    protected function mergeFlagMaps( array $bags ) {
        $map = [];
        foreach ( $bags as $bag ) {
            foreach ( $bag->attrMap as $attr => $rank ) {
                if ( isset( $map[$attr] ) ) {
                    $map[$attr] = min( $map[$attr], $rank );
                } else {
                    $map[$attr] = $rank;
                }
            }
        }

        return $map;
    }

    public function getCurrentTime() {
        return $this->wallClockOverride ?: microtime( true );
    }

    public function setMockTime( &$time ) {
        $this->wallClockOverride =& $time;
    }

    protected function serialize( $value ) {
        return is_int( $value ) ? $value : serialize( $value );
    }

    protected function unserialize( $value ) {
        return $this->isInteger( $value ) ? (int)$value : unserialize( $value );
    }
}