This is an automated email from the ASF dual-hosted git repository. dstandish pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new 095c5fe313 Ensure __exit__ is called in decorator context managers (#38383) 095c5fe313 is described below commit 095c5fe3137e2cb6d45e8f3184bae149cb2850d1 Author: Daniel Standish <15932138+dstand...@users.noreply.github.com> AuthorDate: Thu Mar 21 14:06:49 2024 -0700 Ensure __exit__ is called in decorator context managers (#38383) In #36800 author fixed zombie scheduler issue arising from context manager exit not being called, thus sub process not getting terminated. It was fixed by explicitly calling the `close` function on an ExitStack-managed context manager. Simpler / better / cleaner / more standard solution is to "fix" the underlying context managers by wrapping the yield in a try / finally. --- airflow/cli/commands/celery_command.py | 8 ++++--- airflow/cli/commands/scheduler_command.py | 32 +++++++++++--------------- airflow/providers/celery/cli/celery_command.py | 8 ++++--- tests/cli/commands/test_scheduler_command.py | 3 --- 4 files changed, 23 insertions(+), 28 deletions(-) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index ae8b9c1925..ed6c0dbbd8 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -91,9 +91,11 @@ def _serve_logs(skip_serve_logs: bool = False): if skip_serve_logs is False: sub_proc = Process(target=serve_logs) sub_proc.start() - yield - if sub_proc: - sub_proc.terminate() + try: + yield + finally: + if sub_proc: + sub_proc.terminate() @after_setup_logger.connect() diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index 0b5cac8857..2a55ca2373 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -20,7 +20,7 @@ from __future__ import annotations import logging from argparse import Namespace -from contextlib import ExitStack, contextmanager +from contextlib import contextmanager from multiprocessing import Process from airflow import settings @@ -45,18 +45,8 @@ def _run_scheduler_job(args) -> None: ExecutorLoader.validate_database_executor_compatibility(job_runner.job.executor) InternalApiConfig.force_database_direct_access() enable_health_check = conf.getboolean("scheduler", "ENABLE_HEALTH_CHECK") - with ExitStack() as stack: - stack.enter_context(_serve_logs(args.skip_serve_logs)) - stack.enter_context(_serve_health_check(enable_health_check)) - - try: - run_job(job=job_runner.job, execute_callable=job_runner._execute) - except Exception: - log.exception("Exception when running scheduler job") - raise - finally: - # Ensure that the contexts are closed - stack.close() + with _serve_logs(args.skip_serve_logs), _serve_health_check(enable_health_check): + run_job(job=job_runner.job, execute_callable=job_runner._execute) @cli_utils.action_cli @@ -84,9 +74,11 @@ def _serve_logs(skip_serve_logs: bool = False): if skip_serve_logs is False: sub_proc = Process(target=serve_logs) sub_proc.start() - yield - if sub_proc: - sub_proc.terminate() + try: + yield + finally: + if sub_proc: + sub_proc.terminate() @contextmanager @@ -96,6 +88,8 @@ def _serve_health_check(enable_health_check: bool = False): if enable_health_check: sub_proc = Process(target=serve_health_check) sub_proc.start() - yield - if sub_proc: - sub_proc.terminate() + try: + yield + finally: + if sub_proc: + sub_proc.terminate() diff --git a/airflow/providers/celery/cli/celery_command.py b/airflow/providers/celery/cli/celery_command.py index fff46090aa..f7682b9abf 100644 --- a/airflow/providers/celery/cli/celery_command.py +++ b/airflow/providers/celery/cli/celery_command.py @@ -107,9 +107,11 @@ def _serve_logs(skip_serve_logs: bool = False): if skip_serve_logs is False: sub_proc = Process(target=serve_logs) sub_proc.start() - yield - if sub_proc: - sub_proc.terminate() + try: + yield + finally: + if sub_proc: + sub_proc.terminate() @after_setup_logger.connect() diff --git a/tests/cli/commands/test_scheduler_command.py b/tests/cli/commands/test_scheduler_command.py index 2853763563..b6d6a9d921 100644 --- a/tests/cli/commands/test_scheduler_command.py +++ b/tests/cli/commands/test_scheduler_command.py @@ -165,10 +165,8 @@ class TestSchedulerCommand: @mock.patch("airflow.cli.commands.scheduler_command.SchedulerJobRunner") @mock.patch("airflow.cli.commands.scheduler_command.Process") @mock.patch("airflow.cli.commands.scheduler_command.run_job", side_effect=Exception("run_job failed")) - @mock.patch("airflow.cli.commands.scheduler_command.log") def test_run_job_exception_handling( self, - mock_log, mock_run_job, mock_process, mock_scheduler_job, @@ -183,7 +181,6 @@ class TestSchedulerCommand: job=mock_scheduler_job().job, execute_callable=mock_scheduler_job()._execute, ) - mock_log.exception.assert_called_once_with("Exception when running scheduler job") mock_process.assert_called_once_with(target=serve_logs) mock_process().terminate.assert_called_once_with()