From b6b04f2739ab6111ce4073ddbdda5ceb22ecd98f Mon Sep 17 00:00:00 2001 From: Roland Date: Sat, 9 Jul 2022 18:51:45 -0600 Subject: [PATCH] Removing unused code. Updating comments. Removing timeout. Updating tests now that downloads work. Bumping version for async downloads. Downloads run asynchronously. Using aiohttp to make archive http request so ArchiveDownloader.do_download can be asynchronous. Removing session for create_download(s) usages. Reworkingk create_downloads to use optional_session. Canceling download workers after tests and when stopping the manager. All tests passing. Tests pass! Downloader tests pass. DownloadManager uses multiple workers to download domains in parallel. --- main.py | 18 +- modules/archive/__init__.py | 9 +- modules/archive/api.py | 3 +- modules/archive/lib.py | 34 +- modules/archive/test/test_lib.py | 135 ++++--- modules/videos/channel/lib.py | 1 - modules/videos/channel/test/test_api.py | 13 +- modules/videos/conftest.py | 4 +- modules/videos/downloader.py | 21 +- modules/videos/test/test_downloader.py | 73 ++-- requirements.txt | 12 +- wrolpi/common.py | 13 +- wrolpi/conftest.py | 16 +- wrolpi/downloader.py | 515 +++++++++++++++--------- wrolpi/root_api.py | 12 +- wrolpi/test/common.py | 3 +- wrolpi/test/test_db.py | 10 - wrolpi/test/test_downloader.py | 215 +++++----- wrolpi/test/test_root_api.py | 2 + wrolpi/test/test_rss.py | 41 +- wrolpi/test/test_timeout.py | 47 --- wrolpi/timeout.py | 114 ------ wrolpi/version.py | 2 +- 23 files changed, 650 insertions(+), 663 deletions(-) delete mode 100644 wrolpi/test/test_timeout.py delete mode 100644 wrolpi/timeout.py diff --git a/main.py b/main.py index 3be0cf2c..2d4c420c 100755 --- a/main.py +++ b/main.py @@ -7,12 +7,13 @@ import sys import pytz from sanic import Sanic +from sanic.signals import Event from wrolpi import root_api, BEFORE_STARTUP_FUNCTIONS, after_startup, limit_concurrent, admin from wrolpi.common import logger, get_config, import_modules, check_media_directory from wrolpi.dates import set_timezone from wrolpi.downloader import download_manager -from wrolpi.vars import PROJECT_DIR, DOCKERIZED +from wrolpi.vars import PROJECT_DIR, DOCKERIZED, PYTEST from wrolpi.version import get_version_string logger = logger.getChild('wrolpi-main') @@ -191,16 +192,21 @@ def periodic_downloads(app: Sanic, loop): logger.info('Starting download manager.') async def _periodic_download(): - download_manager.renew_recurring_downloads() - download_manager.delete_old_once_downloads() + download_manager.start_workers() await download_manager.do_downloads() - seconds = 60 - await asyncio.sleep(seconds) - app.add_task(_periodic_download()) + await asyncio.sleep(60) + app.add_task(_periodic_download()) # noqa app.add_task(_periodic_download()) +@root_api.api_app.signal(Event.SERVER_SHUTDOWN_BEFORE) +def handle_server_shutdown(*args, **kwargs): + """Stop downloads when server is shutting down.""" + if not PYTEST: + download_manager.stop() + + if __name__ == '__main__': loop_ = asyncio.get_event_loop() result_ = loop_.run_until_complete(main(loop_)) diff --git a/modules/archive/__init__.py b/modules/archive/__init__.py index 1a0afe57..1d059e8a 100644 --- a/modules/archive/__init__.py +++ b/modules/archive/__init__.py @@ -30,17 +30,16 @@ class ArchiveDownloader(Downloader, ABC): """ return True, None - def do_download(self, download: Download) -> DownloadResult: + async def do_download(self, download: Download) -> DownloadResult: if download.attempts > 3: raise UnrecoverableDownloadError(f'Max download attempts reached for {download.url}') - archive = lib.do_archive(download.url) + archive = await lib.do_archive(download.url) return DownloadResult(success=True, location=f'/archive/{archive.id}') @optional_session - def already_downloaded(self, url: str, session: Session = None) -> Optional[Archive]: - archive = session.query(Archive).filter_by(url=url).one_or_none() - return archive + def already_downloaded(self, url: str, session: Session = None) -> bool: + return bool(session.query(Archive).filter_by(url=url).count()) # Archive downloader is the last downloader which should be used. diff --git a/modules/archive/api.py b/modules/archive/api.py index b8a62e28..de25c386 100644 --- a/modules/archive/api.py +++ b/modules/archive/api.py @@ -1,8 +1,7 @@ import asyncio from http import HTTPStatus -from requests import Request -from sanic import response +from sanic import response, Request from sanic_ext import validate from sanic_ext.extensions.openapi import openapi diff --git a/modules/archive/lib.py b/modules/archive/lib.py index 7c9590c2..1a325a96 100644 --- a/modules/archive/lib.py +++ b/modules/archive/lib.py @@ -10,14 +10,15 @@ from datetime import datetime from itertools import groupby from typing import Iterator, Optional, Tuple, List -import requests +import aiohttp from bs4 import BeautifulSoup from selenium import webdriver from sqlalchemy.orm import Session from modules.archive.models import Domain, Archive from wrolpi.cmd import which -from wrolpi.common import get_media_directory, logger, chunks, extract_domain, chdir, escape_file_name, walk +from wrolpi.common import get_media_directory, logger, chunks, extract_domain, chdir, escape_file_name, walk, \ + aiohttp_post from wrolpi.dates import now, Seconds, local_timezone from wrolpi.db import get_db_session, get_db_curs, get_ranked_models, optional_session from wrolpi.errors import InvalidDomain, UnknownURL, InvalidArchive @@ -102,7 +103,7 @@ def get_new_archive_files(url: str, title: Optional[str]) -> ArchiveFiles: ARCHIVE_TIMEOUT = Seconds.minute * 10 # Wait at most 10 minutes for response. -def request_archive(url: str): +async def request_archive(url: str): """ Send a request to the archive service to archive the URL. """ @@ -110,18 +111,17 @@ def request_archive(url: str): data = {'url': url} try: - resp = requests.post(f'{ARCHIVE_SERVICE}/json', json=data, timeout=ARCHIVE_TIMEOUT) + contents, status = await aiohttp_post(f'{ARCHIVE_SERVICE}/json', json_=data, timeout=ARCHIVE_TIMEOUT) + readability = contents['readability'] + # Compressed base64 + singlefile = contents['singlefile'] + screenshot = contents['screenshot'] + + logger.debug(f'archive request status code {status}') except Exception as e: logger.error('Error when requesting archive', exc_info=e) raise - logger.debug(f'archive request status {resp.status_code=}') - - readability = resp.json()['readability'] - # Compressed base64 - singlefile = resp.json()['singlefile'] - screenshot = resp.json()['screenshot'] - if not (screenshot or singlefile or readability): raise Exception('singlefile response was empty!') @@ -147,8 +147,8 @@ def request_archive(url: str): SINGLE_FILE_BIN = which('single-file', - '/usr/bin/single-file', # rpi ubuntu - '/usr/local/bin/single-file', # debian + '/usr/bin/single-file', # rpi ubuntu + '/usr/local/bin/single-file', # debian warn=True) CHROMIUM = which('chromium-browser', 'chromium', '/usr/bin/chromium-browser', # rpi ubuntu @@ -194,8 +194,8 @@ def local_screenshot(url: str) -> bytes: READABILITY_BIN = which('readability-extractor', - '/usr/bin/readability-extractor', # rpi ubuntu - '/usr/local/bin/readability-extractor', # debian + '/usr/bin/readability-extractor', # rpi ubuntu + '/usr/local/bin/readability-extractor', # debian warn=True) @@ -227,7 +227,7 @@ def local_archive(url: str): return singlefile, readability, screenshot -def do_archive(url: str) -> Archive: +async def do_archive(url: str) -> Archive: """ Perform the real archive request to the archiving service. Store the resulting data into files. Create an Archive record in the DB. Create Domain/URL if missing. @@ -236,7 +236,7 @@ def do_archive(url: str) -> Archive: if DOCKERIZED or PYTEST: # Perform the archive in the Archive docker container. (Typically in the development environment). - singlefile, readability, screenshot = request_archive(url) + singlefile, readability, screenshot = await request_archive(url) else: # Perform the archive using locally installed executables. singlefile, readability, screenshot = local_archive(url) diff --git a/modules/archive/test/test_lib.py b/modules/archive/test/test_lib.py index 2bd134dd..560272fa 100644 --- a/modules/archive/test/test_lib.py +++ b/modules/archive/test/test_lib.py @@ -14,11 +14,11 @@ from wrolpi.dates import local_timezone from wrolpi.db import get_db_session from wrolpi.media_path import MediaPath from wrolpi.root_api import CustomJSONEncoder -from wrolpi.test.common import TestAPI, wrap_test_db, skip_circleci +from wrolpi.test.common import skip_circleci def make_fake_request_archive(readability=True, screenshot=True, title=True): - def fake_request_archive(_): + async def fake_request_archive(_): singlefile = '\ntest single-file\nジにてこちら\nsome title' r = dict( content=f'test readability content', @@ -31,63 +31,61 @@ def make_fake_request_archive(readability=True, screenshot=True, title=True): return fake_request_archive -class TestArchive(TestAPI): +@pytest.mark.asyncio +async def test_no_screenshot(test_session): + with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)): + archive = await do_archive('https://example.com') + assert isinstance(archive.singlefile_path, MediaPath) + assert isinstance(archive.readability_path, MediaPath) + # Screenshot was empty + assert archive.screenshot_path is None - def setUp(self) -> None: - super().setUp() - (pathlib.Path(self.tmp_dir.name) / 'archive').mkdir(exist_ok=True) - @wrap_test_db - def test_no_screenshot(self): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(screenshot=False)): - archive = do_archive('https://example.com') - self.assertIsInstance(archive.singlefile_path, MediaPath) - self.assertIsInstance(archive.readability_path, MediaPath) - # Screenshot was empty - self.assertIsNone(archive.screenshot_path) +@pytest.mark.asyncio +async def test_no_readability(test_session): + with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(readability=False)): + archive = await do_archive('https://example.com') + assert isinstance(archive.singlefile_path, MediaPath) + assert isinstance(archive.screenshot_path, MediaPath) + # Readability empty + assert archive.readability_path is None + assert archive.readability_txt_path is None - @wrap_test_db - def test_no_readability(self): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive(readability=False)): - archive = do_archive('https://example.com') - self.assertIsInstance(archive.singlefile_path, MediaPath) - self.assertIsInstance(archive.screenshot_path, MediaPath) - # Readability empty - self.assertIsNone(archive.readability_path) - self.assertIsNone(archive.readability_txt_path) - @wrap_test_db - def test_dict(self): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - d = do_archive('https://example.com').dict() - self.assertIsInstance(d, dict) - json.dumps(d, cls=CustomJSONEncoder) +@pytest.mark.asyncio +async def test_dict(test_session): + with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): + d = (await do_archive('https://example.com')).dict() + assert isinstance(d, dict) + json.dumps(d, cls=CustomJSONEncoder) - @wrap_test_db - def test_relationships(self): - with get_db_session(commit=True) as session: - url = 'https://wrolpi.org:443' - domain = get_or_create_domain(session, url) - archive = Archive( - singlefile_path=f'{self.tmp_dir.name}/wrolpi.org:443/foo', - title='bar', - url=url, - domain_id=domain.id, - ) - session.add(archive) - session.flush() - self.assertEqual(archive.domain, domain) +@pytest.mark.asyncio +async def test_relationships(test_session, test_directory): + with get_db_session(commit=True) as session: + url = 'https://wrolpi.org:443' + domain = get_or_create_domain(session, url) + archive = Archive( + singlefile_path=f'{test_directory.name}/wrolpi.org:443/foo', + title='bar', + url=url, + domain_id=domain.id, + ) + session.add(archive) + session.flush() - @wrap_test_db - def test_validate_paths(self): - with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - archive = do_archive('https://example.com') - try: - with get_db_session(commit=True): - archive.singlefile_path = 'asdf' - except ValueError as e: - self.assertIn('relative', str(e), f'Relative path error was not raised') + assert archive.domain == domain + + +@pytest.mark.asyncio +async def test_validate_paths(test_session): + with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): + archive = await do_archive('https://example.com') + try: + with get_db_session(commit=True): + archive.singlefile_path = 'asdf' + except ValueError as e: + assert 'relative' in str(e), f'Relative path error was not raised' def test_archive_refresh_deleted_archive(test_session, archive_directory, archive_factory): @@ -225,10 +223,11 @@ def test_get_domains(test_session, archive_factory): assert [i['domain'] for i in get_domains()] == [] -def test_new_archive(test_session, fake_now): +@pytest.mark.asycio +async def test_new_archive(test_session, fake_now): with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): fake_now(datetime(2000, 1, 1)) - archive1 = do_archive('https://example.com') + archive1 = await do_archive('https://example.com') # Everything is filled out. assert isinstance(archive1, Archive) assert archive1.archive_datetime is not None @@ -251,18 +250,19 @@ def test_new_archive(test_session, fake_now): assert json.load(fh) == {'title': 'ジにてこちら', 'url': 'https://example.com'} fake_now(datetime(2000, 1, 2)) - archive2 = do_archive('https://example.com') + archive2 = await do_archive('https://example.com') # Domain is reused. assert archive1.domain == archive2.domain -def test_get_title_from_html(test_session, fake_now): +@pytest.mark.asyncio +async def test_get_title_from_html(test_session, fake_now): fake_now(datetime(2000, 1, 1)) with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - archive = do_archive('example.com') + archive = await do_archive('example.com') assert archive.title == 'ジにてこちら' - def fake_request_archive(_): + async def fake_request_archive(_): singlefile = '\ntest single-file\nジにてこちら\nsome title' r = dict( content=f'test readability content', @@ -273,10 +273,10 @@ def test_get_title_from_html(test_session, fake_now): fake_now(datetime(2000, 1, 2)) with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive = do_archive('example.com') + archive = await do_archive('example.com') assert archive.title == 'some title' - def fake_request_archive(_): + async def fake_request_archive(_): singlefile = '' r = dict( content=f'missing a title', @@ -287,7 +287,7 @@ def test_get_title_from_html(test_session, fake_now): fake_now(datetime(2000, 1, 3)) with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive = do_archive('example.com') + archive = await do_archive('example.com') assert archive.title is None @@ -313,14 +313,15 @@ def test_get_new_archive_files(fake_now): @skip_circleci -def test_title_in_filename(test_session, fake_now, test_directory): +@pytest.mark.asyncio +async def test_title_in_filename(test_session, fake_now, test_directory): """ The Archive files have the title in the path. """ fake_now(datetime(2000, 1, 1)) with mock.patch('modules.archive.lib.request_archive', make_fake_request_archive()): - archive1 = do_archive('example.com') + archive1 = await do_archive('example.com') assert archive1.title == 'ジにてこちら' @@ -335,7 +336,7 @@ def test_title_in_filename(test_session, fake_now, test_directory): assert str(archive1.screenshot_path.path.relative_to(test_directory)) == \ 'archive/2000-01-01-00-00-00_ジにてこちら.png' - def fake_request_archive(_): + async def fake_request_archive(_): singlefile = '\ntest single-file\nジにてこちら\n' # no title in HTML r = dict( # No Title from Readability. @@ -346,7 +347,7 @@ def test_title_in_filename(test_session, fake_now, test_directory): return singlefile, r, s with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive2 = do_archive('example.com') + archive2 = await do_archive('example.com') assert str(archive2.singlefile_path.path.relative_to(test_directory)) == \ 'archive/2000-01-01-00-00-00_NA.html' @@ -359,7 +360,7 @@ def test_title_in_filename(test_session, fake_now, test_directory): assert str(archive2.screenshot_path.path.relative_to(test_directory)) == \ 'archive/2000-01-01-00-00-00_NA.png' - def fake_request_archive(_): + async def fake_request_archive(_): singlefile = '\ntest single-file\nジにてこちら\ndangerous ;\\//_title</html></html>' r = dict( # No Title from Readability. @@ -370,7 +371,7 @@ def test_title_in_filename(test_session, fake_now, test_directory): return singlefile, r, s with mock.patch('modules.archive.lib.request_archive', fake_request_archive): - archive3 = do_archive('example.com') + archive3 = await do_archive('example.com') assert str(archive3.singlefile_path.path.relative_to(test_directory)) == \ 'archive/2000-01-01-00-00-00_dangerous ;_title.html' diff --git a/modules/videos/channel/lib.py b/modules/videos/channel/lib.py index 98c0abe2..2dd9efbe 100644 --- a/modules/videos/channel/lib.py +++ b/modules/videos/channel/lib.py @@ -199,4 +199,3 @@ def download_channel(id_: int): raise InvalidDownload(f'Channel {channel.name} does not have a download!') download.renew(reset_attempts=True) session.commit() - download_manager.start_downloads() diff --git a/modules/videos/channel/test/test_api.py b/modules/videos/channel/test/test_api.py index 04756506..e298f9be 100644 --- a/modules/videos/channel/test/test_api.py +++ b/modules/videos/channel/test/test_api.py @@ -6,6 +6,7 @@ from datetime import timedelta from http import HTTPStatus import mock +import pytest from modules.videos.channel.lib import delete_channel from modules.videos.common import get_no_channel_directory @@ -108,7 +109,8 @@ def test_channel_no_download_frequency(test_client, test_session, test_directory assert response.status_code == HTTPStatus.NO_CONTENT, response.json -def test_channel_frequency_update(download_channel, test_client, test_session): +@pytest.mark.asyncio +def test_channel_frequency_update(download_channel, test_client, test_session, test_download_manager): """ A Channel's Download record is updated when the Channel's frequency is updated. """ @@ -336,8 +338,12 @@ def test_channel_empty_url_doesnt_conflict(test_client, test_session, test_direc assert location != response.headers['Location'] -def test_download_channel_no_refresh(test_session, download_channel, video_download_manager): +@pytest.mark.asyncio +async def test_download_channel_no_refresh(test_session, download_channel, video_download_manager): """A Channel cannot be downloaded until it has been refreshed.""" + d = download_channel.get_download() + d.next_download = video_download_manager.calculate_next_download(d) + test_session.commit() def check_refreshed(expected: bool): channel = test_session.query(Channel).one() @@ -349,7 +355,8 @@ def test_download_channel_no_refresh(test_session, download_channel, video_downl with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info: mock_extract_info.return_value = {'entries': [], 'url': 'foo', 'uploader': 'some uploader', 'channel_id': 'the id', 'id': 'the id'} - video_download_manager.do_downloads_sync() + await video_download_manager.do_downloads() + await video_download_manager.wait_for_all_downloads() check_refreshed(True) diff --git a/modules/videos/conftest.py b/modules/videos/conftest.py index ddc4e55f..df2c3e9c 100644 --- a/modules/videos/conftest.py +++ b/modules/videos/conftest.py @@ -140,9 +140,7 @@ def video_factory(test_session, test_directory): @pytest.fixture def video_download_manager(test_download_manager) -> DownloadManager: - """ - Get a DownloadManager ready to download Videos and Channels. - """ + """Get a DownloadManager ready to download Videos and Channels.""" channel_downloader = ChannelDownloader() video_downloader = VideoDownloader(priority=40) diff --git a/modules/videos/downloader.py b/modules/videos/downloader.py index cc439e55..98c95a78 100755 --- a/modules/videos/downloader.py +++ b/modules/videos/downloader.py @@ -81,7 +81,7 @@ class ChannelDownloader(Downloader, ABC): # A playlist will have an id different from its channel. return info['id'] != info['channel_id'] - def do_download(self, download: Download) -> DownloadResult: + async def do_download(self, download: Download) -> DownloadResult: """Update a Channel's catalog, then schedule downloads of every missing video.""" info = extract_info(download.url, process=False) # Resolve the entries generator. @@ -142,8 +142,9 @@ YT_DLP_BIN = which( class VideoDownloader(Downloader, ABC): - """ - Download a single video. Store the video in its channel's directory, otherwise store it in `videos/NO CHANNEL`. + """Downloads a single video. + + Store the video in its channel's directory, otherwise store it in `videos/NO CHANNEL`. """ name = 'video' pretty_name = 'Videos' @@ -152,13 +153,9 @@ class VideoDownloader(Downloader, ABC): return f'<VideoDownloader>' @optional_session - def already_downloaded(self, url: str, session: Session = None) -> Optional[Video]: + def already_downloaded(self, url: str, session: Session = None) -> bool: # We only consider a video record with a video file as "downloaded". - video = session.query(Video).filter( - Video.url == url, - Video.video_path != None, - ).one_or_none() - return video + return bool(session.query(Video).filter(Video.url == url, Video.video_path != None, ).count()) # noqa @classmethod def valid_url(cls, url) -> Tuple[bool, Optional[dict]]: @@ -177,7 +174,7 @@ class VideoDownloader(Downloader, ABC): logger.debug(f'{cls.__name__} not suitable for {url}') return False, None - def do_download(self, download: Download) -> DownloadResult: + async def do_download(self, download: Download) -> DownloadResult: if download.attempts >= 10: raise UnrecoverableDownloadError('Max download attempts reached') @@ -226,7 +223,7 @@ class VideoDownloader(Downloader, ABC): '--compat-options', 'no-live-chat', url, ) - return_code, logs = self.process_runner(url, cmd, out_dir) + return_code, logs = await self.process_runner(url, cmd, out_dir) stdout = logs['stdout'].decode() if hasattr(logs['stdout'], 'decode') else logs['stdout'] stderr = logs['stderr'].decode() if hasattr(logs['stderr'], 'decode') else logs['stderr'] @@ -293,6 +290,7 @@ class VideoDownloader(Downloader, ABC): @staticmethod def prepare_filename(url: str, out_dir: pathlib.Path) -> Tuple[pathlib.Path, dict]: + """Get the full path of a video file from its URL.""" if not out_dir.is_dir(): raise ValueError(f'Output directory does not exist! {out_dir=}') @@ -300,7 +298,6 @@ class VideoDownloader(Downloader, ABC): options = get_downloader_config().dict() options['outtmpl'] = f'{out_dir}/{options["file_name_format"]}' options['merge_output_format'] = PREFERRED_VIDEO_EXTENSION - logger.debug(f'Download video with options: {options}') logger.debug(f'Downloading {url} to {out_dir}') diff --git a/modules/videos/test/test_downloader.py b/modules/videos/test/test_downloader.py index 7950f821..6ed78d2e 100644 --- a/modules/videos/test/test_downloader.py +++ b/modules/videos/test/test_downloader.py @@ -150,8 +150,9 @@ class TestVideosDownloaders(TestAPI): self.assertEqual(self.mgr.get_downloader('https://www.youtube.com/watch?v=31jPEBiAC3c')[0], self.vd) -def test_video_download_no_channel(test_session, video_download_manager, video_factory, test_directory, - mock_video_extract_info, mock_video_process_runner): +@pytest.mark.asyncio +async def test_video_download_no_channel(test_session, video_download_manager, video_factory, test_directory, + mock_video_extract_info, mock_video_process_runner): """A video can be downloaded even if it does not have a Channel.""" channel_dir = test_directory / 'NO CHANNEL' channel_dir.mkdir(parents=True) @@ -168,7 +169,8 @@ def test_video_download_no_channel(test_session, video_download_manager, video_f mock_video_extract_info.return_value = info_json with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename: mock_prepare_filename.return_value = (video_path, {'id': 'foo'}) - video_download_manager.create_download(url, test_session) + video_download_manager.create_download(url) + await video_download_manager.wait_for_all_downloads() mock_video_process_runner.assert_called_once() @@ -176,9 +178,10 @@ def test_video_download_no_channel(test_session, video_download_manager, video_f assert str(video.video_path.path) == f'{channel_dir}/video.mp4' -def test_download_channel(test_session, simple_channel, video_download_manager, video_file, - mock_video_extract_info, mock_video_prepare_filename, - mock_video_process_runner): +@pytest.mark.asyncio +async def test_download_channel(test_session, simple_channel, video_download_manager, video_file, + mock_video_extract_info, mock_video_prepare_filename, + mock_video_process_runner): """Downloading (updating the catalog of) a Channel creates download records for all of it's missing videos.""" url = 'https://www.youtube.com/c/LearningSelfReliance/videos' @@ -186,7 +189,8 @@ def test_download_channel(test_session, simple_channel, video_download_manager, mock_video_extract_info.return_value = example_channel_json with mock.patch('modules.videos.downloader.get_channel') as mock_get_channel: mock_get_channel.return_value = simple_channel - video_download_manager.create_download(url, test_session) + video_download_manager.create_download(url) + await video_download_manager.wait_for_all_downloads() # Two videos are in the example channel. downloads = video_download_manager.get_once_downloads(test_session) @@ -238,18 +242,22 @@ def test_get_or_create_channel(test_session): assert get_or_create_channel(source_id='quux') == channel -def test_download_info_save(test_session, video_download_manager, mock_video_extract_info, mock_video_prepare_filename, - mock_video_process_runner): +@pytest.mark.asyncio +async def test_download_info_save(test_session, video_download_manager, mock_video_extract_info, + mock_video_prepare_filename, + mock_video_process_runner): """ An info dict from Downloader.valid_url is saved to the Downloader and can be used by the Downloader later. """ mock_video_extract_info.return_value = {'some': 'info'} mock_video_prepare_filename.return_value = 'some path' - video_download_manager.create_download('https://www.youtube.com/watch?v=HQ_62YwcA80', test_session) + video_download_manager.create_download('https://www.youtube.com/watch?v=HQ_62YwcA80') + await video_download_manager.wait_for_all_downloads() assert mock_video_extract_info.call_count == 2 # Only called twice. The "info" gathering call was skipped. -def test_selected_downloader(test_session, video_download_manager, successful_download): +@pytest.mark.asyncio +async def test_selected_downloader(test_session, video_download_manager, successful_download): """ A user can specify which Downloader to use, no other should ever be used (even if the Download fails). """ @@ -267,7 +275,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do # DownloadManager will attempt to choose the Downloader. with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info: mock_extract_info.side_effect = UnsupportedError('not this one') - video_download_manager.create_download('https://example.com', test_session) + video_download_manager.create_download('https://example.com') + await video_download_manager.wait_for_all_downloads() mock_http_valid_url.assert_called_once() mock_extract_info.assert_called_once() mock_do_download.assert_called_once() @@ -279,7 +288,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do # DownloadManager will only use the HTTPDownloader. with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info: - video_download_manager.create_download('https://example.com', test_session, downloader='http') + video_download_manager.create_download('https://example.com', downloader='http') + await video_download_manager.wait_for_all_downloads() # DownloadManager does not try to find the Download, it trusts the "downloader" param. mock_extract_info.assert_not_called() mock_http_valid_url.assert_not_called() @@ -294,7 +304,8 @@ def test_selected_downloader(test_session, video_download_manager, successful_do # DownloadManager will only use the HTTPDownloader, even if the download fails. with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info: mock_do_download.side_effect = Exception('oh no!') - video_download_manager.create_download('https://example.com', test_session, downloader='http') + video_download_manager.create_download('https://example.com', downloader='http') + await video_download_manager.wait_for_all_downloads() # DownloadManager does not try to find the Download, it trusts the "downloader" param. mock_extract_info.assert_not_called() mock_http_valid_url.assert_not_called() @@ -312,12 +323,13 @@ def test_channel_downloader_hidden(video_download_manager): assert [i.__json__() for i in downloaders] == [(dict(name='video', pretty_name='Videos')), ] +@pytest.mark.asyncio def test_bad_downloader(test_session, video_download_manager): """ Attempting to use an unknown downloader should raise an error. """ with pytest.raises(InvalidDownload): - video_download_manager.create_download('https://example.com', test_session, 'bad downloader') + video_download_manager.create_download('https://example.com', downloader='bad downloader') def test_channel_download_no_download(test_session, video_download_manager, simple_channel): @@ -332,22 +344,27 @@ def test_channel_download_no_download(test_session, video_download_manager, simp download_channel(simple_channel.id) -def test_invalid_download_url(test_session, test_download_manager, mock_video_extract_info): +@pytest.mark.asyncio +async def test_invalid_download_url(test_session, test_download_manager, mock_video_extract_info): """An invalid url should not be attempted again.""" _, session = get_db_context() with mock.patch('modules.videos.downloader.VideoDownloader.valid_url') as mock_valid_url, \ mock.patch('wrolpi.downloader.DownloadManager.get_downloader') as mock_get_downloader: - mock_get_downloader.return_value = (VideoDownloader(), None) + video_downloader = VideoDownloader() + test_download_manager.register_downloader(video_downloader) + mock_get_downloader.return_value = (video_downloader, None) mock_valid_url.return_value = (True, {}) # url is valid, but is not a video URL. mock_video_extract_info.side_effect = UnsupportedError('oops') - test_download_manager.create_download('invalid url', session) + test_download_manager.create_download('invalid url', downloader=VideoDownloader.name) + await test_download_manager.wait_for_all_downloads() download = session.query(Download).one() assert download.status == 'failed' -def test_video_download_1(test_session, test_directory, simple_channel, video_download_manager, - mock_video_process_runner): +@pytest.mark.asyncio +async def test_video_download_1(test_session, test_directory, simple_channel, video_download_manager, + mock_video_process_runner): simple_channel.source_id = example_video_json['channel_id'] simple_channel.directory = test_directory / 'videos/channel name' simple_channel.directory.mkdir(parents=True) @@ -361,7 +378,8 @@ def test_video_download_1(test_session, test_directory, simple_channel, video_do mock_extract_info.return_value = example_video_json mock_prepare_filename.return_value = (video_file, {'id': 'foo'}) - video_download_manager.create_download(url, test_session) + video_download_manager.create_download(url) + await video_download_manager.wait_for_all_downloads() mock_video_process_runner.assert_called_once() video_url, _, out_dir = mock_video_process_runner.call_args[0] @@ -374,8 +392,9 @@ def test_video_download_1(test_session, test_directory, simple_channel, video_do assert video.video_path.path.is_absolute() -def test_download_result(test_session, test_directory, video_download_manager, mock_video_process_runner, - mock_video_extract_info): +@pytest.mark.asyncio +async def test_download_result(test_session, test_directory, video_download_manager, mock_video_process_runner, + mock_video_extract_info): """VideoDownloader returns a DownloadResult when complete.""" channel_directory = test_directory / 'videos/channel name' channel_directory.mkdir(parents=True) @@ -385,7 +404,8 @@ def test_download_result(test_session, test_directory, video_download_manager, m mock_video_extract_info.return_value = example_video_json with mock.patch('modules.videos.downloader.VideoDownloader.prepare_filename') as mock_prepare_filename: mock_prepare_filename.return_value = [video_file, {'id': 'foo'}] - video_download_manager.create_download('https://example.com', test_session) + video_download_manager.create_download('https://example.com') + await video_download_manager.wait_for_all_downloads() download: Download = test_session.query(Download).one() assert download.url == 'https://example.com' @@ -451,7 +471,8 @@ example_playlist_json = { } -def test_download_playlist(test_session, test_directory, mock_video_extract_info, video_download_manager): +@pytest.mark.asyncio +async def test_download_playlist(test_session, test_directory, mock_video_extract_info, video_download_manager): """ All videos in a playlist can be downloaded for it's Channel. """ @@ -467,7 +488,7 @@ def test_download_playlist(test_session, test_directory, mock_video_extract_info with mock.patch('modules.videos.downloader.VideoDownloader.do_download') as mock_video_do_download: mock_video_do_download.return_value = DownloadResult(success=True) # Don't download the videos. - result = channel_downloader.do_download(download) + result = await channel_downloader.do_download(download) assert result.success is True, 'Download was not successful' assert set(result.downloads) == { 'https://youtube.com/watch?v=video_2_url', diff --git a/requirements.txt b/requirements.txt index 109d6658..06ab2f28 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,23 +1,23 @@ +Pillow==9.0.1 +Pint==0.18 +PyYAML==5.3.1 # this is the only version available on RPi +SQLAlchemy==1.3.22 +aiohttp==3.8.1 alembic==1.7.5 beautifulsoup4==4.10.0 cachetools==5.0.0 feedparser==6.0.9 mock==4.0.3 -Pillow==9.0.1 -Pint==0.18 psycopg2-binary==2.9.3 pytest-asyncio==0.17.2 pytest==6.2.5 python-dotenv==0.19.2 pytz==2021.3 -PyYAML==5.3.1 # this is the only version available on RPi -requests==2.27.1 sanic-ext==22.1.2 sanic-testing==0.8.2 sanic==21.12.1 selenium==4.1.0 -SQLAlchemy==1.3.22 srt==3.5.1 websockets==10.1 webvtt-py==0.4.6 -yt-dlp==2022.6.29 +yt-dlp==2022.7.18 diff --git a/wrolpi/common.py b/wrolpi/common.py index b0e40cd0..317bd3f8 100644 --- a/wrolpi/common.py +++ b/wrolpi/common.py @@ -20,6 +20,7 @@ from pathlib import Path from typing import Union, Callable, Tuple, Dict, List, Iterable, Optional, Generator, Any from urllib.parse import urlunsplit, urlparse +import aiohttp import yaml from sanic import Blueprint from sanic.request import Request @@ -370,9 +371,7 @@ def set_wrol_mode(enable: bool): def wrol_mode_check(func): - """ - Wraps a function so that it cannot be called when WROL Mode is enabled. - """ + """Wraps a function so that it cannot be called when WROL Mode is enabled.""" @wraps(func) def check(*a, **kw): @@ -786,3 +785,11 @@ def any_extensions(filename: str, extensions: Iterable = ()): Matches lower or upper case of the extension. """ return any(filename.lower().endswith(ext) for ext in extensions) + + +async def aiohttp_post(url: str, json_, timeout: int = None) -> Tuple[Dict, int]: + """Perform an async aiohttp post request. Return the json contents.""" + timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(url, json=json_) as response: + return await response.json(), response.status diff --git a/wrolpi/conftest.py b/wrolpi/conftest.py index 47cea0e1..70b5ecdc 100644 --- a/wrolpi/conftest.py +++ b/wrolpi/conftest.py @@ -1,6 +1,7 @@ """ Fixtures for Pytest tests. """ +import asyncio import pathlib import tempfile from typing import Tuple, Set @@ -121,9 +122,20 @@ def test_download_manager_config(test_directory): @pytest.fixture -def test_download_manager(test_download_manager_config): +async def test_download_manager( + test_session, # session is required because downloads can start without the test DB in place. + test_download_manager_config, +): manager = DownloadManager() - return manager + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.get_event_loop() + manager.start_workers(loop=loop) + + yield manager + + manager.cancel_workers() @pytest.fixture diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index 61950a76..4cd1581b 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -1,16 +1,17 @@ import asyncio import multiprocessing import pathlib -import subprocess import traceback from abc import ABC +from asyncio import Queue, Task, QueueEmpty from dataclasses import dataclass, field from datetime import timedelta, datetime from enum import Enum from functools import partial from operator import attrgetter -from typing import List, Dict +from typing import List, Dict, Generator from typing import Tuple, Optional +from urllib.parse import urlparse import feedparser from feedparser import FeedParserDict @@ -20,15 +21,12 @@ from sqlalchemy.orm import Session from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, ConfigFile, WROLPI_CONFIG from wrolpi.dates import TZDateTime, now, Seconds, local_timezone, recursive_replace_tz -from wrolpi.db import get_db_session, get_db_curs, get_db_context, optional_session +from wrolpi.db import get_db_session, get_db_curs, optional_session from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError from wrolpi.vars import PYTEST logger = logger.getChild(__name__) -DOWNLOAD_IN_PROGRESS = multiprocessing.Semaphore() -DEFAULT_RETRY_FREQUENCY = timedelta(hours=1) - class DownloadFrequency(int, Enum): hourly = 3600 @@ -52,7 +50,7 @@ class DownloadResult: class Download(ModelHelper, Base): """Model that is used to schedule downloads.""" - __tablename__ = 'download' + __tablename__ = 'download' # noqa id = Column(Integer, primary_key=True) url = Column(String, nullable=False) @@ -66,19 +64,18 @@ class Download(ModelHelper, Base): next_download = Column(TZDateTime) status = Column(String, default='new') sub_downloader = Column(Text) + _manager = None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.manager: DownloadManager = None - def __repr__(self): if self.next_download or self.frequency: return f'<Download id={self.id} status={self.status} url={repr(self.url)} ' \ - f'next_download={self.next_download} frequency={self.frequency} attempts={self.attempts} ' \ - f'error={self.error}>' + f'next_download={repr(self.next_download)} frequency={self.frequency} attempts={self.attempts} ' \ + f'error={bool(self.error)}>' return f'<Download id={self.id} status={self.status} url={repr(self.url)} attempts={self.attempts} ' \ - f'error={self.error}>' + f'error={bool(self.error)}>' def __json__(self): d = dict( @@ -126,6 +123,20 @@ class Download(ModelHelper, Base): return self.manager.get_downloader(self.url) + @property + def domain(self): + return urlparse(self.url).netloc + + @property + def manager(self): + if self._manager: + return self._manager + raise ValueError('No manager has been set!') + + @manager.setter + def manager(self, value): + self._manager = value + class Downloader: name: str = None @@ -156,11 +167,14 @@ class Downloader: def __json__(self): return dict(name=self.name, pretty_name=self.pretty_name) + def __repr__(self): + return f'<Downloader name={self.name}>' + @classmethod def valid_url(cls, url: str) -> Tuple[bool, Optional[dict]]: raise NotImplementedError() - def do_download(self, download: Download) -> DownloadResult: + async def do_download(self, download: Download) -> DownloadResult: raise NotImplementedError() @optional_session @@ -187,8 +201,8 @@ class Downloader: if self._kill.is_set(): self._kill.clear() - def process_runner(self, url: str, cmd: Tuple[str, ...], cwd: pathlib.Path, timeout: int = None, - **kwargs) -> Tuple[int, dict]: + async def process_runner(self, url: str, cmd: Tuple[str, ...], cwd: pathlib.Path, timeout: int = None, + **kwargs) -> Tuple[int, dict]: """ Run a subprocess using the provided arguments. This process can be killed by the Download Manager. @@ -196,7 +210,11 @@ class Downloader: """ logger.debug(f'{self} launching download process with args: {" ".join(cmd)}') start = now() - proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd, **kwargs) + proc = await asyncio.create_subprocess_exec(*cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=cwd, + **kwargs) pid = proc.pid # Timeout can be any positive integer. Global download timeout takes precedence, unless it is 0. @@ -204,37 +222,53 @@ class Downloader: timeout = WROLPI_CONFIG.download_timeout or timeout or self.timeout logger.debug(f'{self} launched download process {pid=} {timeout=} for {url}') + stdout, stderr = None, None try: while True: - try: - proc.wait(timeout=1.0) - # Download finished. + task = asyncio.Task(proc.communicate()) + done, pending = await asyncio.wait([task, ], timeout=1) + # Cancel communicate. + if pending: + pending.pop().cancel() + if done: + # Process finished, get the result of proc.communicate(). + stdout, stderr = done.pop().result() + break + if proc.returncode: + # Process died. + stdout, stderr = await proc.communicate() break - except subprocess.TimeoutExpired: - pass - if timeout and (elapsed := now() - start).total_seconds() > timeout: + elapsed = (now() - start).total_seconds() + if timeout and elapsed > timeout: logger.warning(f'Download has exceeded its timeout {elapsed=}') self.kill() if self._kill.is_set(): - logger.warning(f'Killing download {pid=}') + logger.warning(f'Killing download {pid=}, {elapsed} seconds elapsed (timeout was not exceeded).') proc.kill() - proc.poll() break + except Exception as e: + logger.error(f'{self}.process_runner had a download error', exc_info=e) + raise finally: self.clear() # Output all logs from the process. # TODO is there a way to stream this output while the process is running? logger.debug(f'Download exited with {proc.returncode}') - stdout, stderr = proc.communicate() logs = {'stdout': stdout, 'stderr': stderr} return proc.returncode, logs class DownloadManager: + """ + Runs a collection of workers which will download any URLs in the `download` table. + + Only one download from each domain will be downloaded at a time. If we only have one domain (example.com) in the + URLs to be downloaded, then only one worker will be busy. + """ priority_sorter = partial(sorted, key=attrgetter('priority')) def __init__(self): @@ -242,6 +276,12 @@ class DownloadManager: self._instances = dict() self.disabled = multiprocessing.Event() self.stopped: bool = False + self.download_queue: Queue = Queue() + self.workers: List[Task] = [] + self.worker_count: int = 4 + self.data = multiprocessing.Manager().dict() + # We haven't started downloads yet, so no domains are downloading. + self.data['processing_domains'] = [] def register_downloader(self, instance: Downloader): if not isinstance(instance, Downloader): @@ -255,6 +295,132 @@ class DownloadManager: self.instances = tuple(self.priority_sorter(i)) self._instances[instance.name] = instance + async def download_worker(self, num: int, queue: Queue): + """Fetch a download from the queue, perform the download then store the results. + + Calls DownloadManger.start_downloads() after a download completes. + """ + logger.debug(f'DownloadManager.download_worker {num} starting up') + from wrolpi.db import get_db_session + + while True: + if self.stopped: + return + + try: + download: Download = queue.get_nowait() + download_id, url = download.id, download.url + logger.debug(f'Download worker {num} got download {download}') + + downloader: Downloader = download.get_downloader() + if not downloader: + logger.warning(f'Could not find downloader for {download.downloader=}') + + with get_db_session(commit=True) as session: + # Mark the download as started in new session so the change is committed. + download = session.query(Download).filter_by(id=download_id).one() + download.started() + + self.data['processing_domains'].append(download.domain) + + try_again = True + try: + if asyncio.iscoroutinefunction(downloader.do_download): + result = await downloader.do_download(download) + else: + result = downloader.do_download(download) + except UnrecoverableDownloadError as e: + # Download failed and should not be retried. + logger.warning(f'UnrecoverableDownloadError for {url}', exc_info=e) + result = DownloadResult(success=False, error=str(traceback.format_exc())) + try_again = False + except Exception as e: + logger.warning(f'Failed to download {url}. Will be tried again later.', exc_info=e) + result = DownloadResult(success=False, error=str(traceback.format_exc())) + + error_len = len(result.error) if result.error else 0 + logger.debug(f'Got success={result.success} from {downloader} with {error_len=}') + + with get_db_session(commit=True) as session: + # Modify the download in a new session because downloads may take a long time. + download = session.query(Download).filter_by(id=download_id).one() + # Use a new location if provided, keep the old location if no new location is provided, otherwise + # clear out an outdated location. + download.location = result.location or download.location or None + # Clear any old errors if the download succeeded. + download.error = result.error if result.error else None + download.next_download = self.calculate_next_download(download, session) + + if result.downloads: + logger.info(f'Adding {len(result.downloads)} downloads from result of {download.url}') + self.create_downloads(result.downloads, session, + downloader=download.sub_downloader) + + if try_again is False and not download.frequency: + # Only once-downloads can fail. + download.fail() + elif result.success: + download.complete() + else: + download.defer() + + queue.task_done() + # Remove this domain from the running list. + self._remove_domain(download.domain) + # Request any new downloads be added to the queue. + asyncio.create_task(self.queue_downloads()) + except asyncio.CancelledError: + logger.warning(f'download_worker canceled!') + queue.task_done() + return + except QueueEmpty: + # No work yet. + await asyncio.sleep(0.1) + continue + except Exception as e: + logger.warning(f'Download worker had unexpected error', exc_info=e) + + def _add_domain(self, domain: str): + """Add a domain to the processing list. + + Raises ValueError if the domain is already processing.""" + processing_domains = self.data['processing_domains'] + if domain in processing_domains: + raise ValueError(f'Domain already being downloaded! {domain}') + self.data['processing_domains'] = [*processing_domains, domain] + + def _remove_domain(self, domain: str): + """Remove a domain from the processing list.""" + self.data['processing_domains'] = [i for i in self.data['processing_domains'] if i != domain] + + @wrol_mode_check + def start_workers(self, loop=None): + """Start all download worker tasks. Does nothing if they are already running.""" + if not self.workers: + if not loop: + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # Loop isn't running, start one. Probably testing. + loop = asyncio.get_event_loop() + + logger.debug(f'Starting {self.worker_count} download workers') + for i in range(self.worker_count): + coro = self.download_worker(i, self.download_queue) + task = loop.create_task(coro) + self.workers.append(task) + + def workers_running(self): + for task in self.workers: + if not task.done(): + return True + return False + + def cancel_workers(self): + if self.workers_running(): + for task in self.workers: + task.cancel() + def get_downloader(self, url: str) -> Tuple[Downloader, Dict]: for i in self.instances: valid, info = i.valid_url(url) @@ -262,29 +428,10 @@ class DownloadManager: return i, info raise InvalidDownload(f'Invalid URL {url=}') - def get_downloader_by_name(self, name: str) -> Downloader: - """ - Tries to find a registered Downloader by its name. Returns None if it cannot be found. - """ + def get_downloader_by_name(self, name: str) -> Optional[Downloader]: + """Attempt to find a registered Downloader by its name. Returns None if it cannot be found.""" return self._instances.get(name) - def get_new_downloads(self, session: Session) -> List[Download]: - """ - Get all "new" downloads. This method fetches the first download each iteration, so it will fetch downloads - that were created after calling it. - """ - last = None - while True: - download = session.query(Download).filter_by(status='new').order_by(Download.id).first() - if not download: - return - if download == last: - # Got the last download again. Is something wrong? - return - last = download - download.manager = self - yield download - def get_or_create_download(self, url: str, session: Session) -> Download: """Get a Download by its URL, if it cannot be found create one.""" if not url: @@ -301,159 +448,136 @@ class DownloadManager: download.manager = self return download - def create_download(self, url: str, session, downloader: str = None, skip_download: bool = False, - reset_attempts: bool = False, sub_downloader: str = None) -> Download: + @optional_session + def create_download(self, url: str, session: Session = None, downloader: str = None, reset_attempts: bool = False, + sub_downloader: str = None) -> Download: """Schedule a URL for download. If the URL failed previously, it may be retried.""" - downloads = self.create_downloads([url], session, downloader, skip_download=skip_download, + downloads = self.create_downloads([url], session=session, downloader=downloader, reset_attempts=reset_attempts, sub_downloader=sub_downloader) return downloads[0] - def create_downloads(self, urls: List[str], session: Session = None, downloader_name: str = None, - skip_download: bool = False, reset_attempts: bool = False, sub_downloader: str = None) \ + @optional_session + def create_downloads(self, urls: List[str], session: Session = None, downloader: str = None, + reset_attempts: bool = False, sub_downloader: str = None) \ -> List[Download]: """Schedule all URLs for download. If one cannot be downloaded, none will be added.""" - if not session: - _, session = get_db_context() + if not all(urls): + raise ValueError('Download must have a URL') downloads = [] - downloader = self.get_downloader_by_name(downloader_name) if downloader_name else None - if downloader_name and not downloader: - # Could not find the downloader - raise InvalidDownload(f'Unknown downloader {downloader_name}') + shared_downloader = self.get_downloader_by_name(downloader) if downloader else None + if downloader and not shared_downloader: + # Could not look up the downloader. + raise InvalidDownload(f'Unknown downloader {downloader}') with session.transaction: for url in urls: - if not url: - raise ValueError('Download must have a URL') if url in DOWNLOAD_MANAGER_CONFIG.skip_urls and reset_attempts: # User manually entered this download, remove it from the skip list. - DOWNLOAD_MANAGER_CONFIG.skip_urls = [i for i in DOWNLOAD_MANAGER_CONFIG.skip_urls if i != url] - DOWNLOAD_MANAGER_CONFIG.save() + self.remove_from_skip_list(url) elif url in DOWNLOAD_MANAGER_CONFIG.skip_urls: logger.warning(f'Skipping {url} because it is in the download_manager.yaml skip list.') continue info_json = None - if not downloader: + if not shared_downloader: # User has requested automatic downloader selection, try and find it. - downloader_, info_json = self.get_downloader(url) + local_downloader, info_json = self.get_downloader(url) else: - downloader_ = downloader + # One Downloader provided for all URLs, use it. + local_downloader = shared_downloader download = self.get_or_create_download(url, session) # Download may have failed, try again. download.renew(reset_attempts=reset_attempts) - download.downloader = downloader_.name + download.downloader = local_downloader.name download.info_json = info_json or download.info_json download.sub_downloader = sub_downloader downloads.append(download) - if skip_download is True: - return downloads - - if PYTEST: - # Download now (for testing). - self._do_downloads(session) - else: - # Start downloading ASAP. - self.start_downloads() + # Start downloading ASAP. + try: + asyncio.create_task(self.queue_downloads()) + except RuntimeError: + # Event loop isn't running. Probably testing? + if not PYTEST: + logger.info(f'Unable to queue downloads after creating {len(downloads)} download(s).') return downloads + def recurring_download(self, url: str, frequency: int, downloader: str = None, + sub_downloader: str = None, reset_attempts: bool = False) -> Download: + """Schedule a recurring download.""" + if not frequency: + raise ValueError('Recurring download must have a frequency!') + + with get_db_session(commit=True) as session: + download, = self.create_downloads([url, ], session=session, downloader=downloader, + sub_downloader=sub_downloader, reset_attempts=reset_attempts) + download.frequency = frequency + + return download + @wrol_mode_check @optional_session - def _do_downloads(self, session: Session = None): - """This method calls the Downloader's do_download method.""" + async def queue_downloads(self, session: Session = None): + """Put all downloads in queue. Will only queue downloads if there are workers to take them. Each worker + only receives one domain, this is to prevent downloading from one domain many times at once.""" if self.disabled.is_set(): raise InvalidDownload('DownloadManager is disabled') - try: - downloads = self.get_new_downloads(session) - - download_count = 0 - for download in downloads: - if self.disabled.is_set(): - raise InvalidDownload('DownloadManager is disabled') - if self.stopped: - return - - download_count += 1 - download_id = download.id - url = download.url - - if url in DOWNLOAD_MANAGER_CONFIG.skip_urls: - logger.info(f'Skipping download in skip_urls list: {url}') - session.query(Download).filter_by(id=download.id).delete() - session.commit() - continue - - downloader = self.get_downloader_by_name(download.downloader) - if not downloader: - downloader, info_json = self.get_downloader(download.url) - download.downloader = downloader.name - download.info_json = info_json - downloader.clear() - - download.started() - session.commit() - - try_again = True - try: - result = downloader.do_download(download) - except UnrecoverableDownloadError as e: - # Download failed and should not be retried. - logger.warning(f'UnrecoverableDownloadError for {url}', exc_info=e) - result = DownloadResult(success=False, error=str(traceback.format_exc())) - try_again = False - except Exception as e: - logger.warning(f'Failed to download {url}. Will be tried again later.', exc_info=e) - result = DownloadResult(success=False, error=str(traceback.format_exc())) - - error_len = len(result.error) if result.error else 0 - logger.debug(f'Got success={result.success} from {downloader} with {error_len=}') - - # Long-running, get the download again. - with get_db_session(commit=True) as session_: - download = self.get_download(session_, id_=download_id) - # Use a new location if provided, keep the old location if no new location is provided, otherwise - # clear out an outdated location. - download.location = result.location or download.location or None - # Clear any old errors if the download succeeded. - download.error = result.error if result.error else None - download.next_download = self.calculate_next_download(download, session) - - if result.downloads: - logger.info(f'Adding {len(result.downloads)} downloads from result of {download.url}') - self.create_downloads(result.downloads, session_, downloader_name=download.sub_downloader) - - if try_again is False: - download.fail() - elif result.success: - download.complete() - else: - download.defer() - - if download_count: - logger.info(f'Done doing {download_count} downloads.') - finally: - session.close() - - def do_downloads_sync(self): - acquired = DOWNLOAD_IN_PROGRESS.acquire(block=False) - if not acquired: + if len(self.data['processing_domains']) >= len(self.workers): return - try: - self._do_downloads() - finally: - DOWNLOAD_IN_PROGRESS.release() + # Find download whose domain isn't already being downloaded. + new_downloads = session.query(Download).filter( + Download.status == 'new', + Download.domain not in self.data['processing_domains'], + ).order_by(Download.id) # noqa + count = 0 + for download in new_downloads: + download.manager = self # Assign this Download to this manager. + domain = download.domain + if domain not in self.data['processing_domains']: + self._add_domain(domain) + await self.download_queue.put(download) + count += 1 + if count: + logger.debug(f'Added {count} downloads to queue.') async def do_downloads(self): - return self.do_downloads_sync() + """Schedule any downloads that are new. - def start_downloads(self): - """Start an async task to do downloads. This does nothing when testing!""" - if not PYTEST: - asyncio.create_task(self.do_downloads()) + Warning: Downloads will still be running even after this returns! See `wait_for_all_downloads`. + """ + if self.disabled.is_set(): + return + self.renew_recurring_downloads() + self.delete_old_once_downloads() + await self.queue_downloads() + + async def wait_for_all_downloads(self): + """Wait for all Downloads in queue AND any new Downloads to complete. + + THIS METHOD IS FOR TESTING. + """ + while True: + await asyncio.sleep(0.1) + + if not self.workers_running(): + raise ValueError('No workers are running!') + + await self.queue_downloads() + + try: + next(self.get_new_downloads()) + continue + except StopIteration: + pass + + if self.download_queue.empty(): + # Queue is empty. + break + await self.download_queue.join() @staticmethod def reset_downloads(): @@ -461,23 +585,26 @@ class DownloadManager: with get_db_curs(commit=True) as curs: curs.execute("UPDATE download SET status='new' WHERE status='pending' OR status='deferred'") - def recurring_download(self, url: str, frequency: int, skip_download: bool = False, downloader: str = None, - sub_downloader: str = None, reset_attempts: bool = False) -> Download: - """Schedule a recurring download.""" - with get_db_session(commit=True) as session: - download = self.create_download(url, session=session, skip_download=True, downloader=downloader, - sub_downloader=sub_downloader, reset_attempts=reset_attempts) - download.frequency = frequency - - if skip_download is False and PYTEST: - self.do_downloads_sync() - elif skip_download is False: - self.start_downloads() - - return download - DOWNLOAD_SORT = ('pending', 'failed', 'new', 'deferred', 'complete') + @optional_session + def get_new_downloads(self, session: Session) -> Generator[Download, None, None]: + """ + Get all "new" downloads. This method fetches the first download each iteration, so it will fetch downloads + that were created after calling it. + """ + last = None + while True: + download = session.query(Download).filter_by(status='new').order_by(Download.id).first() + if not download: + return + if download == last: + # Got the last download again. Is something wrong? + return + last = download + download.manager = self + yield download + @optional_session def get_recurring_downloads(self, session: Session = None, limit: int = None): """Get all Downloads that will be downloaded in the future.""" @@ -519,11 +646,7 @@ class DownloadManager: renewed = True if renewed: - logger.debug('Renewed recurring downloads') session.commit() - self.start_downloads() - else: - logger.debug('No recurring downloads to renew') def get_downloads(self, session: Session) -> List[Download]: downloads = list(session.query(Download).all()) @@ -566,11 +689,13 @@ class DownloadManager: def stop(self): """Stop all Downloaders, defer all pending downloads.""" + logger.warning('Stopping all downloads') self.stopped = True for downloader in self.instances: downloader.kill() for download in self.get_pending_downloads(): download.defer() + self.cancel_workers() def kill(self): """Kill all downloads. Do not start new downloads.""" @@ -582,15 +707,19 @@ class DownloadManager: def enable(self): """Enable downloading. Start downloading.""" + logger.info('Enabling downloading') + for downloader in self.instances: + downloader.clear() self.disabled.clear() - self.start_downloads() + self.start_workers() + asyncio.create_task(self.do_downloads()) FINISHED_STATUSES = ('complete', 'failed') def delete_old_once_downloads(self): - """ - Delete all once-downloads that have expired. Do not delete downloads that are new, or should be tried again. - """ + """Delete all once-downloads that have expired. + + Do not delete downloads that are new, or should be tried again.""" with get_db_session(commit=True) as session: downloads = self.get_once_downloads(session) one_month = now() - timedelta(days=30) @@ -674,8 +803,9 @@ class DownloadManager: download.manager = self return downloads + @staticmethod @optional_session - def calculate_next_download(self, download: Download, session: Session) -> Optional[datetime]: + def calculate_next_download(download: Download, session: Session) -> Optional[datetime]: """ If the download is "deferred", download soon. But, slowly increase the time between attempts. @@ -738,13 +868,22 @@ class DownloadManager: # Add all downloads to permanent skip list. ids = [i.id for i in failed_downloads] - urls = list(set(DOWNLOAD_MANAGER_CONFIG.skip_urls) | {i.url for i in failed_downloads}) - DOWNLOAD_MANAGER_CONFIG.skip_urls = urls + self.add_to_skip_list(*(i.url for i in failed_downloads)) # Delete all failed once-downloads. session.execute('DELETE FROM download WHERE id = ANY(:ids)', {'ids': ids}) session.commit() + @staticmethod + def add_to_skip_list(*urls: str): + DOWNLOAD_MANAGER_CONFIG.skip_urls = list(set(DOWNLOAD_MANAGER_CONFIG.skip_urls) | set(urls)) + DOWNLOAD_MANAGER_CONFIG.save() + + @staticmethod + def remove_from_skip_list(url: str): + DOWNLOAD_MANAGER_CONFIG.skip_urls = [i for i in DOWNLOAD_MANAGER_CONFIG.skip_urls if i != url] + DOWNLOAD_MANAGER_CONFIG.save() + # The global DownloadManager. This should be used everywhere! download_manager = DownloadManager() @@ -824,14 +963,12 @@ class RSSDownloader(Downloader, ABC): logger.warning(f'RSS entry {idx} did not have a link!') session = Session.object_session(download) - self.manager.create_downloads(urls, session=session, downloader_name=download.sub_downloader) + self.manager.create_downloads(urls, session=session, downloader=download.sub_downloader) return DownloadResult(success=True) @staticmethod def acceptable_bozo_errors(feed): - """ - Feedparser can report some errors, some we can ignore. - """ + """Feedparser can report some errors, some we can ignore.""" if 'document declared as' in str(feed['bozo_exception']): # Feed's encoding does not match what is declared, this is fine. return True diff --git a/wrolpi/root_api.py b/wrolpi/root_api.py index d481a1fc..f672e0a3 100644 --- a/wrolpi/root_api.py +++ b/wrolpi/root_api.py @@ -13,7 +13,6 @@ from sanic import Sanic, response, Blueprint, __version__ as sanic_version from sanic.blueprint_group import BlueprintGroup from sanic.request import Request from sanic.response import HTTPResponse -from sanic.signals import Event from sanic_ext import validate from sanic_ext.extensions.openapi import openapi @@ -26,7 +25,7 @@ from wrolpi.downloader import download_manager from wrolpi.errors import WROLModeEnabled, InvalidTimezone, API_ERRORS, APIError, ValidationError, HotspotError from wrolpi.media_path import MediaPath from wrolpi.schema import RegexRequest, RegexResponse, SettingsRequest, SettingsResponse, DownloadRequest, EchoResponse -from wrolpi.vars import DOCKERIZED, PYTEST +from wrolpi.vars import DOCKERIZED from wrolpi.version import __version__ logger = logger.getChild(__name__) @@ -210,7 +209,7 @@ async def post_download(_: Request, body: DownloadRequest): download_manager.recurring_download(urls[0], body.frequency, downloader=downloader, sub_downloader=body.sub_downloader, reset_attempts=True) else: - download_manager.create_downloads(urls, downloader_name=downloader, sub_downloader=body.sub_downloader, + download_manager.create_downloads(urls, downloader=downloader, sub_downloader=body.sub_downloader, reset_attempts=True) return response.empty() @@ -376,10 +375,3 @@ def json_error_handler(request, exception: Exception): api_app.error_handler.add(APIError, json_error_handler) - - -@api_app.signal(Event.SERVER_SHUTDOWN_BEFORE) -def handle_server_shutdown(*args, **kwargs): - """Stop downloads when server is shutting down.""" - if not PYTEST: - download_manager.stop() diff --git a/wrolpi/test/common.py b/wrolpi/test/common.py index 344ee401..6fa03feb 100644 --- a/wrolpi/test/common.py +++ b/wrolpi/test/common.py @@ -10,7 +10,6 @@ from typing import List import mock import pytest -from requests import Response from wrolpi.common import set_test_media_directory, get_media_directory from wrolpi.conftest import ROUTES_ATTACHED, test_db, test_client # noqa @@ -128,7 +127,7 @@ class TestAPI(ExtendedTestCase): def tearDown(self) -> None: set_test_media_directory(None) - def assertHTTPStatus(self, response: Response, status: int): + def assertHTTPStatus(self, response, status: int): self.assertEqual(response.status_code, status) assertOK = partialmethod(assertHTTPStatus, status=HTTPStatus.OK) diff --git a/wrolpi/test/test_db.py b/wrolpi/test/test_db.py index 8d2d43ca..821c9188 100644 --- a/wrolpi/test/test_db.py +++ b/wrolpi/test/test_db.py @@ -35,13 +35,3 @@ def test_optional_session(test_session): func(session=test_session) func(test_session) - -def test_test_session_is_required(): - """ - fixture `test_session` must be used in all tests to avoid using the real database. - """ - try: - with get_db_session() as session: - raise Exception('test engine was not used!!') - except ValueError as e: - assert 'test database' in str(e) diff --git a/wrolpi/test/test_downloader.py b/wrolpi/test/test_downloader.py index 6e3423dd..8fdd773f 100644 --- a/wrolpi/test/test_downloader.py +++ b/wrolpi/test/test_downloader.py @@ -1,12 +1,14 @@ +import asyncio from abc import ABC from datetime import datetime from itertools import zip_longest +from typing import Tuple, Optional from unittest import mock from unittest.mock import MagicMock import pytest -from wrolpi.dates import local_timezone, Seconds +from wrolpi.dates import local_timezone, Seconds, now from wrolpi.db import get_db_context from wrolpi.downloader import Downloader, Download, DownloadFrequency, DownloadResult from wrolpi.errors import UnrecoverableDownloadError, InvalidDownload, WROLModeEnabled @@ -35,6 +37,7 @@ class HTTPDownloader(Downloader): return url.startswith('https://') or url.startswith('http://'), None +@pytest.mark.asyncio def test_register_downloader(test_session, test_download_manager): """Downloaders can be registered and have specific priorities.""" assert test_download_manager.instances == tuple() @@ -54,84 +57,39 @@ def test_register_downloader(test_session, test_download_manager): test_download_manager.register_downloader(permissive_downloader) -def test_do_downloads(test_session, test_download_manager): - _, session = get_db_context() - - http_downloader = HTTPDownloader() - http_downloader.do_download = MagicMock() - http_downloader.do_download.return_value = DownloadResult(success=True) - test_download_manager.register_downloader(http_downloader) - - permissive_downloader = PermissiveDownloader(priority=100) - permissive_downloader.do_download = MagicMock() - permissive_downloader.do_download.return_value = DownloadResult(success=False) - test_download_manager.register_downloader(permissive_downloader) - - # https is handled by the HTTP Downloader. - test_download_manager.create_download('https://example.com', session) - http_downloader.do_download.assert_called_once() - permissive_downloader.do_download.assert_not_called() - assert test_download_manager.get_download(session, 'https://example.com') is not None - - http_downloader.do_download.reset_mock() - - # try the permissive download, which returns a failure. - test_download_manager.create_download('foo', session) - http_downloader.do_download.assert_not_called() - permissive_downloader.do_download.assert_called_once() - download = test_download_manager.get_download(session, 'foo') - assert download.attempts == 1 - - # try again - test_download_manager.create_download('foo', session) - test_download_manager._do_downloads(session) - download = test_download_manager.get_download(session, 'foo') - assert download.attempts == 2 - - # finally success - permissive_downloader.do_download.return_value = DownloadResult(success=True) - test_download_manager.create_download('foo', session) - test_download_manager._do_downloads(session) - download = test_download_manager.get_download(session, 'foo') - assert download.status == 'complete' - - # No downloads left. - assert list(test_download_manager.get_new_downloads(session)) == [] - - -def test_delete_old_once_downloads(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_delete_old_once_downloads(test_session, test_download_manager): """Once-downloads over a month old should be deleted.""" with mock.patch('wrolpi.downloader.now') as mock_now: mock_now.return_value = local_timezone(datetime(2020, 6, 5, 0, 0)) permissive_downloader = PermissiveDownloader(priority=0) test_download_manager.register_downloader(permissive_downloader) - _, session = get_db_context() # Recurring downloads should not be deleted. - d1 = test_download_manager.create_download('https://example.com/1', session, skip_download=True) - d2 = test_download_manager.create_download('https://example.com/2', session, skip_download=True) + d1 = test_download_manager.create_download('https://example.com/1') + d2 = test_download_manager.create_download('https://example.com/2') d1.frequency = 1 d2.frequency = 1 d2.started() # Should be deleted. - d3 = test_download_manager.create_download('https://example.com/3', session, skip_download=True) - d4 = test_download_manager.create_download('https://example.com/4', session, skip_download=True) + d3 = test_download_manager.create_download('https://example.com/3') + d4 = test_download_manager.create_download('https://example.com/4') d3.complete() d4.complete() d3.last_successful_download = local_timezone(datetime(2020, 1, 1, 0, 0, 0)) d4.last_successful_download = local_timezone(datetime(2020, 5, 1, 0, 0, 0)) # Not a month old. - d5 = test_download_manager.create_download('https://example.com/5', session, skip_download=True) + d5 = test_download_manager.create_download('https://example.com/5') d5.last_successful_download = local_timezone(datetime(2020, 6, 1, 0, 0, 0)) # An old, but pending download should not be deleted. - d6 = test_download_manager.create_download('https://example.com/6', session, skip_download=True) + d6 = test_download_manager.create_download('https://example.com/6') d6.last_successful_download = local_timezone(datetime(2020, 4, 1, 0, 0, 0)) d6.started() test_download_manager.delete_old_once_downloads() # Two old downloads are deleted. - downloads = session.query(Download).order_by(Download.id).all() + downloads = test_session.query(Download).order_by(Download.id).all() expected = [ dict(url='https://example.com/1', frequency=1, attempts=0, status='new'), dict(url='https://example.com/2', frequency=1, attempts=1, status='pending'), @@ -142,7 +100,8 @@ def test_delete_old_once_downloads(test_session, test_download_manager): assert_dict_contains(download.dict(), expected) -def test_ensure_download(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_ensure_download(test_session, test_download_manager): assert len(test_download_manager.get_downloads(test_session)) == 0 http_downloader = HTTPDownloader() @@ -150,7 +109,7 @@ def test_ensure_download(test_session, test_download_manager): http_downloader.do_download.return_value = DownloadResult(success=True) test_download_manager.register_downloader(http_downloader) - test_download_manager.create_download('https://example.com', test_session) + test_download_manager.create_download('https://example.com') download = test_download_manager.get_download(test_session, 'https://example.com') assert download.url == 'https://example.com' @@ -159,19 +118,20 @@ def test_ensure_download(test_session, test_download_manager): assert test_download_manager.get_download(test_session, 'https://example.com/not downloading') is None -def test_create_downloads(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_create_downloads(test_session, test_download_manager): """Multiple downloads can be scheduled using DownloadManager.create_downloads.""" http_downloader = HTTPDownloader(priority=0) test_download_manager.register_downloader(http_downloader) # Both URLs are valid and are scheduled. - test_download_manager.create_downloads(['https://example.com/1', 'https://example.com/2'], skip_download=True) + test_download_manager.create_downloads(['https://example.com/1', 'https://example.com/2']) downloads = test_download_manager.get_downloads(test_session) assert {i.url for i in downloads} == {'https://example.com/1', 'https://example.com/2'} # One URL is bad, neither should be scheduled. with pytest.raises(InvalidDownload): - test_download_manager.create_downloads(['https://example.com/3', 'bad url should fail'], skip_download=True) + test_download_manager.create_downloads(['https://example.com/3', 'bad url should fail']) downloads = test_download_manager.get_downloads(test_session) assert {i.url for i in downloads} == {'https://example.com/1', 'https://example.com/2'} @@ -195,12 +155,13 @@ def test_downloader_must_have_name(): @pytest.mark.asyncio async def test_download_wrol_mode(test_session, test_download_manager): with pytest.raises(WROLModeEnabled): - test_download_manager.do_downloads_sync() + await test_download_manager.do_downloads() with pytest.raises(WROLModeEnabled): - await test_download_manager.do_downloads_sync() + await test_download_manager.do_downloads() -def test_download_get_downloader(test_session, test_download_manager, ): +@pytest.mark.asyncio +async def test_download_get_downloader(test_session, test_download_manager): """ A Download can find it's Downloader. """ @@ -209,14 +170,15 @@ def test_download_get_downloader(test_session, test_download_manager, ): test_download_manager.register_downloader(permissive_downloader) test_download_manager.register_downloader(http_downloader) - download1 = test_download_manager.create_download('https://example.com', test_session, skip_download=True) + download1 = test_download_manager.create_download('https://example.com') assert download1.get_downloader() == permissive_downloader - download2 = test_download_manager.create_download('https://example.com', test_session, skip_download=True) + download2 = test_download_manager.create_download('https://example.com') download2.downloader = 'bad downloader' assert download2.get_downloader() is None +@pytest.mark.asyncio def test_calculate_next_download(test_session, test_download_manager, fake_now): fake_now(datetime(2000, 1, 1)) download = Download() @@ -249,79 +211,78 @@ def test_calculate_next_download(test_session, test_download_manager, fake_now): assert test_download_manager.calculate_next_download(d3) == local_timezone(datetime(2000, 1, 9, 18)) -def test_recurring_downloads(test_session, test_download_manager, fake_now): - """ - A recurring Download should be downloaded repeatedly forever. - """ - _, session = get_db_context() - +@pytest.mark.asyncio +async def test_recurring_downloads(test_session, test_download_manager, fake_now): + """A recurring Download should be downloaded forever.""" http_downloader = HTTPDownloader() http_downloader.do_download = MagicMock() http_downloader.do_download.return_value = DownloadResult(success=True) test_download_manager.register_downloader(http_downloader) # Download every hour. - test_download_manager.recurring_download('https://example.com', Seconds.hour, skip_download=True) + test_download_manager.recurring_download('https://example.com', Seconds.hour) # One download is scheduled. - downloads = test_download_manager.get_new_downloads(session) + downloads = test_download_manager.get_new_downloads(test_session) assert [(i.url, i.frequency) for i in downloads] == [('https://example.com', Seconds.hour)] - now = fake_now(datetime(2020, 1, 1, 0, 0, 0)) + now_ = fake_now(datetime(2020, 1, 1, 0, 0, 0)) # Download is processed and successful, no longer "new". - test_download_manager.do_downloads_sync() + await test_download_manager.wait_for_all_downloads() http_downloader.do_download.assert_called_once() - assert list(test_download_manager.get_new_downloads(session)) == [] - downloads = list(test_download_manager.get_recurring_downloads(session)) + assert list(test_download_manager.get_new_downloads(test_session)) == [] + downloads = list(test_download_manager.get_recurring_downloads(test_session)) assert len(downloads) == 1 download = downloads[0] expected = local_timezone(datetime(2020, 1, 1, 1, 0, 0)) assert download.next_download == expected - assert download.last_successful_download == now + assert download.last_successful_download == now_ # Download is not due for an hour. - test_download_manager.renew_recurring_downloads(session) - assert list(test_download_manager.get_new_downloads(session)) == [] - assert download.last_successful_download == now + test_download_manager.renew_recurring_downloads(test_session) + await test_download_manager.wait_for_all_downloads() + assert list(test_download_manager.get_new_downloads(test_session)) == [] + assert download.last_successful_download == now_ # Download is due an hour later. fake_now(datetime(2020, 1, 1, 2, 0, 1)) - test_download_manager.renew_recurring_downloads(session) - (download,) = list(test_download_manager.get_new_downloads(session)) + test_download_manager.renew_recurring_downloads(test_session) + (download,) = list(test_download_manager.get_new_downloads(test_session)) # Download is "new" but has not been downloaded a second time. assert download.next_download == expected - assert download.last_successful_download == now + assert download.last_successful_download == now_ assert download.status == 'new' # Try the download, but it fails. http_downloader.do_download.reset_mock() http_downloader.do_download.return_value = DownloadResult(success=False) - test_download_manager.do_downloads_sync() + await test_download_manager.wait_for_all_downloads() http_downloader.do_download.assert_called_once() - download = session.query(Download).one() + download = test_session.query(Download).one() # Download is deferred, last successful download remains the same. assert download.status == 'deferred' - assert download.last_successful_download == now + assert download.last_successful_download == now_ # Download should be retried after the DEFAULT_RETRY_FREQUENCY. expected = local_timezone(datetime(2020, 1, 1, 3, 0, 0, 997200)) assert download.next_download == expected # Try the download again, it finally succeeds. http_downloader.do_download.reset_mock() - now = fake_now(datetime(2020, 1, 1, 4, 0, 1)) + now_ = fake_now(datetime(2020, 1, 1, 4, 0, 1)) http_downloader.do_download.return_value = DownloadResult(success=True) - test_download_manager.renew_recurring_downloads(session) - test_download_manager.do_downloads_sync() + test_download_manager.renew_recurring_downloads(test_session) + await test_download_manager.wait_for_all_downloads() http_downloader.do_download.assert_called_once() - download = session.query(Download).one() + download = test_session.query(Download).one() assert download.status == 'complete' - assert download.last_successful_download == now + assert download.last_successful_download == now_ # Floats cause slightly wrong date. assert download.next_download == local_timezone(datetime(2020, 1, 1, 5, 0, 0, 997200)) -def test_max_attempts(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_max_attempts(test_session, test_download_manager): """A Download will only be attempted so many times, it will eventually be deleted.""" _, session = get_db_context() @@ -330,23 +291,27 @@ def test_max_attempts(test_session, test_download_manager): http_downloader.do_download.return_value = DownloadResult(success=True) test_download_manager.register_downloader(http_downloader) - test_download_manager.create_download('https://example.com', session) + test_download_manager.create_download('https://example.com') + await test_download_manager.wait_for_all_downloads() download = session.query(Download).one() assert download.attempts == 1 - test_download_manager.create_download('https://example.com', session) + test_download_manager.create_download('https://example.com') + await test_download_manager.wait_for_all_downloads() download = session.query(Download).one() assert download.attempts == 2 # There are no further attempts. http_downloader.do_download.side_effect = UnrecoverableDownloadError() - test_download_manager.create_download('https://example.com', session) + test_download_manager.create_download('https://example.com') + await test_download_manager.wait_for_all_downloads() download = session.query(Download).one() assert download.attempts == 3 assert download.status == 'failed' -def test_skip_urls(test_session, test_download_manager, test_client, assert_download_urls): +@pytest.mark.asyncio +async def test_skip_urls(test_session, test_download_manager, test_client, assert_download_urls): """The DownloadManager will not create downloads for URLs in it's skip list.""" _, session = get_db_context() from wrolpi.downloader import DOWNLOAD_MANAGER_CONFIG @@ -361,41 +326,79 @@ def test_skip_urls(test_session, test_download_manager, test_client, assert_down 'https://example.com/1', 'https://example.com/skipme', 'https://example.com/2', - ], downloader_name=HTTPDownloader.name) + ], downloader=HTTPDownloader.name) + await test_download_manager.wait_for_all_downloads() assert_download_urls({'https://example.com/1', 'https://example.com/2'}) assert DOWNLOAD_MANAGER_CONFIG.skip_urls == ['https://example.com/skipme'] test_download_manager.delete_completed() # The user can start a download even if the URL is in the skip list. - test_download_manager.create_download('https://example.com/skipme', session=test_session, reset_attempts=True) + test_download_manager.create_download('https://example.com/skipme', reset_attempts=True) assert_download_urls({'https://example.com/skipme'}) assert DOWNLOAD_MANAGER_CONFIG.skip_urls == [] -def test_process_runner_timeout(test_directory): +@pytest.mark.asyncio +async def test_process_runner_timeout(test_directory): """A Downloader can cancel it's download using a timeout.""" # Default timeout of 3 seconds. downloader = Downloader(0, 'downloader', timeout=3) # Default timeout is obeyed. start = datetime.now() - downloader.process_runner('https://example.com', ('sleep', '8'), test_directory) + await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory) elapsed = datetime.now() - start assert 3 < elapsed.total_seconds() < 4 # One-off timeout is obeyed. start = datetime.now() - downloader.process_runner('https://example.com', ('sleep', '8'), test_directory, timeout=1) + await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory, timeout=1) elapsed = datetime.now() - start assert 1 < elapsed.total_seconds() < 2 - # This Downloader's timeout is ignored in favor of the global timeout. + # Global timeout is obeyed. from wrolpi.common import WROLPI_CONFIG WROLPI_CONFIG.download_timeout = 3 - - # Global timeout is obeyed. start = datetime.now() - downloader.process_runner('https://example.com', ('sleep', '8'), test_directory) + await downloader.process_runner('https://example.com', ('sleep', '8'), test_directory) elapsed = datetime.now() - start assert 2 < elapsed.total_seconds() < 4 + + +@pytest.mark.asyncio +async def test_multi_domain_download(test_session, test_directory, test_download_manager): + class TestDownloader(Downloader, ABC): + name = 'test_downloader' + + @classmethod + def valid_url(cls, url: str) -> Tuple[bool, Optional[dict]]: + return True, {} + + async def do_download(self, download: Download) -> DownloadResult: + await asyncio.sleep(1) + return DownloadResult(success=True) + + async def wait_and_assert(urls_, max_seconds): + downloads = test_download_manager.create_downloads(urls_, downloader=test_downloader.name) + assert downloads and all(isinstance(i, Download) for i in downloads), 'Did not get Downloads' + start = now() + await test_download_manager.wait_for_all_downloads() + assert (elapsed := now() - start).total_seconds() < max_seconds, \ + f'Downloads took {elapsed} seconds. Were they in parallel?' + downloads = test_download_manager.get_downloads(test_session) + incomplete = [i.url for i in downloads if i.status != 'complete'] + assert not incomplete, f'Some downloads did not complete: {incomplete}' + + test_downloader = TestDownloader() + test_download_manager.register_downloader(test_downloader) + + # Each domain's downloads will be processed one at a time. + # example.com and example.org will both download in parallel. + urls = ['https://example.com/1', 'https://example.com/2', 'https://example.com/3', 'https://example.org/1', + 'https://example.org/2'] + await wait_and_assert(urls, 4) + + # Downloading more domains that workers is possible. + urls = [f'https://example.{i}' for i in range(test_download_manager.worker_count + 2)] + await wait_and_assert(urls, 3) diff --git a/wrolpi/test/test_root_api.py b/wrolpi/test/test_root_api.py index df488018..120398b3 100644 --- a/wrolpi/test/test_root_api.py +++ b/wrolpi/test/test_root_api.py @@ -15,6 +15,7 @@ from wrolpi.test.common import TestAPI, wrap_test_db, skip_circleci class TestRootAPI(TestAPI): + @wrap_test_db def test_index(self): """ Index should have some details in an HTML response @@ -61,6 +62,7 @@ class TestRootAPI(TestAPI): assert response.status_code == HTTPStatus.BAD_REQUEST assert json.loads(response.body) == {'valid': False, 'regex': '.*(missing parenthesis.*'} + @wrap_test_db def test_get_settings(self): with mock.patch('wrolpi.admin.CPUFREQ_INFO_BIN', '/usr/bin/cpufreq-info'), \ mock.patch('wrolpi.admin.NMCLI_BIN', '/usr/bin/nmcli'): diff --git a/wrolpi/test/test_rss.py b/wrolpi/test/test_rss.py index 962cf337..df984b7f 100644 --- a/wrolpi/test/test_rss.py +++ b/wrolpi/test/test_rss.py @@ -4,6 +4,7 @@ from itertools import zip_longest from typing import List import mock +import pytest from wrolpi.db import optional_session from wrolpi.downloader import Download, DownloadResult, RSSDownloader @@ -20,13 +21,8 @@ class RSSHTTPDownloader(HTTPDownloader): return download -def test_rss_download_invalid(test_session, test_download_manager): - """Feeds are not downloaded when FeedParser reports an issue.""" - rss_downloader = RSSDownloader() - test_download_manager.register_downloader(rss_downloader) - - -def test_rss_download(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_rss_download(test_session, test_download_manager): """An RSS Downloader will create new Downloads for every link in the feed.""" rss_downloader = RSSDownloader() test_download_manager.register_downloader(rss_downloader) @@ -50,7 +46,8 @@ def test_rss_download(test_session, test_download_manager): dict(link='https://example.com/c'), ] ) - test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http') + test_download_manager.create_download('https://example.com/feed', sub_downloader='http') + await test_download_manager.wait_for_all_downloads() # Feed download is complete. check_downloads([ @@ -68,7 +65,8 @@ def test_rss_download(test_session, test_download_manager): dict(link='https://example.com/d'), ] ) - test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http') + test_download_manager.create_download('https://example.com/feed', sub_downloader='http') + await test_download_manager.wait_for_all_downloads() # Only the new URLs are Archived. check_downloads([ @@ -80,28 +78,8 @@ def test_rss_download(test_session, test_download_manager): ]) -def test_post_rss_download(test_session, test_client, test_download_manager): - """An RSS can be downloaded in the UI""" - rss_downloader = RSSDownloader() - test_download_manager.register_downloader(rss_downloader) - http_downloader = RSSHTTPDownloader() - test_download_manager.register_downloader(http_downloader) - - with mock.patch('wrolpi.downloader.parse_feed') as mock_parse_feed: - mock_parse_feed.return_value = dict(bozo=0, entries=[]) - content = dict(urls='https://example.com/feed', downloader='rss', frequency=100) - request, response = test_client.post('/api/download', content=json.dumps(content)) - assert response.status == HTTPStatus.NO_CONTENT - - # Download was attempted. - (download,) = test_download_manager.get_downloads(test_session) - assert download.url == 'https://example.com/feed' - assert download.frequency == 100 - assert download.downloader == 'rss' - assert download.attempts == 1 - - -def test_rss_no_entries(test_session, test_download_manager): +@pytest.mark.asyncio +async def test_rss_no_entries(test_session, test_download_manager): """An RSS feed with no entries is handled.""" rss_downloader = RSSDownloader() test_download_manager.register_downloader(rss_downloader) @@ -109,6 +87,7 @@ def test_rss_no_entries(test_session, test_download_manager): with mock.patch('wrolpi.downloader.parse_feed') as mock_parse_feed: mock_parse_feed.return_value = dict(bozo=0, ) # missing `entries` test_download_manager.create_download('https://example.com/feed', test_session, sub_downloader='http') + await test_download_manager.wait_for_all_downloads() (download,) = test_download_manager.get_downloads(test_session) assert 'entries' in download.error diff --git a/wrolpi/test/test_timeout.py b/wrolpi/test/test_timeout.py deleted file mode 100644 index a84d098c..00000000 --- a/wrolpi/test/test_timeout.py +++ /dev/null @@ -1,47 +0,0 @@ -import time -import unittest -from unittest import mock - -from wrolpi.timeout import timeout, WorkerException - - -class TestTimeout(unittest.TestCase): - - def test_timeout(self): - """ - The process will be killed if it takes too long. - """ - - @timeout(2) - def func(): - time.sleep(5) - raise Exception('Was not killed!') - - self.assertRaises(TimeoutError, func) - - def test_timeout_exception(self): - """ - If the wrapped function has an error, it will raise a WorkerException. - """ - - @timeout(10) - def func(): - raise ValueError('Cause a worker error!') - - self.assertRaises(WorkerException, func) - - def test_test_timeout(self): - """Tests can overwrite the timeout""" - - @timeout(0.1) - def func(): - time.sleep(1) - return True - - # Killed after 0.1 seconds - self.assertRaises(TimeoutError, func) - - with mock.patch('wrolpi.timeout.TEST_TIMEOUT', 2): - # New timeout allows the function to finish. - result = func() - self.assertTrue(result) diff --git a/wrolpi/timeout.py b/wrolpi/timeout.py deleted file mode 100644 index 3322014e..00000000 --- a/wrolpi/timeout.py +++ /dev/null @@ -1,114 +0,0 @@ -""" -Wrap a function with `timeout` to kill it after some amount of seconds, the wrapped function is run in it's own -process. If the timeout is reached, kill the process and raise a TimeoutError. - -Log any errors in the global logger. If the worker process raises an Exception, raise an Exception sharing it's -description, log the traceback. - ->>> @timeout(5) ->>> def long_running_function(): ->>> while True: ->>> # Run infinitely ->>> time.sleep(1) - ->>> long_running_function() -TimeoutError() -""" -import multiprocessing.connection -import time -import traceback -from datetime import timedelta -from functools import wraps -from multiprocessing import Process, Pipe - -from wrolpi.common import logger -from wrolpi.dates import utc_now - -logger = logger.getChild(__name__) - - -class WorkerException(Exception): - pass - - -TEST_TIMEOUT = None -KILL_ATTEMPTS = 100 - - -def timeout(seconds: float) -> callable: - """ - Create a worker process, let it run until the timeout, then kill it. If the worker isn't killed, return - it's result or exception. If it is killed, raise a TimeoutError. - """ - - def wrapper(func): - - def worker_wrapper(pipe: multiprocessing.connection.Connection, *args, **kwargs): - """ - Wraps a worker function so we can catch it's Exception, or return its results. - """ - to_send = {} - try: - result = func(*args, **kwargs) - to_send['result'] = result - except Exception as e: - to_send.update({'exception': str(e), 'traceback': traceback.format_exc()}) - finally: - pipe.send(to_send) - - @wraps(func) - def wrapped(*args, **kwargs): - # unittest may override the timeout - _seconds = TEST_TIMEOUT or seconds - - # Start the worker_wrapper, give it a way to communicate back. - parent_conn, child_conn = Pipe() - args = (child_conn,) + args - worker = Process(target=worker_wrapper, args=args, kwargs=kwargs) - worker.start() - - # Continually attempt to get a result until the timeout is reached. - kill_time = utc_now() + timedelta(seconds=_seconds) - while True: - if not worker.is_alive(): - # Worker exited. - break - if not utc_now() < kill_time: - # Timeout was reached. - break - if parent_conn.poll(0.1): # This prevents a "busy loop" - # Result is waiting. - break - - # Get the result, if any. - result = None - if parent_conn.poll(0): - result = parent_conn.recv() - - attempts = 0 - while worker.is_alive(): - # Worker has run out of time. - worker.kill() - time.sleep(0.1) - attempts += 1 - if attempts > KILL_ATTEMPTS: - logger.warning(f'Failed to kill worker! {func}') - break - - worker.join() - - if result: - if 'result' in result: - return result['result'] - elif 'exception' in result: - exception = result['exception'] - _traceback = result['traceback'] - logger.warning(f'Traceback from worker: {_traceback}') - raise WorkerException(f'Exception from worker: {exception}') - - if attempts > 0: - raise TimeoutError(f'Killed worker after {_seconds} seconds: {func}') - - return wrapped - - return wrapper diff --git a/wrolpi/version.py b/wrolpi/version.py index dabeb96b..11320bab 100644 --- a/wrolpi/version.py +++ b/wrolpi/version.py @@ -1,4 +1,4 @@ -__version__ = '0.5.3-beta' +__version__ = '0.6.0-beta' import subprocess