|
19 | 19 | from __future__ import annotations |
20 | 20 |
|
21 | 21 | import json |
| 22 | +import os |
| 23 | +import signal |
| 24 | +from unittest import mock |
22 | 25 |
|
23 | 26 | import pytest |
24 | 27 |
|
|
28 | 31 | BaseCoordinator, |
29 | 32 | CoordinatorManager, |
30 | 33 | _PythonCoordinator, |
| 34 | + _warm_shutdown_signals, |
31 | 35 | get_coordinator_manager, |
32 | 36 | reset_coordinator_manager, |
33 | 37 | ) |
@@ -229,3 +233,136 @@ def test_every_example_coordinator_constructs(self, sdk_config): |
229 | 233 | for queue, key in queue_to_coordinator.items(): |
230 | 234 | coordinator = manager.for_queue(queue) |
231 | 235 | assert isinstance(coordinator, import_string(specs[key]["classpath"])) |
| 236 | + |
| 237 | + |
| 238 | +class TestWarmShutdownSignals: |
| 239 | + """Tests for the warm-shutdown signal handling that wraps task supervision.""" |
| 240 | + |
| 241 | + @pytest.fixture(autouse=True) |
| 242 | + def _restore_disposition(self): |
| 243 | + """Guarantee SIGTERM/SIGINT dispositions are restored even if a test leaks one.""" |
| 244 | + original_term = signal.getsignal(signal.SIGTERM) |
| 245 | + original_int = signal.getsignal(signal.SIGINT) |
| 246 | + yield |
| 247 | + signal.signal(signal.SIGTERM, original_term) |
| 248 | + signal.signal(signal.SIGINT, original_int) |
| 249 | + |
| 250 | + def test_installs_handlers_inside_context(self): |
| 251 | + """While the context is active a warm-shutdown handler is installed for both signals.""" |
| 252 | + sentinel_term = signal.getsignal(signal.SIGTERM) |
| 253 | + sentinel_int = signal.getsignal(signal.SIGINT) |
| 254 | + |
| 255 | + with _warm_shutdown_signals(): |
| 256 | + inside_term = signal.getsignal(signal.SIGTERM) |
| 257 | + inside_int = signal.getsignal(signal.SIGINT) |
| 258 | + |
| 259 | + assert callable(inside_term) |
| 260 | + assert callable(inside_int) |
| 261 | + # The installed handler is the warm-shutdown closure, not the previous disposition. |
| 262 | + assert inside_term is not sentinel_term |
| 263 | + assert inside_int is not sentinel_int |
| 264 | + # Both signals share the same warm-shutdown closure. |
| 265 | + assert inside_term is inside_int |
| 266 | + |
| 267 | + def test_restores_previous_dispositions_on_exit(self): |
| 268 | + """The exact previous dispositions are restored when the context exits normally.""" |
| 269 | + |
| 270 | + def _prev_term(signum, frame): # pragma: no cover - never invoked |
| 271 | + pass |
| 272 | + |
| 273 | + def _prev_int(signum, frame): # pragma: no cover - never invoked |
| 274 | + pass |
| 275 | + |
| 276 | + signal.signal(signal.SIGTERM, _prev_term) |
| 277 | + signal.signal(signal.SIGINT, _prev_int) |
| 278 | + |
| 279 | + with _warm_shutdown_signals(): |
| 280 | + pass |
| 281 | + |
| 282 | + assert signal.getsignal(signal.SIGTERM) is _prev_term |
| 283 | + assert signal.getsignal(signal.SIGINT) is _prev_int |
| 284 | + |
| 285 | + def test_restores_previous_dispositions_on_exception(self): |
| 286 | + """Dispositions are restored even if the wrapped body raises.""" |
| 287 | + |
| 288 | + def _prev_term(signum, frame): # pragma: no cover - never invoked |
| 289 | + pass |
| 290 | + |
| 291 | + signal.signal(signal.SIGTERM, _prev_term) |
| 292 | + |
| 293 | + with pytest.raises(RuntimeError, match="boom"), _warm_shutdown_signals(): |
| 294 | + raise RuntimeError("boom") |
| 295 | + |
| 296 | + assert signal.getsignal(signal.SIGTERM) is _prev_term |
| 297 | + |
| 298 | + def test_sigterm_inside_context_does_not_kill(self): |
| 299 | + """ |
| 300 | + A SIGTERM delivered while supervising must be swallowed, not kill the process. |
| 301 | +
|
| 302 | + This is the regression guard: with the default SIGTERM disposition (SIG_DFL) |
| 303 | + in place as the *previous* handler, sending SIGTERM to ourselves would |
| 304 | + terminate the process. The warm-shutdown handler installed by the context |
| 305 | + manager must absorb it so the running task is allowed to finish. |
| 306 | + """ |
| 307 | + # Make the pre-context disposition the default so a missing warm-shutdown |
| 308 | + # handler would actually kill this process (and fail the test by dying). |
| 309 | + signal.signal(signal.SIGTERM, signal.SIG_DFL) |
| 310 | + |
| 311 | + reached_after_signal = False |
| 312 | + with _warm_shutdown_signals(): |
| 313 | + os.kill(os.getpid(), signal.SIGTERM) |
| 314 | + # If the handler did not absorb the signal, we never get here. |
| 315 | + reached_after_signal = True |
| 316 | + |
| 317 | + assert reached_after_signal |
| 318 | + # And the default disposition is put back afterwards. |
| 319 | + assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL |
| 320 | + |
| 321 | + |
| 322 | +class TestPythonCoordinatorWarmShutdown: |
| 323 | + """The Python coordinator must wrap start() and wait() in the warm-shutdown handlers.""" |
| 324 | + |
| 325 | + def test_execute_task_wraps_start_and_wait(self, monkeypatch): |
| 326 | + """ |
| 327 | + Handlers are installed for the whole start()+wait() window and restored after. |
| 328 | +
|
| 329 | + Capturing the SIGTERM disposition at the moment ``start()`` and ``wait()`` |
| 330 | + run proves the handler spans the RUNNING transition (start) and the |
| 331 | + terminal-state report / log upload (wait), with no window left uncovered. |
| 332 | + """ |
| 333 | + original_term = signal.getsignal(signal.SIGTERM) |
| 334 | + captured: dict[str, object] = {} |
| 335 | + |
| 336 | + class _FakeProcess: |
| 337 | + final_state = "success" |
| 338 | + |
| 339 | + def wait(self_inner): |
| 340 | + captured["wait"] = signal.getsignal(signal.SIGTERM) |
| 341 | + return 0 |
| 342 | + |
| 343 | + def _fake_start(*args, **kwargs): |
| 344 | + captured["start"] = signal.getsignal(signal.SIGTERM) |
| 345 | + return _FakeProcess() |
| 346 | + |
| 347 | + import airflow.sdk.execution_time.supervisor as supervisor_mod |
| 348 | + |
| 349 | + monkeypatch.setattr(supervisor_mod.ActivitySubprocess, "start", staticmethod(_fake_start)) |
| 350 | + |
| 351 | + coordinator = _PythonCoordinator() |
| 352 | + result = coordinator.execute_task( |
| 353 | + what=mock.MagicMock(), |
| 354 | + dag_rel_path="some_dag.py", |
| 355 | + bundle_info=mock.MagicMock(), |
| 356 | + client=mock.MagicMock(), |
| 357 | + subprocess_logs_to_stdout=False, |
| 358 | + ) |
| 359 | + |
| 360 | + assert result.exit_code == 0 |
| 361 | + assert result.final_state == "success" |
| 362 | + # During both start() and wait() a warm-shutdown handler was installed... |
| 363 | + assert callable(captured["start"]) |
| 364 | + assert callable(captured["wait"]) |
| 365 | + assert captured["start"] is not original_term |
| 366 | + assert captured["start"] is captured["wait"] |
| 367 | + # ...and the original disposition is restored once execute_task returns. |
| 368 | + assert signal.getsignal(signal.SIGTERM) is original_term |
0 commit comments