HashRing.php

Go to the documentation of this file.

<?php
namespace Wikimedia\HashRing;
 
use InvalidArgumentException;
use RuntimeException;
use UnexpectedValueException;
 
class HashRing {
    protected $algo;
    protected $weightByLocation;
    protected $ejectExpiryByLocation;
 
    protected $baseRing;
    protected $liveRing;
 
    private const HASHES_PER_LOCATION = 40;
    private const SECTORS_PER_HASH = 4;
 
    public const KEY_POS = 0;
    public const KEY_LOCATION = 1;
 
    public const RING_ALL = 0;
    public const RING_LIVE = 1;
 
    public function __construct( array $map, $algo = 'sha1', array $ejections = [] ) {
        $this->init( $map, $algo, $ejections );
    }
 
    protected function init( array $map, $algo, array $ejections ) {
        if ( !in_array( $algo, hash_algos(), true ) ) {
            throw new RuntimeException( __METHOD__ . ": unsupported '$algo' hash algorithm." );
        }
 
        $weightByLocation = array_filter( $map );
        if ( $weightByLocation === [] ) {
            throw new UnexpectedValueException( "No locations with non-zero weight." );
        } elseif ( min( $map ) < 0 ) {
            throw new InvalidArgumentException( "Location weight cannot be negative." );
        }
 
        $this->algo = $algo;
        $this->weightByLocation = $weightByLocation;
        $this->ejectExpiryByLocation = $ejections;
        $this->baseRing = $this->buildLocationRing( $this->weightByLocation );
    }
 
    final public function getLocation( $item ) {
        return $this->getLocations( $item, 1 )[0];
    }
 
    public function getLocations( $item, $limit, $from = self::RING_ALL ) {
        if ( $from === self::RING_ALL ) {
            $ring = $this->baseRing;
        } elseif ( $from === self::RING_LIVE ) {
            $ring = $this->getLiveRing();
        } else {
            throw new InvalidArgumentException( "Invalid ring source specified." );
        }
 
        // Short-circuit for the common single-location case. Note that if there was only one
        // location and it was ejected from the live ring, getLiveRing() would have error out.
        if ( count( $this->weightByLocation ) == 1 ) {
            return ( $limit > 0 ) ? [ $ring[0][self::KEY_LOCATION] ] : [];
        }
 
        // Locate the node index for this item's position on the hash ring
        $itemIndex = $this->findNodeIndexForPosition( $this->getItemPosition( $item ), $ring );
 
        $locations = [];
        $currentIndex = null;
        while ( count( $locations ) < $limit ) {
            if ( $currentIndex === null ) {
                $currentIndex = $itemIndex;
            } else {
                $currentIndex = $this->getNextClockwiseNodeIndex( $currentIndex, $ring );
                if ( $currentIndex === $itemIndex ) {
                    break; // all nodes visited
                }
            }
            // @phan-suppress-next-line PhanTypeMismatchDimFetchNullable False positive
            $nodeLocation = $ring[$currentIndex][self::KEY_LOCATION];
            if ( !in_array( $nodeLocation, $locations, true ) ) {
                // Ignore other nodes for the same locations already added
                $locations[] = $nodeLocation;
            }
        }
 
        return $locations;
    }
 
    private function findNodeIndexForPosition( $position, $ring ) {
        $count = count( $ring );
        if ( $count === 0 ) {
            return null;
        }
 
        $index = null;
        $lowPos = 0;
        $highPos = $count;
        while ( true ) {
            $midPos = (int)( ( $lowPos + $highPos ) / 2 );
            if ( $midPos === $count ) {
                $index = 0;
                break;
            }
 
            $midVal = $ring[$midPos][self::KEY_POS];
            $midMinusOneVal = ( $midPos === 0 ) ? 0 : $ring[$midPos - 1][self::KEY_POS];
            if ( $position <= $midVal && $position > $midMinusOneVal ) {
                $index = $midPos;
                break;
            }
 
            if ( $midVal < $position ) {
                $lowPos = $midPos + 1;
            } else {
                $highPos = $midPos - 1;
            }
 
            if ( $lowPos > $highPos ) {
                $index = 0;
                break;
            }
        }
 
        return $index;
    }
 
    public function getLocationWeights() {
        return $this->weightByLocation;
    }
 
    public function ejectFromLiveRing( $location, $ttl ) {
        if ( !isset( $this->weightByLocation[$location] ) ) {
            throw new UnexpectedValueException( "No location '$location' in the ring." );
        }
 
        $expiry = $this->getCurrentTime() + $ttl;
        $this->ejectExpiryByLocation[$location] = $expiry;
 
        $this->liveRing = null; // invalidate ring cache
 
        return ( count( $this->ejectExpiryByLocation ) < count( $this->weightByLocation ) );
    }
 
    final public function getLiveLocation( $item ) {
        return $this->getLocations( $item, 1, self::RING_LIVE )[0];
    }
 
    final public function getLiveLocations( $item, $limit ) {
        return $this->getLocations( $item, $limit, self::RING_LIVE );
    }
 
    public function getLiveLocationWeights() {
        $now = $this->getCurrentTime();
 
        return array_diff_key(
            $this->weightByLocation,
            array_filter(
                $this->ejectExpiryByLocation,
                static function ( $expiry ) use ( $now ) {
                    return ( $expiry > $now );
                }
            )
        );
    }
 
    private function buildLocationRing( array $weightByLocation ) {
        $locationCount = count( $weightByLocation );
        $totalWeight = array_sum( $weightByLocation );
 
        $ring = [];
        // Assign nodes to all locations based on location weight
        $claimed = []; // (position as string => (node, index))
        foreach ( $weightByLocation as $location => $weight ) {
            $ratio = $weight / $totalWeight;
            // There $locationCount * (HASHES_PER_LOCATION * 4) nodes available;
            // assign a few groups of nodes to this location based on its weight.
            $nodesQuartets = intval( $ratio * self::HASHES_PER_LOCATION * $locationCount );
            for ( $qi = 0; $qi < $nodesQuartets; ++$qi ) {
                // For efficiency, get 4 points per hash call and 4X node count.
                // If $algo is MD5, then this matches that of with libketama.
                // See https://github.com/RJ/ketama/blob/master/libketama/ketama.c
                $positions = $this->getNodePositionQuartet( "{$location}-{$qi}" );
                foreach ( $positions as $gi => $position ) {
                    $node = ( $qi * self::SECTORS_PER_HASH + $gi ) . "@$location";
                    $posKey = (string)$position; // large integer
                    if ( isset( $claimed[$posKey] ) ) {
                        // Disallow duplicates  (name decides precedence)
                        if ( $claimed[$posKey]['node'] > $node ) {
                            continue;
                        } else {
                            unset( $ring[$claimed[$posKey]['index']] );
                        }
                    }
                    $ring[] = [
                        self::KEY_POS => $position,
                        self::KEY_LOCATION => $location
                    ];
                    $claimed[$posKey] = [ 'node' => $node, 'index' => count( $ring ) - 1 ];
                }
            }
        }
        // Sort the locations into clockwise order based on the hash ring position
        usort( $ring, static function ( $a, $b ) {
            if ( $a[self::KEY_POS] === $b[self::KEY_POS] ) {
                throw new UnexpectedValueException( 'Duplicate node positions.' );
            }
 
            return ( $a[self::KEY_POS] < $b[self::KEY_POS] ? -1 : 1 );
        } );
 
        return $ring;
    }
 
    private function getItemPosition( $item ) {
        // If $algo is MD5, then this matches that of with libketama.
        // See https://github.com/RJ/ketama/blob/master/libketama/ketama.c
        $octets = substr( hash( $this->algo, (string)$item, true ), 0, 4 );
        if ( strlen( $octets ) != 4 ) {
            throw new UnexpectedValueException( __METHOD__ . ": {$this->algo} is < 32 bits." );
        }
 
        $pos = unpack( 'V', $octets )[1];
        if ( $pos < 0 ) {
            // Most-significant-bit is set, causing unpack() to return a negative integer due
            // to the fact that it returns a signed int. Cast it to an unsigned integer string.
            $pos = sprintf( '%u', $pos );
        }
 
        return (float)$pos;
    }
 
    private function getNodePositionQuartet( $nodeGroupName ) {
        $octets = substr( hash( $this->algo, (string)$nodeGroupName, true ), 0, 16 );
        if ( strlen( $octets ) != 16 ) {
            throw new UnexpectedValueException( __METHOD__ . ": {$this->algo} is < 128 bits." );
        }
 
        $positions = [];
        foreach ( unpack( 'V4', $octets ) as $signed ) {
            $positions[] = (float)sprintf( '%u', $signed );
        }
 
        return $positions;
    }
 
    private function getNextClockwiseNodeIndex( $i, $ring ) {
        if ( !isset( $ring[$i] ) ) {
            throw new UnexpectedValueException( __METHOD__ . ": reference index is invalid." );
        }
 
        $next = $i + 1;
 
        return ( $next < count( $ring ) ) ? $next : 0;
    }
 
    protected function getLiveRing() {
        if ( !$this->ejectExpiryByLocation ) {
            return $this->baseRing; // nothing ejected
        }
 
        $now = $this->getCurrentTime();
 
        if ( $this->liveRing === null || min( $this->ejectExpiryByLocation ) <= $now ) {
            // Live ring needs to be regenerated...
            $this->ejectExpiryByLocation = array_filter(
                $this->ejectExpiryByLocation,
                static function ( $expiry ) use ( $now ) {
                    return ( $expiry > $now );
                }
            );
 
            if ( count( $this->ejectExpiryByLocation ) ) {
                // Some locations are still ejected from the ring
                $liveRing = [];
                foreach ( $this->baseRing as $nodeInfo ) {
                    $location = $nodeInfo[self::KEY_LOCATION];
                    if ( !isset( $this->ejectExpiryByLocation[$location] ) ) {
                        $liveRing[] = $nodeInfo;
                    }
                }
            } else {
                $liveRing = $this->baseRing;
            }
 
            $this->liveRing = $liveRing;
        }
 
        if ( !$this->liveRing ) {
            throw new UnexpectedValueException( "The live ring is currently empty." );
        }
 
        return $this->liveRing;
    }
 
    protected function getCurrentTime() {
        return time();
    }
 
    public function __serialize() {
        return [
            'algorithm' => $this->algo,
            'locations' => $this->weightByLocation,
            'ejections' => $this->ejectExpiryByLocation
        ];
    }
 
    public function __unserialize( $data ) {
        if ( is_array( $data ) ) {
            $this->init( $data['locations'] ?? [], $data['algorithm'] ?? 'sha1', $data['ejections'] ?? [] );
        } else {
            throw new UnexpectedValueException( __METHOD__ . ": unable to decode JSON." );
        }
    }
}
 
class_alias( HashRing::class, 'HashRing' );