Coverage for an_website/utils/background_tasks.py: 44.231%
52 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-10 18:56 +0000
1# This program is free software: you can redistribute it and/or modify
2# it under the terms of the GNU Affero General Public License as
3# published by the Free Software Foundation, either version 3 of the
4# License, or (at your option) any later version.
5#
6# This program is distributed in the hope that it will be useful,
7# but WITHOUT ANY WARRANTY; without even the implied warranty of
8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
9# GNU Affero General Public License for more details.
10#
11# You should have received a copy of the GNU Affero General Public License
12# along with this program. If not, see <https://www.gnu.org/licenses/>.
13"""Tasks running in the background."""
15import asyncio
16import logging
17import os
18import time
19from collections.abc import Iterable, Set
20from functools import wraps
21from typing import TYPE_CHECKING, Final, Protocol, assert_type, cast
23import typed_stream
24from elasticsearch import AsyncElasticsearch
25from redis.asyncio import Redis
26from tornado.web import Application
28from .. import EVENT_ELASTICSEARCH, EVENT_REDIS, EVENT_SHUTDOWN
29from .elasticsearch_setup import setup_elasticsearch_configs
31if TYPE_CHECKING:
32 from .utils import ModuleInfo
34LOGGER: Final = logging.getLogger(__name__)
36HEARTBEAT: float = 0
39class BackgroundTask(Protocol):
40 """A protocol representing a background task."""
42 async def __call__(self, *, app: Application, worker: int | None) -> None:
43 """Start the background task."""
45 @property
46 def __name__(self) -> str: # pylint: disable=bad-dunder-name
47 """The name of the task."""
50async def check_elasticsearch(
51 app: Application, worker: int | None
52) -> None: # pragma: no cover
53 """Check Elasticsearch."""
54 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
55 es: AsyncElasticsearch = cast(
56 AsyncElasticsearch, app.settings.get("ELASTICSEARCH")
57 )
58 if not await es.ping():
59 EVENT_ELASTICSEARCH.clear()
60 LOGGER.error(
61 "Connecting to Elasticsearch failed on worker: %s", worker
62 )
63 elif not EVENT_ELASTICSEARCH.is_set():
64 try:
65 await setup_elasticsearch_configs(
66 es, app.settings["ELASTICSEARCH_PREFIX"]
67 )
68 except Exception: # pylint: disable=broad-except
69 LOGGER.exception(
70 "An exception occured while configuring Elasticsearch on worker: %s", # noqa: B950
71 worker,
72 )
73 else:
74 EVENT_ELASTICSEARCH.set()
75 await asyncio.sleep(20)
78async def check_if_ppid_changed(ppid: int) -> None:
79 """Check whether Technoblade hates us."""
80 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
81 if os.getppid() != ppid:
82 EVENT_SHUTDOWN.set()
83 return
84 await asyncio.sleep(1)
87async def check_redis(
88 app: Application, worker: int | None
89) -> None: # pragma: no cover
90 """Check Redis."""
91 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
92 redis: Redis[str] = cast("Redis[str]", app.settings.get("REDIS"))
93 try:
94 await redis.ping()
95 except Exception: # pylint: disable=broad-except
96 EVENT_REDIS.clear()
97 LOGGER.exception("Connecting to Redis failed on worker %s", worker)
98 else:
99 EVENT_REDIS.set()
100 await asyncio.sleep(20)
103async def heartbeat() -> None:
104 """Heartbeat."""
105 global HEARTBEAT # pylint: disable=global-statement
106 while HEARTBEAT: # pylint: disable=while-used
107 HEARTBEAT = time.monotonic()
108 await asyncio.sleep(0.05)
111async def wait_for_shutdown() -> None: # pragma: no cover
112 """Wait for the shutdown event."""
113 loop = asyncio.get_running_loop()
114 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
115 await asyncio.sleep(0.05)
116 loop.stop()
119def start_background_tasks( # pylint: disable=too-many-arguments
120 *,
121 app: Application,
122 processes: int,
123 module_infos: Iterable[ModuleInfo],
124 loop: asyncio.AbstractEventLoop,
125 main_pid: int,
126 elasticsearch_is_enabled: bool,
127 redis_is_enabled: bool,
128 worker: int | None,
129) -> Set[asyncio.Task[None]]:
130 """Start all required background tasks."""
132 async def execute_background_task(task: BackgroundTask, /) -> None:
133 """Execute a background task with error handling."""
134 try:
135 await task(app=app, worker=worker)
136 except asyncio.exceptions.CancelledError:
137 pass
138 except BaseException as exc: # pylint: disable=broad-exception-caught
139 LOGGER.exception(
140 "A %s exception occured while executing background task %s.%s",
141 exc.__class__.__name__,
142 task.__module__,
143 task.__name__,
144 )
145 if not isinstance(exc, Exception):
146 raise
147 else:
148 LOGGER.debug(
149 "Background task %s.%s finished executing",
150 task.__module__,
151 task.__name__,
152 )
154 background_tasks: set[asyncio.Task[None]] = set()
156 def create_task(fun: BackgroundTask, /) -> asyncio.Task[None]:
157 """Create an asyncio.Task object from a BackgroundTask."""
158 name = f"{fun.__module__}.{fun.__name__}"
159 if not worker: # log only once
160 LOGGER.info("starting %s background task", name)
161 task = loop.create_task(execute_background_task(fun), name=name)
162 task.add_done_callback(background_tasks.discard)
163 return task
165 task_stream: typed_stream.Stream[asyncio.Task[None]] = assert_type(
166 typed_stream.Stream(module_infos)
167 .flat_map(lambda info: info.required_background_tasks)
168 .chain(
169 typed_stream.Stream((heartbeat, wait_for_shutdown)).map(
170 lambda fun: wraps(fun)(lambda **_: fun())
171 )
172 )
173 .chain(
174 [
175 wraps(check_if_ppid_changed)(
176 lambda **k: check_if_ppid_changed(main_pid)
177 )
178 ]
179 if processes
180 else ()
181 )
182 .chain([check_elasticsearch] if elasticsearch_is_enabled else ())
183 .chain([check_redis] if redis_is_enabled else ())
184 .distinct()
185 .map(create_task),
186 typed_stream.Stream[asyncio.Task[None]],
187 )
189 background_tasks.update(task_stream)
191 return background_tasks