Batch scans to prevent auto DDOS

This commit is contained in:
Zoe Roux 2023-04-05 01:02:40 +09:00
parent bb716ab6b6
commit 8f22785d2f
2 changed files with 21 additions and 2 deletions

View File

@ -8,7 +8,7 @@ from guessit import guessit
from providers.provider import Provider from providers.provider import Provider
from providers.types.episode import Episode, PartialShow from providers.types.episode import Episode, PartialShow
from providers.types.season import Season, SeasonTranslation from providers.types.season import Season, SeasonTranslation
from .utils import log_errors, provider_cache, set_in_cache from .utils import batch, log_errors, provider_cache, set_in_cache
class Scanner: class Scanner:
@ -23,8 +23,11 @@ class Scanner:
self.languages = languages self.languages = languages
async def scan(self, path: str): async def scan(self, path: str):
logging.info("Starting the scan. It can take some times...")
videos = filter(lambda p: p.is_file(), Path(path).rglob("*")) videos = filter(lambda p: p.is_file(), Path(path).rglob("*"))
await asyncio.gather(*map(self.identify, videos)) # We batch videos by 30 because too mutch at once kinda DDOS everything.
for group in batch(videos, 30):
await asyncio.gather(*map(self.identify, group))
async def is_registered(self, path: Path) -> bool: async def is_registered(self, path: Path) -> bool:
# TODO: Once movies are separated from the api, a new endpoint should be created to check for paths. # TODO: Once movies are separated from the api, a new endpoint should be created to check for paths.

View File

@ -1,9 +1,25 @@
import asyncio import asyncio
import logging import logging
from functools import wraps from functools import wraps
from itertools import islice
from typing import Iterator, List, TypeVar
from providers.utils import ProviderError from providers.utils import ProviderError
T = TypeVar("T")
def batch(iterable: Iterator[T], n: int) -> Iterator[List[T]]:
"Batch data into lists of length n. The last batch may be shorter."
# batched('ABCDEFG', 3) --> ABC DEF G
it = iter(iterable)
while True:
batch = list(islice(it, n))
if not batch:
return
yield batch
def log_errors(f): def log_errors(f):
@wraps(f) @wraps(f)
async def internal(*args, **kwargs): async def internal(*args, **kwargs):