From 1230a3056d5aafea053fa428587ec401232b52cc Mon Sep 17 00:00:00 2001 From: Roland Date: Sat, 22 Oct 2022 14:26:59 -0600 Subject: [PATCH] Creating background_task function to keep set of tasks from being cleaned up during garbage collection. --- modules/map/api.py | 4 ++-- modules/map/test/test_api.py | 13 ++---------- wrolpi/common.py | 18 ++++++++++++++++- wrolpi/downloader.py | 19 +++++++++--------- wrolpi/files/api.py | 4 ++-- wrolpi/files/lib.py | 5 +++-- wrolpi/files/test/test_lib.py | 2 +- wrolpi/root_api.py | 1 + wrolpi/status.py | 38 ++++++++++++++++++++++------------- wrolpi/test/test_common.py | 2 ++ wrolpi/vars.py | 2 +- 11 files changed, 65 insertions(+), 43 deletions(-) diff --git a/modules/map/api.py b/modules/map/api.py index 21980b54..e514ef88 100644 --- a/modules/map/api.py +++ b/modules/map/api.py @@ -7,7 +7,7 @@ from sanic_ext import validate from sanic_ext.extensions.openapi import openapi from modules.map import lib, schema -from wrolpi.common import wrol_mode_check, get_media_directory +from wrolpi.common import wrol_mode_check, get_media_directory, background_task from wrolpi.errors import ValidationError from wrolpi.root_api import get_blueprint, json_response from wrolpi.vars import PYTEST, DOCKERIZED @@ -34,7 +34,7 @@ async def import_pbfs(_: Request, body: schema.ImportPost): if PYTEST: await coro else: - asyncio.create_task(coro) + background_task(coro) return response.empty() diff --git a/modules/map/test/test_api.py b/modules/map/test/test_api.py index 23e98706..00542429 100644 --- a/modules/map/test/test_api.py +++ b/modules/map/test/test_api.py @@ -86,17 +86,8 @@ def test_status_and_import(test_client, test_session, test_directory): def test_empty_import(test_client, test_session, test_directory): """Some files must be requested.""" with mock.patch('modules.map.lib.import_file') as mock_import_file: - body = {'pbfs': []} - request, response = test_client.post('/api/map/import', content=json.dumps(body)) - assert response.status_code == HTTPStatus.BAD_REQUEST - mock_import_file.assert_not_called() - - body = {'dumps': []} - request, response = test_client.post('/api/map/import', content=json.dumps(body)) - assert response.status_code == HTTPStatus.BAD_REQUEST - mock_import_file.assert_not_called() - - body = {'pbfs': [], 'dumps': []} + body = {'files': []} request, response = test_client.post('/api/map/import', content=json.dumps(body)) assert response.status_code == HTTPStatus.BAD_REQUEST mock_import_file.assert_not_called() + assert 'validate' in response.json['error'] diff --git a/wrolpi/common.py b/wrolpi/common.py index 26b1e734..1594ef40 100644 --- a/wrolpi/common.py +++ b/wrolpi/common.py @@ -10,6 +10,7 @@ import re import string import sys import tempfile +from asyncio import Task from copy import deepcopy from datetime import datetime, date from decimal import Decimal @@ -87,7 +88,8 @@ __all__ = [ 'timer', 'cum_timer', 'limit_concurrent', - 'truncate_object_bytes' + 'truncate_object_bytes', + 'background_task', ] # Base is used for all SQLAlchemy models. @@ -996,3 +998,17 @@ def truncate_object_bytes(obj: Union[List[str], str, None], maximum_bytes: int) index = -1 * round(len(obj) * .2) obj = obj[:index or -1] return truncate_object_bytes(obj, maximum_bytes) + + +BACKGROUND_TASKS = set() + + +def background_task(coro) -> Task: + """Convenience function which creates an asyncio task for the provided coroutine. + + The task is stored in a global set of background tasks so the task will not be discarded by the garbage collector. + """ + task = asyncio.create_task(coro) + BACKGROUND_TASKS.add(task) + task.add_done_callback(BACKGROUND_TASKS.discard) + return task diff --git a/wrolpi/downloader.py b/wrolpi/downloader.py index 7acd3757..c8cf48c4 100644 --- a/wrolpi/downloader.py +++ b/wrolpi/downloader.py @@ -21,7 +21,8 @@ from sqlalchemy import Column, Integer, String, Text from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Session -from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, ConfigFile, WROLPI_CONFIG +from wrolpi.common import Base, ModelHelper, logger, wrol_mode_check, zig_zag, ConfigFile, WROLPI_CONFIG, \ + background_task from wrolpi.dates import TZDateTime, now, Seconds, local_timezone, recursive_replace_tz from wrolpi.db import get_db_session, get_db_curs, optional_session from wrolpi.errors import InvalidDownload, UnrecoverableDownloadError @@ -323,8 +324,8 @@ class DownloadManager: name = f'{pid}.{num}' worker_logger = logger.getChild(f'download_worker.{name}') - disabled = 'is disabled' if self.disabled.is_set() else 'is NOT disabled' - worker_logger.info(f'Starting up. DownloadManager {disabled}.') + disabled = 'disabled' if self.disabled.is_set() else 'enabled' + worker_logger.info(f'Starting up. DownloadManager is {disabled}.') last_heartbeat = now() while True: @@ -406,7 +407,7 @@ class DownloadManager: # 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()) + background_task(self.queue_downloads()) except asyncio.CancelledError: worker_logger.warning('Canceled!') self.download_queue.task_done() @@ -475,9 +476,9 @@ class DownloadManager: await download_manager.do_downloads() await asyncio.sleep(30) - asyncio.create_task(_perpetual_download()) + background_task(_perpetual_download()) - asyncio.create_task(_perpetual_download()) + background_task(_perpetual_download()) def get_downloader(self, url: str) -> Tuple[Downloader, Dict]: for i in self.instances: @@ -555,7 +556,7 @@ class DownloadManager: # Start downloading ASAP. try: - asyncio.create_task(self.queue_downloads()) + background_task(self.queue_downloads()) except RuntimeError: # Event loop isn't running. Probably testing? if not PYTEST: @@ -788,8 +789,8 @@ class DownloadManager: self.start_workers(loop) try: - asyncio.create_task(self.perpetual_download()) - asyncio.create_task(self.do_downloads()) + background_task(self.perpetual_download()) + background_task(self.do_downloads()) except RuntimeError: # This may not work while testing. if not PYTEST: diff --git a/wrolpi/files/api.py b/wrolpi/files/api.py index 50627d00..62dc71b0 100644 --- a/wrolpi/files/api.py +++ b/wrolpi/files/api.py @@ -7,7 +7,7 @@ from sanic import response from sanic_ext import validate from sanic_ext.extensions.openapi import openapi -from wrolpi.common import get_media_directory, wrol_mode_check +from wrolpi.common import get_media_directory, wrol_mode_check, background_task from wrolpi.errors import InvalidFile from wrolpi.root_api import get_blueprint, json_response from . import lib, schema @@ -88,7 +88,7 @@ async def refresh_directory(_: Request, body: schema.DirectoryRefreshRequest): if PYTEST: await lib.refresh_directory_files_recursively(directory) else: - asyncio.create_task(lib.refresh_directory_files_recursively(directory)) + background_task(lib.refresh_directory_files_recursively(directory)) return response.empty() diff --git a/wrolpi/files/lib.py b/wrolpi/files/lib.py index 00a72194..cc46ecde 100644 --- a/wrolpi/files/lib.py +++ b/wrolpi/files/lib.py @@ -14,7 +14,8 @@ from sqlalchemy.orm import Session from wrolpi.cmd import which from wrolpi.common import get_media_directory, wrol_mode_check, logger, limit_concurrent, \ - get_files_and_directories, apply_modelers, apply_after_refresh, get_model_by_table_name, chunks_by_name + get_files_and_directories, apply_modelers, apply_after_refresh, get_model_by_table_name, chunks_by_name, \ + background_task from wrolpi.dates import now from wrolpi.db import get_db_session, get_db_curs, get_ranked_models from wrolpi.errors import InvalidFile @@ -203,7 +204,7 @@ def cancelable_wrapper(func: callable): if PYTEST: return await func(*args, **kwargs) - task = asyncio.create_task(func(*args, **kwargs)) + task = background_task(func(*args, **kwargs)) REFRESH_TASKS.append(task) return wrapped diff --git a/wrolpi/files/test/test_lib.py b/wrolpi/files/test/test_lib.py index db84e49e..a008514e 100644 --- a/wrolpi/files/test/test_lib.py +++ b/wrolpi/files/test/test_lib.py @@ -212,7 +212,7 @@ async def test_refresh_cancel(test_session, make_files_structure, test_directory task_.cancel() with pytest.raises(asyncio.CancelledError): await task_ - assert (datetime.now() - before).total_seconds() < 0.5, 'Task took too long. Was the refresh canceled?' + assert (datetime.now() - before).total_seconds() < 0.8, 'Task took too long. Was the refresh canceled?' task = asyncio.create_task(lib.refresh_files()) await assert_cancel(task) diff --git a/wrolpi/root_api.py b/wrolpi/root_api.py index 4c67b69b..2f8257c1 100644 --- a/wrolpi/root_api.py +++ b/wrolpi/root_api.py @@ -344,6 +344,7 @@ async def get_status(_: Request): bandwidth=s.bandwidth, cpu_info=s.cpu_info, disk_bandwidth=s.disk_bandwidth, + dockerized=DOCKERIZED, downloads=downloads, drives=s.drives, hotspot_status=admin.hotspot_status().name, diff --git a/wrolpi/status.py b/wrolpi/status.py index 6676de3e..7b75741d 100755 --- a/wrolpi/status.py +++ b/wrolpi/status.py @@ -12,6 +12,7 @@ from typing import List, Optional, Tuple from wrolpi.cmd import which from wrolpi.common import logger, limit_concurrent from wrolpi.dates import now +from wrolpi.vars import DOCKERIZED try: import psutil @@ -358,21 +359,22 @@ def get_nic_names() -> List[str]: @dataclass class DiskBandwidthInfo: - bytes_read: int = None - bytes_write: int = None + bytes_read_ps: int = None + bytes_write_ps: int = None elapsed: int = None name: str = None def __json__(self): return dict( - bytes_read=self.bytes_read, - bytes_write=self.bytes_write, + bytes_read_ps=self.bytes_read_ps, + bytes_write_ps=self.bytes_write_ps, elapsed=self.elapsed, name=self.name, ) -DRIVE_BANDWIDTH = multiprocessing.Manager().dict() +IGNORED_DISK_NAMES = ('loop',) +DISKS_BANDWIDTH = multiprocessing.Manager().dict() async def get_bandwidth_info() -> Tuple[List[NICBandwidthInfo], List[DiskBandwidthInfo]]: @@ -392,14 +394,22 @@ async def get_bandwidth_info() -> Tuple[List[NICBandwidthInfo], List[DiskBandwid name=name, speed=nic['speed'], )) - for name in sorted(DRIVE_BANDWIDTH.keys()): - disk = DRIVE_BANDWIDTH[name] + used_disks = [] + for name in sorted(DISKS_BANDWIDTH.keys()): + disk = DISKS_BANDWIDTH[name] if 'bytes_read_ps' not in disk: # No status collected yet continue + if any(name.startswith(i) for i in used_disks): + # Report only the first disk. Do not report it's partitions. Disks are sorted to make this work! + # i.e. report "sda" but not "sda1" or "sda2". + continue + if any(name.startswith(i) for i in IGNORED_DISK_NAMES): + continue + used_disks.append(name) disks_info.append(DiskBandwidthInfo( - bytes_read=disk['bytes_read_ps'], - bytes_write=disk['bytes_write_ps'], + bytes_read_ps=disk['bytes_read_ps'], + bytes_write_ps=disk['bytes_write_ps'], elapsed=disk['elapsed'], name=name, )) @@ -458,11 +468,11 @@ async def bandwidth_worker(count: int = None): timestamp = now().timestamp() for name_, disk in psutil.disk_io_counters(perdisk=True).items(): tic = timestamp, disk.read_bytes, disk.write_bytes - if bw := DRIVE_BANDWIDTH.get(name_): + if bw := DISKS_BANDWIDTH.get(name_): bw['historical'] = (bw['historical'] + [tic, ])[-21:] - DRIVE_BANDWIDTH.update({name_: bw}) + DISKS_BANDWIDTH.update({name_: bw}) else: - DRIVE_BANDWIDTH.update({ + DISKS_BANDWIDTH.update({ name_: dict(historical=[tic, ]), }) @@ -489,13 +499,13 @@ async def bandwidth_worker(count: int = None): speed=historical[-1][-1], # Use the most recent speed. ) }) - for name, stats in DRIVE_BANDWIDTH.items(): + for name, stats in DISKS_BANDWIDTH.items(): if 'historical' not in stats: continue historical = stats['historical'] bytes_read_ps, bytes_write_ps, elapsed = _calculate_bytes_per_second(historical) - DRIVE_BANDWIDTH.update({ + DISKS_BANDWIDTH.update({ name: dict( historical=historical, bytes_read_ps=bytes_read_ps, diff --git a/wrolpi/test/test_common.py b/wrolpi/test/test_common.py index 8d64f152..ffb06f35 100644 --- a/wrolpi/test/test_common.py +++ b/wrolpi/test/test_common.py @@ -425,7 +425,9 @@ def test_chunks_by_name(test_directory, make_files_structure): ('True', True), ('1', True), ('yes', True), + ('y', True), ('no', False), + ('n', False), ('0', False), ('False', False), ('false', False), diff --git a/wrolpi/vars.py b/wrolpi/vars.py index 69e834bd..f33fa206 100644 --- a/wrolpi/vars.py +++ b/wrolpi/vars.py @@ -40,7 +40,7 @@ def truthy_arg(value: str) -> bool: False """ v = str(value).lower() - return v == 'true' or v == 't' or v == '1' or v == 'yes' + return v == 'true' or v == 't' or v == '1' or v == 'yes' or v == 'y' # Get the media directory from the environment.