Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/gunicorn_app.py: 0%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17""" 

18Gunicorn application with integrated worker monitoring. 

19 

20This module provides a custom Gunicorn application that integrates worker 

21monitoring directly into the arbiter process loop. By subclassing the ``Arbiter``, 

22we can perform rolling worker restarts without needing a separate monitoring 

23thread or subprocess. 

24 

25The pattern follows gunicorn's recommended extension approach: 

26- Subclass ``BaseApplication`` to configure gunicorn programmatically 

27- Override ``run()`` to use a custom ``Arbiter`` 

28- Custom ``Arbiter`` hooks into manage_workers() for monitoring logic 

29""" 

30 

31from __future__ import annotations 

32 

33import os 

34import signal 

35import sys 

36import time 

37import warnings 

38from typing import TYPE_CHECKING, Any 

39 

40import structlog 

41from gunicorn.app.base import BaseApplication 

42from gunicorn.arbiter import Arbiter 

43from gunicorn.glogging import Logger as GunicornLogger 

44from uvicorn.workers import UvicornWorker 

45 

46from airflow.configuration import conf 

47from airflow.exceptions import RemovedInAirflow4Warning 

48 

49if TYPE_CHECKING: 

50 from ssl import VerifyMode 

51 

52 from fastapi import FastAPI 

53 from gunicorn.app.base import Application 

54 

55log = structlog.get_logger(__name__) 

56 

57 

58class AirflowGunicornLogger(GunicornLogger): 

59 """ 

60 Gunicorn logger that routes all output through Airflow's logging setup. 

61 

62 Gunicorn's default Logger.setup() installs its own StreamHandler with a custom 

63 formatter on ``gunicorn.error`` and ``gunicorn.access``, bypassing any root-level 

64 handler (including our structlog ProcessorFormatter). Overriding setup() to do 

65 nothing lets records propagate to root where Airflow's handler picks them up. 

66 """ 

67 

68 def setup(self, cfg) -> None: 

69 self.error_log.propagate = True 

70 self.access_log.propagate = True 

71 

72 

73class AirflowUvicornWorker(UvicornWorker): 

74 """ 

75 Uvicorn worker that preserves Airflow's structlog-based logging setup. 

76 

77 Uvicorn workers normally call ``logging.config.dictConfig(LOGGING_CONFIG)`` on startup 

78 which would override any structlog configuration applied before gunicorn starts. 

79 Setting ``log_config=None`` prevents that. ``access_log=False`` disables uvicorn's 

80 built-in access logger because ``HttpAccessLogMiddleware`` handles access logging. 

81 """ 

82 

83 CONFIG_KWARGS = { 

84 **UvicornWorker.CONFIG_KWARGS, 

85 "log_config": None, 

86 "access_log": False, 

87 } 

88 

89 

90class AirflowArbiter(Arbiter): 

91 """Custom ``Arbiter`` with rolling worker restarts via manage_workers() hook.""" 

92 

93 def __init__(self, app: Application): 

94 super().__init__(app) 

95 

96 # Worker refresh configuration 

97 self.worker_refresh_interval = conf.getint("api", "worker_refresh_interval", fallback=0) 

98 self.worker_refresh_batch_size = conf.getint("api", "worker_refresh_batch_size", fallback=1) 

99 

100 # State tracking for rolling restarts 

101 self._last_refresh_time = time.monotonic() 

102 self._refresh_in_progress = False 

103 self._workers_to_replace: set[int] = set() 

104 

105 # Validate configuration 

106 if self.worker_refresh_batch_size > self.num_workers: 

107 log.warning( 

108 "worker_refresh_batch_size (%d) > num_workers (%d), reducing batch size", 

109 self.worker_refresh_batch_size, 

110 self.num_workers, 

111 ) 

112 self.worker_refresh_batch_size = self.num_workers 

113 

114 if self.worker_refresh_interval > 0: 

115 log.info( 

116 "Worker refresh enabled: interval=%ds, batch_size=%d", 

117 self.worker_refresh_interval, 

118 self.worker_refresh_batch_size, 

119 ) 

120 

121 def manage_workers(self) -> None: 

122 """Maintain worker count and perform rolling restarts if configured.""" 

123 super().manage_workers() 

124 

125 # Check if worker refresh is enabled and due 

126 if self.worker_refresh_interval > 0: 

127 self._check_worker_refresh() 

128 

129 def _check_worker_refresh(self) -> None: 

130 """Check if it's time to start or continue a worker refresh cycle.""" 

131 elapsed = time.monotonic() - self._last_refresh_time 

132 

133 if not self._refresh_in_progress: 

134 # Check if it's time to start a new refresh cycle 

135 if elapsed >= self.worker_refresh_interval: 

136 self._start_refresh_cycle() 

137 else: 

138 # Continue ongoing refresh cycle 

139 self._continue_refresh_cycle() 

140 

141 def _start_refresh_cycle(self) -> None: 

142 """Start a new rolling worker refresh cycle.""" 

143 if not self.WORKERS: 

144 return 

145 

146 self._refresh_in_progress = True 

147 self._workers_to_replace = set(self.WORKERS.keys()) 

148 log.info( 

149 "Starting worker refresh cycle: %d workers to replace", 

150 len(self._workers_to_replace), 

151 ) 

152 self._continue_refresh_cycle() 

153 

154 def _continue_refresh_cycle(self) -> None: 

155 """Continue rolling refresh: spawn new workers, kill old ones in batches.""" 

156 # Remove workers that have already exited 

157 current_pids = set(self.WORKERS.keys()) 

158 self._workers_to_replace &= current_pids 

159 

160 if not self._workers_to_replace: 

161 # All original workers have been replaced 

162 log.info("Worker refresh cycle completed") 

163 self._refresh_in_progress = False 

164 self._last_refresh_time = time.monotonic() 

165 return 

166 

167 # Check if we have capacity to spawn new workers 

168 # We temporarily exceed num_workers during rolling restart 

169 current_count = len(self.WORKERS) 

170 max_during_refresh = self.num_workers + self.worker_refresh_batch_size 

171 

172 if current_count < max_during_refresh: 

173 # Spawn new workers up to batch size 

174 workers_to_spawn = min( 

175 self.worker_refresh_batch_size, 

176 max_during_refresh - current_count, 

177 ) 

178 if workers_to_spawn > 0: 

179 log.debug("Spawning %d new worker(s) for refresh", workers_to_spawn) 

180 for _ in range(workers_to_spawn): 

181 self.spawn_worker() 

182 

183 # If we have more workers than target, kill old ones 

184 if current_count > self.num_workers: 

185 workers_to_kill = min( 

186 current_count - self.num_workers, 

187 self.worker_refresh_batch_size, 

188 len(self._workers_to_replace), 

189 ) 

190 

191 # Kill oldest workers first (FIFO) 

192 sorted_workers = sorted( 

193 [(pid, w) for pid, w in self.WORKERS.items() if pid in self._workers_to_replace], 

194 key=lambda x: x[1].age, 

195 ) 

196 

197 for pid, worker in sorted_workers[:workers_to_kill]: 

198 log.info("Killing old worker %d (age: %s) for refresh", pid, worker.age) 

199 self.kill_worker(pid, signal.SIGTERM) 

200 self._workers_to_replace.discard(pid) 

201 

202 

203class AirflowGunicornApp(BaseApplication): 

204 """Gunicorn application that uses AirflowArbiter for worker management.""" 

205 

206 def __init__(self, options: dict[str, Any] | None = None): 

207 self.options = options or {} 

208 self.application: FastAPI | None = None 

209 super().__init__() 

210 

211 def load_config(self) -> None: 

212 """Load configuration from options dict, then GUNICORN_CMD_ARGS env var.""" 

213 for key, value in self.options.items(): 

214 if key in self.cfg.settings and value is not None: 

215 self.cfg.set(key.lower(), value) 

216 

217 cmd_args = self.cfg.get_cmd_args_from_env() 

218 if cmd_args: 

219 log.info("Applying GUNICORN_CMD_ARGS: %s", cmd_args) 

220 parser = self.cfg.parser() 

221 env_args = parser.parse_args(cmd_args) 

222 for k, v in vars(env_args).items(): 

223 if v is None or k == "args": 

224 continue 

225 self.cfg.set(k.lower(), v) 

226 

227 def load(self) -> Any: 

228 """Load and return the WSGI/ASGI application.""" 

229 if self.application is None: 

230 from airflow.api_fastapi.main import app 

231 

232 self.application = app 

233 return self.application 

234 

235 def run(self) -> None: 

236 """Run the application with AirflowArbiter.""" 

237 try: 

238 AirflowArbiter(self).run() 

239 except RuntimeError as e: 

240 log.error("Gunicorn failed to start", error=str(e)) 

241 sys.exit(1) 

242 

243 

244def create_gunicorn_app( 

245 host: str, 

246 port: int, 

247 num_workers: int, 

248 worker_timeout: int, 

249 ssl_cert: str | None = None, 

250 ssl_key: str | None = None, 

251 ssl_ca_file: str | None = None, 

252 ssl_cert_reqs: VerifyMode | None = None, 

253 log_level: str = "info", 

254 proxy_headers: bool = False, 

255) -> AirflowGunicornApp: 

256 """ 

257 Create a configured AirflowGunicornApp instance. 

258 

259 :param host: Host to bind to 

260 :param port: Port to bind to 

261 :param num_workers: Number of worker processes 

262 :param worker_timeout: Worker timeout in seconds 

263 :param ssl_cert: Path to SSL certificate file 

264 :param ssl_key: Path to SSL key file 

265 :param ssl_ca_file: Path to the SSL CA certs file 

266 :param ssl_cert_reqs: SSL client certificate requirements 

267 :param log_level: Log level (debug, info, warning, error, critical) 

268 :param proxy_headers: Whether to trust proxy headers 

269 """ 

270 options = { 

271 "bind": f"{host}:{port}", 

272 "workers": num_workers, 

273 "worker_class": "airflow.api_fastapi.gunicorn_app.AirflowUvicornWorker", 

274 "timeout": worker_timeout, 

275 "graceful_timeout": worker_timeout, 

276 "keepalive": worker_timeout, 

277 "loglevel": log_level, 

278 "logger_class": "airflow.api_fastapi.gunicorn_app.AirflowGunicornLogger", 

279 "preload_app": True, 

280 # Use our gunicorn_config module for hooks (post_worker_init, worker_exit) 

281 "config": "python:airflow.api_fastapi.gunicorn_config", 

282 } 

283 

284 if ssl_cert and ssl_key: 

285 options["certfile"] = ssl_cert 

286 options["keyfile"] = ssl_key 

287 if ssl_ca_file: 

288 options["ca_certs"] = ssl_ca_file 

289 if ssl_cert_reqs is not None: 

290 options["cert_reqs"] = ssl_cert_reqs 

291 

292 if not proxy_headers: 

293 # ``UvicornWorker`` leaves uvicorn's ``proxy_headers`` at its default of True, so 

294 # trusting nobody is the only way to keep the worker off X-Forwarded-*. 

295 options["forwarded_allow_ips"] = "" 

296 elif "FORWARDED_ALLOW_IPS" not in os.environ: 

297 # Deployments that never set the variable are still relying on this server type 

298 # trusting every client, so keep that until the next major version. 

299 warnings.warn( 

300 "Under gunicorn the API server trusts X-Forwarded-* from all clients by default. " 

301 "Set FORWARDED_ALLOW_IPS to restrict trusted proxies; this default will become " 

302 "restrictive in Airflow 4.", 

303 category=RemovedInAirflow4Warning, 

304 stacklevel=2, 

305 ) 

306 options["forwarded_allow_ips"] = "*" 

307 

308 return AirflowGunicornApp(options)