diff --git a/reddit2telegram/supplier.py b/reddit2telegram/supplier.py index edd7d33f..399b6e87 100644 --- a/reddit2telegram/supplier.py +++ b/reddit2telegram/supplier.py @@ -27,7 +27,8 @@ def send_to_channel_from_subreddit(how_to_post, channel_to_post, subreddit, subm client_id=config['reddit']['client_id'], client_secret=config['reddit']['client_secret'], username=config['reddit']['username'], - password=config['reddit']['password'] + password=config['reddit']['password'], + requestor_kwargs={'timeout': utils.REDDIT_TIMEOUT} ) if submissions_ranking == 'top': submissions = reddit.subreddit(subreddit).top(limit=submissions_limit) diff --git a/reddit2telegram/task_queue.py b/reddit2telegram/task_queue.py index 91bde0a6..c92b2983 100644 --- a/reddit2telegram/task_queue.py +++ b/reddit2telegram/task_queue.py @@ -17,6 +17,8 @@ TASK_SUPPLY = 'supply' +ABANDONED_TASK_MIN_AGE_SECONDS = 60 + logger = logging.getLogger(__name__) @@ -67,6 +69,29 @@ def select_all_available_tasks(collection: Collection): }, sort=[('created_at', pymongo.ASCENDING)])) +def recover_abandoned_tasks(collection: Collection): + cutoff = time.time() - ABANDONED_TASK_MIN_AGE_SECONDS + result = collection.update_many( + { + 'status': { + '$in': [ + TaskStatus.IN_PROGRESS.value, + TaskStatus.SCHEDULED.value, + ] + }, + 'updated_at': {'$lt': cutoff}, + }, + { + '$set': { + 'status': TaskStatus.NEW.value, + 'updated_at': time.time(), + }, + } + ) + if result.modified_count: + logger.warning('Recovered %d abandoned tasks', result.modified_count) + + def execute_task(collection: Collection, id: ObjectId, name: str, args: Mapping): update_task_status(collection, id, TaskStatus.IN_PROGRESS) try: @@ -90,6 +115,7 @@ def start_consumer( ): running = True collection = mongo_database[COLLECTION] + recover_abandoned_tasks(collection) logger.info('Starting consumer %s', executor) while running: try: diff --git a/reddit2telegram/utils/__init__.py b/reddit2telegram/utils/__init__.py index 87b3d56a..c9d42375 100644 --- a/reddit2telegram/utils/__init__.py +++ b/reddit2telegram/utils/__init__.py @@ -1,10 +1,9 @@ # encoding:utf-8 -import urllib from urllib.parse import urlparse from html import unescape import requests -from requests.exceptions import InvalidSchema, MissingSchema +from requests.exceptions import InvalidSchema, MissingSchema, RequestException import os import imghdr import random @@ -15,6 +14,7 @@ import enum import subprocess import asyncio +import threading from imgurpython import ImgurClient import yaml @@ -61,6 +61,16 @@ ERRORS_CNT_LIMIT = 2 +HTTP_TIMEOUT = (5, 20) +HEAD_TIMEOUT = (3, 8) +TELEGRAM_TIMEOUT = 30 +REDDIT_TIMEOUT = 20 + + +_MONGO_DATABASES = {} +_MONGO_LOCK = threading.Lock() + + @enum.unique class SupplyResult(enum.Enum): SUCCESSFULLY = 0 @@ -77,6 +87,20 @@ def _normalize_reddit_media_url(url): return url.replace('auto=webp', 'auto=jpg') +def _http_request(method, url, *, timeout=HTTP_TIMEOUT, **kwargs): + kwargs.setdefault('timeout', timeout) + return requests.request(method, url, **kwargs) + + +def _get_shared_database(config): + key = (config['db']['host'], config['db']['name']) + with _MONGO_LOCK: + if key not in _MONGO_DATABASES: + client = pymongo.MongoClient(host=key[0]) + _MONGO_DATABASES[key] = client[key[1]] + return _MONGO_DATABASES[key] + + def get_url(submission, mp4_instead_gif=True): ''' return TYPE, URL @@ -84,11 +108,13 @@ def get_url(submission, mp4_instead_gif=True): ''' def what_is_inside(url): - header = requests.head(url).headers - if 'Content-Type' in header: - return header['Content-Type'] - else: + try: + response = _http_request('head', url, timeout=HEAD_TIMEOUT, allow_redirects=True) + except RequestException: return '' + if 'Content-Type' in response.headers: + return response.headers['Content-Type'] + return '' # If reddit native gallery if hasattr(submission, 'gallery_data'): @@ -248,7 +274,7 @@ def imgur_album_to_story(images): elif 'gfycat.com' in urlparse(url).netloc: rname = re.findall(r'gfycat.com\/(?:detail\/)?(\w*)', url)[0] try: - r = requests.get(GFYCAT_GET + rname) + r = _http_request('get', GFYCAT_GET + rname) if r.status_code != 200: logging.info('Gfy fail prevented!') return TYPE_OTHER, url @@ -257,7 +283,7 @@ def imgur_album_to_story(images): return TYPE_GIF, urls['mp4Url'] else: return TYPE_GIF, urls['max5mbGif'] - except KeyError: + except (KeyError, RequestException): logging.info('Gfy fail prevented!') return TYPE_OTHER, url else: @@ -267,20 +293,23 @@ def imgur_album_to_story(images): def download_file(url, filename): # http://stackoverflow.com/questions/16694907/how-to-download-large-file-in-python-with-requests-py # NOTE the stream=True parameter - r = requests.get(url, stream=True) - if r.status_code >= 400: + try: + with _http_request('get', url, stream=True) as r: + if r.status_code >= 400: + return False + chunk_counter = 0 + chunk_size = 1024 + with open(filename, 'wb') as f: + for chunk in r.iter_content(chunk_size=chunk_size): + if chunk: # filter out keep-alive new chunks + f.write(chunk) + #f.flush() commented by recommendation from J.F.Sebastian + chunk_counter += 1 + # It is not possible to send greater than 50 MB via Telegram + if chunk_counter > TELEGRAM_VIDEO_LIMIT / chunk_size: + return False + except RequestException: return False - chunk_counter = 0 - chunk_size = 1024 - with open(filename, 'wb') as f: - for chunk in r.iter_content(chunk_size=chunk_size): - if chunk: # filter out keep-alive new chunks - f.write(chunk) - #f.flush() commented by recommendation from J.F.Sebastian - chunk_counter += 1 - # It is not possible to send greater than 50 MB via Telegram - if chunk_counter > TELEGRAM_VIDEO_LIMIT / chunk_size: - return False return True @@ -314,20 +343,23 @@ def clean_after_module(submodule_name=None): def md5_sum_from_url(url): try: - r = requests.get(url, stream=True) + r = _http_request('get', url, stream=True) except InvalidSchema: return None except MissingSchema: return None + except RequestException: + return None chunk_counter = 0 hash_store = hashlib.md5() - for chunk in r.iter_content(chunk_size=1024): - if chunk: # filter out keep-alive new chunks - hash_store.update(chunk) - chunk_counter += 1 - # It is not possible to send greater than 50 MB via Telegram - if chunk_counter > 50 * 1024: - return None + with r: + for chunk in r.iter_content(chunk_size=1024): + if chunk: # filter out keep-alive new chunks + hash_store.update(chunk) + chunk_counter += 1 + # It is not possible to send greater than 50 MB via Telegram + if chunk_counter > 50 * 1024: + return None return hash_store.hexdigest() @@ -343,12 +375,11 @@ def weighted_random_subreddit(weights): def get_url_size(url): # https://stackoverflow.com/questions/55226378/how-can-i-get-the-file-size-from-a-link-without-downloading-it-in-python - req = urllib.request.Request(url, method='HEAD') try: - f = urllib.request.urlopen(req) - except Exception: + response = _http_request('head', url, timeout=HEAD_TIMEOUT, allow_redirects=True) + except RequestException: return 0 - content_length = f.headers.get('Content-Length') + content_length = response.headers.get('Content-Length') if not content_length: return 0 return int(content_length) @@ -364,7 +395,13 @@ def __init__(self, t_channel=None, config=None): with open(os.path.join('configs', 'prod.yml')) as f: config = yaml.safe_load(f.read()) self.config = config - request = HTTPXRequest(connection_pool_size=8, pool_timeout=30) + request = HTTPXRequest( + connection_pool_size=8, + pool_timeout=TELEGRAM_TIMEOUT, + read_timeout=TELEGRAM_TIMEOUT, + write_timeout=TELEGRAM_TIMEOUT, + connect_timeout=TELEGRAM_TIMEOUT, + ) self.telegram_bot = Bot(self.config['telegram']['token'], request=request) self._loop = asyncio.new_event_loop() if t_channel is None: @@ -385,12 +422,13 @@ def _run_async(self, coro): return self._loop.run_until_complete(coro) def _make_mongo_connections(self): - self.stats = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['stats'] - self.urls = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['urls'] - self.contents = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['contents'] - self.errors = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['errors'] - self.tasks = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['tasks'] - self.settings = pymongo.MongoClient(host=self.config['db']['host'])[self.config['db']['name']]['settings'] + database = _get_shared_database(self.config) + self.stats = database['stats'] + self.urls = database['urls'] + self.contents = database['contents'] + self.errors = database['errors'] + self.tasks = database['tasks'] + self.settings = database['settings'] def _get_file_name(self, ext='file'): os.makedirs(TEMP_FOLDER, exist_ok=True) @@ -536,7 +574,7 @@ def send_gif(self, url, text, parse_mode=None): def _get_dash_audio_url(self, dash_url): try: - resp = requests.get(dash_url) + resp = _http_request('get', dash_url) if resp.status_code >= 400: return None base_urls = re.findall(r'([^<]+)', resp.text) diff --git a/tests/test_resilience.py b/tests/test_resilience.py new file mode 100644 index 00000000..22f48bd1 --- /dev/null +++ b/tests/test_resilience.py @@ -0,0 +1,92 @@ +import time +from types import SimpleNamespace +from unittest import mock + +import requests + +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parent.parent +APP_DIR = REPO_ROOT / 'reddit2telegram' +sys.path.insert(0, str(APP_DIR)) + +import task_queue +import utils + + +class FakeCollection: + def __init__(self, modified_count=0): + self.modified_count = modified_count + self.calls = [] + + def update_many(self, selector, update): + self.calls.append((selector, update)) + return SimpleNamespace(modified_count=self.modified_count) + + +def test_recover_abandoned_tasks_requeues_stale_non_terminal_work(): + collection = FakeCollection(modified_count=3) + + with mock.patch('task_queue.time.time', return_value=1_000): + task_queue.recover_abandoned_tasks(collection) + + assert len(collection.calls) == 1 + selector, update = collection.calls[0] + assert selector['status']['$in'] == [ + task_queue.TaskStatus.IN_PROGRESS.value, + task_queue.TaskStatus.SCHEDULED.value, + ] + assert selector['updated_at']['$lt'] == 940 + assert update['$set']['status'] == task_queue.TaskStatus.NEW.value + assert update['$set']['updated_at'] == 1_000 + + +def test_get_url_returns_other_when_head_request_fails(): + submission = SimpleNamespace( + url='https://example.com/image.jpg', + is_video=False, + media=None, + crosspost_parent_list=[], + is_self=False, + ) + + with mock.patch( + 'utils._http_request', + side_effect=requests.RequestException('timeout'), + ): + what, url = utils.get_url(submission) + + assert what == utils.TYPE_OTHER + assert url == 'https://example.com/image.jpg' + + +def test_get_url_size_returns_zero_when_head_request_fails(): + with mock.patch( + 'utils._http_request', + side_effect=requests.RequestException('timeout'), + ): + assert utils.get_url_size('https://example.com/file.mp4') == 0 + + +def test_shared_database_reuses_single_mongo_client(): + utils._MONGO_DATABASES.clear() + calls = [] + + class FakeClient: + def __init__(self, host): + calls.append(host) + + def __getitem__(self, name): + return {'db_name': name} + + config = {'db': {'host': 'localhost', 'name': 'reddit2telegram'}} + + with mock.patch('utils.pymongo.MongoClient', FakeClient): + db_one = utils._get_shared_database(config) + db_two = utils._get_shared_database(config) + + assert db_one == {'db_name': 'reddit2telegram'} + assert db_two == db_one + assert calls == ['localhost']