Creating background_task function to keep set of tasks from being cleaned up during garbage collection.

This commit is contained in:
Roland 2022-10-22 14:26:59 -06:00
parent 6fa6be50fe
commit 1230a3056d
No known key found for this signature in database
GPG key ID: 9D6DBD3ACF11AB93
11 changed files with 65 additions and 43 deletions

View file

@ -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()

View file

@ -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']

View file

@ -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

View file

@ -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:

View file

@ -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()

View file

@ -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

View file

@ -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)

View file

@ -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,

View file

@ -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,

View file

@ -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),

View file

@ -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.