terrassumptions/server/app/point_gen/__init__.py

import asyncio
import collections
import logging
import random
from typing import List, Tuple, Dict, Union

from .random_street_view import call_random_street_view, VALID_COUNTRIES as RSV_COUNTRIES
from .urban_centers import urban_coord, VALID_COUNTRIES as URBAN_COUNTRIES
from .shared import aiohttp_client

from ..schemas import GameConfig, GenMethodEnum, CountryCode, CacheInfo, GeneratorInfo

logger = logging.getLogger(__name__)

generator_info = [
    GeneratorInfo(
        generation_method=GenMethodEnum.rsv,
        country_locks=RSV_COUNTRIES
    ),
    GeneratorInfo(
        generation_method=GenMethodEnum.urban,
        country_locks=URBAN_COUNTRIES
    ),
]

cache_names = {
    GenMethodEnum.rsv: "RSV",
    GenMethodEnum.urban: "Urban",
}


class ExhaustedSourceError(Exception):
    pass


class PointStore:
    def __init__(self, 
                 cache_targets: Dict[Tuple[GenMethodEnum, CountryCode], int],
                 rsv_country_retries: int = 5,
                 urban_country_pool_size: int = 30,
                 urban_country_retries: int = 30,
                 urban_city_retries: int = 50):
        self.cache_targets = cache_targets
        self.rsv_country_retries = rsv_country_retries
        self.urban_country_pool_size = urban_country_pool_size
        self.urban_country_retries = urban_country_retries
        self.urban_city_retries = urban_city_retries
        self.store = collections.defaultdict(collections.deque)

    async def get_point(self, generator: GenMethodEnum, country: Union[CountryCode, None], force_generate: bool = False) -> Tuple[str, float, float]:
        if country is None:
            # generating points across the whole world
            # for current generators, this means selecting a country at random
            if generator == GenMethodEnum.rsv:
                for _ in range(self.rsv_country_retries):
                    # try a few countries before giving up, just in case one has no data
                    country = random.choice(RSV_COUNTRIES)
                    # RSV point function returns a collection of points, which should be cached
                    points = await call_random_street_view(country)
                    if len(points) > 0:
                        point = points.pop()
                        self.store[(generator, country)].extend(points)
                        return point
                else:
                    raise ExhaustedSourceError
            elif generator == GenMethodEnum.urban:
                # try many countries since finding an urban center point is harder
                countries = random.sample(URBAN_COUNTRIES, k=min(self.urban_country_pool_size, len(URBAN_COUNTRIES)))
                for country in countries:
                    logger.info(f"Selecting urban centers from {country}")
                    pt = await urban_coord(country)
                    if pt is not None:
                        return pt
                else:
                    raise ExhaustedSourceError
            else:
                raise ExhaustedSourceError
        else:
            # generating points for a specific country
            # if we already have a point ready, just return it immediately
            if not force_generate:
                stock = self.store[(generator, country)]
                if len(stock) > 0:
                    return stock.popleft()
            
            # otherwise, need to actually generate a new point
            if generator == GenMethodEnum.rsv:
                # RSV point function returns a collection of points, which should be cached
                points = await call_random_street_view(country)
                if len(points) == 0:
                    raise ExhaustedSourceError
                point = points.pop()
                self.store[(generator, country)].extend(points)
                return point
            elif generator == GenMethodEnum.urban:
                for i in range(self.urban_country_retries):
                    logger.info(f"Attempt #{i + 1} to select urban centers from {country}")
                    pt = await urban_coord(country, city_retries=self.urban_city_retries)
                    if pt is not None:
                        return pt
                else:
                    raise ExhaustedSourceError
            else:
                raise ExhaustedSourceError

    async def get_points(self, config: GameConfig) -> List[Tuple[str, float, float]]:
        """
        Provide points according to the GameConfig.

        Return a list of valid geo points, as 
        (2 character country code, latitude, longitude) tuples.

        In the event that the configured source cannot reasonably supply enough points,
        most likely due to time constraints, this will raise an ExhaustedSourceError.
        """
        try:
            point_tasks = [self.get_point(config.generation_method, config.country_lock) for _ in range(config.rounds)]
            gathered = asyncio.gather(*point_tasks)
            return await asyncio.wait_for(gathered, 60)
        except asyncio.TimeoutError:
            raise ExhaustedSourceError

    def get_cache_info(self) -> List[CacheInfo]:
        """
        Get CacheInfo for all caches.
        """
        return [CacheInfo(cache_name=f"{cache_names[g]}-{c}", size=len(ps)) for (g, c), ps in self.store.items()]

    async def _restock_source_impl(self, generator: GenMethodEnum, country: CountryCode):
        key = (generator, country)
        target = self.cache_targets.get(key, 0)
        stock = self.store[key]
        while len(stock) < target:  # this check allows for RSV to do its multi-point restock
            stock.append(await self.get_point(*key, force_generate=True))

    async def restock_source(self, config: GameConfig):
        """
        Restock any caches associated with the GameConfig.
        """
        if config.country_lock is None:
            return
        try:
            await self._restock_source_impl(config.generation_method, config.country_lock)
        except ExhaustedSourceError:
            # if the cache can't be restocked, that is bad, but not fatal
            logger.exception(f"Failed to fully restock point cache for {config}")

    async def restock_all(self, timeout: Union[int, float, None] = None):
        """
        Restock all caches.
        """
        restock_tasks = [self._restock_source_impl(gen, cc) for (gen, cc) in self.cache_targets.keys()]
        gathered = asyncio.gather(*restock_tasks)
        try:
            await asyncio.wait_for(gathered, timeout)
        except (asyncio.TimeoutError, ExhaustedSourceError):
            # if this task times out, it's fine, as it's just intended to be a best effort
            logger.exception(f"Failed to fully restock point cache for {config}")


points = PointStore({
    (GenMethodEnum.urban, "us"): 10,
})