ch_part2_tokens_impl.py

class ConsistentHashingWithTokens:
    def __init__(
        self,
        server_tokens: Dict[ServerId, int],
        seed: int = SEED,
    ):
        assert all(tokens >= 1 for tokens in server_tokens.values())
        self._seed = seed
        self._server_tokens = server_tokens
        self._allocate_servers()
        
    def key_lookup(self, key: str) -> ServerToken:
        key_hash = ServerToken(
            server_id=key,
            token_idx=0,  # dummy value
            _seed=self._seed,
        ).hash_value
        distances = (
            (
                key_hash - server_token_obj.hash_value,
                server_id,
            )
            for (server_id, _), server_token_obj
            in self._server_allocations.items()
        )
        return self._find_closest(distances=distances)
    
    def _find_closest(
        self,
        distances: Iterable[Tuple[int, ServerId]],
    ) -> ServerId:
        # find the server with the smallest positive distance
        # assign None if no such server exists
        # which means that the server with the largest hash is
        # responsible for this key
        closest = min((
            (dist, server_id)
            for dist, server_id in distances
            if dist >= 0
        ), default=None)
        return (
            # get the ID of the server with the largest hash value
            max(self._server_allocations.values()).server_id
            if closest is None
            else closest[1]
        )
        
    def _allocate_servers(self):
        self._server_allocations = {
            (server_id, token_idx): ServerToken(
                server_id=server_id,
                token_idx=token_idx,
                _seed=self._seed,
            )
            for server_id, token_count in self._server_tokens.items()
            for token_idx in range(token_count)
        }