Coverage for an_website / utils / background_tasks.py: 44.231%
52 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-24 17:35 +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 try:
59 await es.transport.perform_request("HEAD", "/")
60 except Exception: # pylint: disable=broad-except
61 EVENT_ELASTICSEARCH.clear()
62 LOGGER.exception(
63 "Connecting to Elasticsearch failed on worker: %s", worker
64 )
65 else:
66 if not EVENT_ELASTICSEARCH.is_set():
67 try:
68 await setup_elasticsearch_configs(
69 es, app.settings["ELASTICSEARCH_PREFIX"]
70 )
71 except Exception: # pylint: disable=broad-except
72 LOGGER.exception(
73 "An exception occured while configuring Elasticsearch on worker: %s", # noqa: B950
74 worker,
75 )
76 else:
77 EVENT_ELASTICSEARCH.set()
78 await asyncio.sleep(20)
81async def check_if_ppid_changed(ppid: int) -> None:
82 """Check whether Technoblade hates us."""
83 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
84 if os.getppid() != ppid:
85 EVENT_SHUTDOWN.set()
86 return
87 await asyncio.sleep(1)
90async def check_redis(
91 app: Application, worker: int | None
92) -> None: # pragma: no cover
93 """Check Redis."""
94 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
95 redis: Redis[str] = cast("Redis[str]", app.settings.get("REDIS"))
96 try:
97 await redis.ping()
98 except Exception: # pylint: disable=broad-except
99 EVENT_REDIS.clear()
100 LOGGER.exception("Connecting to Redis failed on worker %s", worker)
101 else:
102 EVENT_REDIS.set()
103 await asyncio.sleep(20)
106async def heartbeat() -> None:
107 """Heartbeat."""
108 global HEARTBEAT # pylint: disable=global-statement
109 while HEARTBEAT: # pylint: disable=while-used
110 HEARTBEAT = time.monotonic()
111 await asyncio.sleep(0.05)
114async def wait_for_shutdown() -> None: # pragma: no cover
115 """Wait for the shutdown event."""
116 loop = asyncio.get_running_loop()
117 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used
118 await asyncio.sleep(0.05)
119 loop.stop()
122def start_background_tasks( # pylint: disable=too-many-arguments
123 *,
124 app: Application,
125 processes: int,
126 module_infos: Iterable[ModuleInfo],
127 loop: asyncio.AbstractEventLoop,
128 main_pid: int,
129 elasticsearch_is_enabled: bool,
130 redis_is_enabled: bool,
131 worker: int | None,
132) -> Set[asyncio.Task[None]]:
133 """Start all required background tasks."""
135 async def execute_background_task(task: BackgroundTask, /) -> None:
136 """Execute a background task with error handling."""
137 try:
138 await task(app=app, worker=worker)
139 except asyncio.exceptions.CancelledError:
140 pass
141 except BaseException as exc: # pylint: disable=broad-exception-caught
142 LOGGER.exception(
143 "A %s exception occured while executing background task %s.%s",
144 exc.__class__.__name__,
145 task.__module__,
146 task.__name__,
147 )
148 if not isinstance(exc, Exception):
149 raise
150 else:
151 LOGGER.debug(
152 "Background task %s.%s finished executing",
153 task.__module__,
154 task.__name__,
155 )
157 background_tasks: set[asyncio.Task[None]] = set()
159 def create_task(fun: BackgroundTask, /) -> asyncio.Task[None]:
160 """Create an asyncio.Task object from a BackgroundTask."""
161 name = f"{fun.__module__}.{fun.__name__}"
162 if not worker: # log only once
163 LOGGER.info("starting %s background task", name)
164 task = loop.create_task(execute_background_task(fun), name=name)
165 task.add_done_callback(background_tasks.discard)
166 return task
168 task_stream: typed_stream.Stream[asyncio.Task[None]] = assert_type(
169 typed_stream.Stream(module_infos)
170 .flat_map(lambda info: info.required_background_tasks)
171 .chain(
172 typed_stream.Stream((heartbeat, wait_for_shutdown)).map(
173 lambda fun: wraps(fun)(lambda **_: fun())
174 )
175 )
176 .chain(
177 [
178 wraps(check_if_ppid_changed)(
179 lambda **k: check_if_ppid_changed(main_pid)
180 )
181 ]
182 if processes
183 else ()
184 )
185 .chain([check_elasticsearch] if elasticsearch_is_enabled else ())
186 .chain([check_redis] if redis_is_enabled else ())
187 .distinct()
188 .map(create_task),
189 typed_stream.Stream[asyncio.Task[None]],
190 )
192 background_tasks.update(task_stream)
194 return background_tasks