NameTableStore.php

Go to the documentation of this file.

<?php
namespace MediaWiki\Storage;
 
use Exception;
use IExpiringStore;
use Psr\Log\LoggerInterface;
use WANObjectCache;
use Wikimedia\Assert\Assert;
use Wikimedia\Rdbms\Database;
use Wikimedia\Rdbms\IDatabase;
use Wikimedia\Rdbms\ILoadBalancer;
use Wikimedia\RequestTimeout\TimeoutException;
 
class NameTableStore {
 
    private $loadBalancer;
 
    private $cache;
 
    private $logger;
 
    private $tableCache = null;
 
    private $domain;
 
    private $cacheTTL;
 
    private $table;
    private $idField;
    private $nameField;
    private $normalizationCallback;
    private $insertCallback;
 
    public function __construct(
        ILoadBalancer $dbLoadBalancer,
        WANObjectCache $cache,
        LoggerInterface $logger,
        $table,
        $idField,
        $nameField,
        callable $normalizationCallback = null,
        $dbDomain = false,
        callable $insertCallback = null
    ) {
        $this->loadBalancer = $dbLoadBalancer;
        $this->cache = $cache;
        $this->logger = $logger;
        $this->table = $table;
        $this->idField = $idField;
        $this->nameField = $nameField;
        $this->normalizationCallback = $normalizationCallback;
        $this->domain = $dbDomain;
        $this->cacheTTL = IExpiringStore::TTL_MONTH;
        $this->insertCallback = $insertCallback;
    }
 
    private function getDBConnection( $index, $flags = 0 ) {
        return $this->loadBalancer->getConnectionRef( $index, [], $this->domain, $flags );
    }
 
    private function getCacheKey() {
        return $this->cache->makeGlobalKey(
            'NameTableSqlStore',
            $this->table,
            $this->loadBalancer->resolveDomainID( $this->domain )
        );
    }
 
    private function normalizeName( $name ) {
        if ( $this->normalizationCallback === null ) {
            return $name;
        }
        return call_user_func( $this->normalizationCallback, $name );
    }
 
    public function acquireId( string $name ) {
        $name = $this->normalizeName( $name );
 
        $table = $this->getTableFromCachesOrReplica();
        $searchResult = array_search( $name, $table, true );
        if ( $searchResult === false ) {
            $id = $this->store( $name );
            if ( $id === null ) {
                // RACE: $name was already in the db, probably just inserted, so load from primary DB.
                // Use DBO_TRX to avoid missing inserts due to other threads or REPEATABLE-READs.
                $table = $this->reloadMap( ILoadBalancer::CONN_TRX_AUTOCOMMIT );
 
                $searchResult = array_search( $name, $table, true );
                if ( $searchResult === false ) {
                    // Insert failed due to IGNORE flag, but DB_PRIMARY didn't give us the data
                    $m = "No insert possible but primary DB didn't give us a record for " .
                        "'{$name}' in '{$this->table}'";
                    $this->logger->error( $m );
                    throw new NameTableAccessException( $m );
                }
            } else {
                if ( isset( $table[$id] ) ) {
                    // This can happen when a transaction is rolled back and acquireId is called in
                    // an onTransactionResolution() callback, which gets executed before retryStore()
                    // has a chance to run. The right thing to do in this case is to discard the old
                    // value. According to the contract of acquireId, the caller should not have
                    // used it outside the transaction, so it should not be persisted anywhere after
                    // the rollback.
                    $m = "Got ID $id for '$name' from insert"
                        . " into '{$this->table}', but ID $id was previously associated with"
                        . " the name '{$table[$id]}'. Overriding the old value, which presumably"
                        . " has been removed from the database due to a transaction rollback.";
 
                    $this->logger->warning( $m );
                }
 
                $table[$id] = $name;
                $searchResult = $id;
 
                // As store returned an ID we know we inserted so delete from WAN cache
                $dbw = $this->getDBConnection( DB_PRIMARY );
                $dbw->onTransactionPreCommitOrIdle( function () {
                    $this->cache->delete( $this->getCacheKey() );
                }, __METHOD__ );
            }
            $this->tableCache = $table;
        }
 
        return $searchResult;
    }
 
    public function reloadMap( $connFlags = 0 ) {
        if ( $connFlags !== 0 && defined( 'MW_PHPUNIT_TEST' ) ) {
            // HACK: We can't use $connFlags while doing PHPUnit tests, because the
            // fake database tables are bound to a single connection.
            $connFlags = 0;
        }
 
        $dbw = $this->getDBConnection( DB_PRIMARY, $connFlags );
        $this->tableCache = $this->loadTable( $dbw );
        $dbw->onTransactionPreCommitOrIdle( function () {
            $this->cache->reap( $this->getCacheKey(), INF );
        }, __METHOD__ );
 
        return $this->tableCache;
    }
 
    public function getId( string $name ) {
        $name = $this->normalizeName( $name );
 
        $table = $this->getTableFromCachesOrReplica();
        $searchResult = array_search( $name, $table, true );
 
        if ( $searchResult !== false ) {
            return $searchResult;
        }
 
        throw NameTableAccessException::newFromDetails( $this->table, 'name', $name );
    }
 
    public function getName( int $id ) {
        $table = $this->getTableFromCachesOrReplica();
        if ( array_key_exists( $id, $table ) ) {
            return $table[$id];
        }
        $fname = __METHOD__;
 
        $table = $this->cache->getWithSetCallback(
            $this->getCacheKey(),
            $this->cacheTTL,
            function ( $oldValue, &$ttl, &$setOpts ) use ( $id, $fname ) {
                // Check if cached value is up-to-date enough to have $id
                if ( is_array( $oldValue ) && array_key_exists( $id, $oldValue ) ) {
                    // Completely leave the cache key alone
                    $ttl = WANObjectCache::TTL_UNCACHEABLE;
                    // Use the old value
                    return $oldValue;
                }
                // Regenerate from replica DB, and primary DB if needed
                foreach ( [ DB_REPLICA, DB_PRIMARY ] as $source ) {
                    // Log a fallback to primary
                    if ( $source === DB_PRIMARY ) {
                        $this->logger->info(
                            $fname . ' falling back to primary select from ' .
                            $this->table . ' with id ' . $id
                        );
                    }
                    $db = $this->getDBConnection( $source );
                    $cacheSetOpts = Database::getCacheSetOptions( $db );
                    $table = $this->loadTable( $db );
                    if ( array_key_exists( $id, $table ) ) {
                        break; // found it
                    }
                }
                // Use the value from last source checked
                $setOpts += $cacheSetOpts;
 
                return $table;
            },
            [ 'minAsOf' => INF ] // force callback run
        );
 
        $this->tableCache = $table;
 
        if ( array_key_exists( $id, $table ) ) {
            return $table[$id];
        }
 
        throw NameTableAccessException::newFromDetails( $this->table, 'id', $id );
    }
 
    public function getMap() {
        return $this->getTableFromCachesOrReplica();
    }
 
    private function getTableFromCachesOrReplica() {
        if ( $this->tableCache !== null ) {
            return $this->tableCache;
        }
 
        $table = $this->cache->getWithSetCallback(
            $this->getCacheKey(),
            $this->cacheTTL,
            function ( $oldValue, &$ttl, &$setOpts ) {
                $dbr = $this->getDBConnection( DB_REPLICA );
                $setOpts += Database::getCacheSetOptions( $dbr );
                return $this->loadTable( $dbr );
            }
        );
 
        $this->tableCache = $table;
 
        return $table;
    }
 
    private function loadTable( IDatabase $db ) {
        $result = $db->select(
            $this->table,
            [
                'id' => $this->idField,
                'name' => $this->nameField
            ],
            [],
            __METHOD__,
            [ 'ORDER BY' => 'id' ]
        );
 
        $assocArray = [];
        foreach ( $result as $row ) {
            $assocArray[$row->id] = $row->name;
        }
 
        return $assocArray;
    }
 
    private function store( string $name ) {
        Assert::parameter( $name !== '', '$name', 'should not be an empty string' );
        // Note: this is only called internally so normalization of $name has already occurred.
 
        $dbw = $this->getDBConnection( DB_PRIMARY );
 
        $id = null;
        $dbw->doAtomicSection(
            __METHOD__,
            function ( IDatabase $unused, $fname )
            use ( $name, &$id, $dbw ) {
                // NOTE: use IDatabase from the parent scope here, not the function parameter.
                // If $dbw is a wrapper around the actual DB, we need to call the wrapper here,
                // not the inner instance.
                $dbw->insert(
                    $this->table,
                    $this->getFieldsToStore( $name ),
                    $fname,
                    [ 'IGNORE' ]
                );
 
                if ( $dbw->affectedRows() === 0 ) {
                    $this->logger->info(
                        'Tried to insert name into table ' . $this->table . ', but value already existed.'
                    );
 
                    return;
                }
 
                $id = $dbw->insertId();
 
                // Any open transaction may still be rolled back. If that happens, we have to re-try the
                // insertion and restore a consistent state of the cached table.
                $dbw->onAtomicSectionCancel(
                    function ( $trigger, IDatabase $unused ) use ( $name, $id, $dbw ) {
                        $this->retryStore( $dbw, $name, $id );
                    },
                $fname );
            },
            IDatabase::ATOMIC_CANCELABLE
        );
 
        return $id;
    }
 
    private function retryStore( IDatabase $dbw, $name, $id ) {
        // NOTE: in the closure below, use the IDatabase from the original method call,
        // not the one passed to the closure as a parameter.
        // If $dbw is a wrapper around the actual DB, we need to call the wrapper,
        // not the inner instance.
 
        try {
            $dbw->doAtomicSection(
                __METHOD__,
                function ( IDatabase $unused, $fname ) use ( $name, $id, $dbw ) {
                    // Try to insert a row with the ID we originally got.
                    // If that fails (because of a key conflict), we will just try to get another ID again later.
                    $dbw->insert(
                        $this->table,
                        $this->getFieldsToStore( $name, $id ),
                        $fname
                    );
 
                    // Make sure we re-load the map in case this gets rolled back again.
                    // We could re-try once more, but that bears the risk of an infinite loop.
                    // So let's just give up on the ID.
                    $dbw->onAtomicSectionCancel(
                        function ( $trigger, IDatabase $unused ) {
                            $this->logger->warning(
                                'Re-insertion of name into table ' . $this->table
                                . ' was rolled back. Giving up and reloading the cache.'
                            );
                            $this->reloadMap( ILoadBalancer::CONN_TRX_AUTOCOMMIT );
                        },
                        $fname
                    );
 
                    $this->logger->info(
                        'Re-insert name into table ' . $this->table . ' after failed transaction.'
                    );
                },
                IDatabase::ATOMIC_CANCELABLE
            );
        } catch ( TimeoutException $e ) {
            throw $e;
        } catch ( Exception $ex ) {
            $this->logger->error(
                'Re-insertion of name into table ' . $this->table . ' failed: ' . $ex->getMessage()
            );
        } finally {
            // NOTE: we reload regardless of whether the above insert succeeded. There is
            // only three possibilities: the insert succeeded, so the new map will have
            // the desired $id/$name mapping. Or the insert failed because another
            // process already inserted that same $id/$name mapping, in which case the
            // new map will also have it. Or another process grabbed the desired ID for
            // another name, or the database refuses to insert the given ID into the
            // auto increment field - in that case, the new map will not have a mapping
            // for $name (or has a different mapping for $name). In that last case, we can
            // only hope that the ID produced within the failed transaction has not been
            // used outside that transaction.
 
            $this->reloadMap( ILoadBalancer::CONN_TRX_AUTOCOMMIT );
        }
    }
 
    private function getFieldsToStore( $name, $id = null ) {
        $fields = [];
 
        $fields[$this->nameField] = $name;
 
        if ( $id !== null ) {
            $fields[$this->idField] = $id;
        }
 
        if ( $this->insertCallback !== null ) {
            $fields = call_user_func( $this->insertCallback, $fields );
        }
        return $fields;
    }
 
}