Coverage for an_website/utils/background_tasks.py: 41.379%

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

14 

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 

22 

23import typed_stream 

24from elasticsearch import ApiError, AsyncElasticsearch, TransportError 

25from redis.asyncio import Redis 

26from tornado.web import Application 

27 

28from .. import EVENT_ELASTICSEARCH, EVENT_REDIS, EVENT_SHUTDOWN 

29from .elasticsearch_setup import setup_elasticsearch_configs 

30 

31if TYPE_CHECKING: 

32 from .utils import ModuleInfo 

33 

34LOGGER: Final = logging.getLogger(__name__) 

35 

36HEARTBEAT: float = 0 

37 

38 

39class BackgroundTask(Protocol): 

40 """A protocol representing a background task.""" 

41 

42 async def __call__(self, *, app: Application, worker: int | None) -> None: 

43 """Start the background task.""" 

44 

45 @property 

46 def __name__(self) -> str: # pylint: disable=bad-dunder-name 

47 """The name of the task.""" 

48 

49 

50async def _try_ping_elastic_search(es: AsyncElasticsearch) -> Exception | None: 

51 """Return an exception if the info() API failed. 

52 

53 See: AsyncElasticsearch.ping for a similar implementation. 

54 """ 

55 try: 

56 await es.perform_request( 

57 "HEAD", 

58 "/", 

59 headers={"accept": "application/json"}, 

60 endpoint_id="ping", 

61 ) 

62 return None 

63 except (ApiError, TransportError) as exc: 

64 return exc 

65 

66 

67async def check_elasticsearch( 

68 app: Application, worker: int | None 

69) -> None: # pragma: no cover 

70 """Check Elasticsearch.""" 

71 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used 

72 es: AsyncElasticsearch = cast( 

73 AsyncElasticsearch, app.settings.get("ELASTICSEARCH") 

74 ) 

75 if exc := await _try_ping_elastic_search(es): 

76 EVENT_ELASTICSEARCH.clear() 

77 LOGGER.error( 

78 "Connecting to Elasticsearch failed on worker: %s", 

79 worker, 

80 exc_info=exc, 

81 ) 

82 elif not EVENT_ELASTICSEARCH.is_set(): 

83 try: 

84 await setup_elasticsearch_configs( 

85 es, app.settings["ELASTICSEARCH_PREFIX"] 

86 ) 

87 except Exception: # pylint: disable=broad-except 

88 LOGGER.exception( 

89 "An exception occured while configuring Elasticsearch on worker: %s", # noqa: B950 

90 worker, 

91 ) 

92 else: 

93 EVENT_ELASTICSEARCH.set() 

94 await asyncio.sleep(25) 

95 

96 

97async def check_if_ppid_changed(ppid: int) -> None: 

98 """Check whether Technoblade hates us.""" 

99 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used 

100 if os.getppid() != ppid: 

101 EVENT_SHUTDOWN.set() 

102 return 

103 await asyncio.sleep(1) 

104 

105 

106async def check_redis( 

107 app: Application, worker: int | None 

108) -> None: # pragma: no cover 

109 """Check Redis.""" 

110 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used 

111 redis: Redis[str] = cast("Redis[str]", app.settings.get("REDIS")) 

112 try: 

113 await redis.ping() 

114 except Exception: # pylint: disable=broad-except 

115 EVENT_REDIS.clear() 

116 LOGGER.exception("Connecting to Redis failed on worker %s", worker) 

117 else: 

118 EVENT_REDIS.set() 

119 await asyncio.sleep(20) 

120 

121 

122async def heartbeat() -> None: 

123 """Heartbeat.""" 

124 global HEARTBEAT # pylint: disable=global-statement 

125 while HEARTBEAT: # pylint: disable=while-used 

126 HEARTBEAT = time.monotonic() 

127 await asyncio.sleep(0.05) 

128 

129 

130async def wait_for_shutdown() -> None: # pragma: no cover 

131 """Wait for the shutdown event.""" 

132 loop = asyncio.get_running_loop() 

133 while not EVENT_SHUTDOWN.is_set(): # pylint: disable=while-used 

134 await asyncio.sleep(0.05) 

135 loop.stop() 

136 

137 

138def start_background_tasks( # pylint: disable=too-many-arguments 

139 *, 

140 app: Application, 

141 processes: int, 

142 module_infos: Iterable[ModuleInfo], 

143 loop: asyncio.AbstractEventLoop, 

144 main_pid: int, 

145 elasticsearch_is_enabled: bool, 

146 redis_is_enabled: bool, 

147 worker: int | None, 

148) -> Set[asyncio.Task[None]]: 

149 """Start all required background tasks.""" 

150 

151 async def execute_background_task(task: BackgroundTask, /) -> None: 

152 """Execute a background task with error handling.""" 

153 try: 

154 await task(app=app, worker=worker) 

155 except asyncio.exceptions.CancelledError: 

156 pass 

157 except BaseException as exc: # pylint: disable=broad-exception-caught 

158 LOGGER.exception( 

159 "A %s exception occured while executing background task %s.%s", 

160 exc.__class__.__name__, 

161 task.__module__, 

162 task.__name__, 

163 ) 

164 if not isinstance(exc, Exception): 

165 raise 

166 else: 

167 LOGGER.debug( 

168 "Background task %s.%s finished executing", 

169 task.__module__, 

170 task.__name__, 

171 ) 

172 

173 background_tasks: set[asyncio.Task[None]] = set() 

174 

175 def create_task(fun: BackgroundTask, /) -> asyncio.Task[None]: 

176 """Create an asyncio.Task object from a BackgroundTask.""" 

177 name = f"{fun.__module__}.{fun.__name__}" 

178 if not worker: # log only once 

179 LOGGER.info("starting %s background task", name) 

180 task = loop.create_task(execute_background_task(fun), name=name) 

181 task.add_done_callback(background_tasks.discard) 

182 return task 

183 

184 task_stream: typed_stream.Stream[asyncio.Task[None]] = assert_type( 

185 typed_stream.Stream(module_infos) 

186 .flat_map(lambda info: info.required_background_tasks) 

187 .chain( 

188 typed_stream.Stream((heartbeat, wait_for_shutdown)).map( 

189 lambda fun: wraps(fun)(lambda **_: fun()) 

190 ) 

191 ) 

192 .chain( 

193 [ 

194 wraps(check_if_ppid_changed)( 

195 lambda **k: check_if_ppid_changed(main_pid) 

196 ) 

197 ] 

198 if processes 

199 else () 

200 ) 

201 .chain([check_elasticsearch] if elasticsearch_is_enabled else ()) 

202 .chain([check_redis] if redis_is_enabled else ()) 

203 .distinct() 

204 .map(create_task), 

205 typed_stream.Stream[asyncio.Task[None]], 

206 ) 

207 

208 background_tasks.update(task_stream) 

209 

210 return background_tasks