Repository: apache/airflow
Pull request: Add `lock` to `BaseDagBundle` and use it in `GitDagBundle`
`initialize` and `refresh` need to be safe to call from different threads/processes, so this adds a simple `lock` context manager that bundles can use if the underlying implementation doesn't allow for that.

For example, `GitDagBundle` is the first taker.

closes: #46444

Problem statement:
Add locking for bundle initialize on workers
### Body

Right now it's possible for more than one task to be trying to initialize a bundle at the same time, which can lead to git failing, for example.

We need to use lock files, like we are building for ttl, to also serialize this.

### Committer

- [x] I acknowledge that I am a maintainer/committer of the Apache Airflow project.

Hints from issue discussion:

Commit message: Add `lock` to `BaseDagBundle` and use it in `GitDagBundle`

`initialize` and `refresh` need to be safe to call from different
threads/processes, so this adds a simple `lock` context manager that
bundles can use if the underlying implementation doesn't allow for that.

For example, `GitDagBundle` is the first taker.

Patch to review:
diff --git a/airflow/dag_processing/bundles/base.py b/airflow/dag_processing/bundles/base.py
index 81bf38d373677..16b39123dac95 100644
--- a/airflow/dag_processing/bundles/base.py
+++ b/airflow/dag_processing/bundles/base.py
@@ -17,8 +17,10 @@
 
 from __future__ import annotations
 
+import fcntl
 import tempfile
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from pathlib import Path
 
 from airflow.configuration import conf
@@ -46,6 +48,7 @@ class BaseDagBundle(ABC):
     """
 
     supports_versioning: bool = False
+    _locked: bool = False
 
     def __init__(
         self,
@@ -67,6 +70,10 @@ def initialize(self) -> None:
         and allows for deferring expensive operations until that point in time. This will
         only be called when Airflow needs the bundle files on disk - some uses only need
         to call the `view_url` method, which can run without initializing the bundle.
+
+        This method must be safe to call concurrently from different threads or processes.
+        If the underlying implementation is not safe, the `lock` context manager can be used to
+        ensure that only one thread or process is initializing the bundle at a time.
         """
         self.is_initialized = True
 
@@ -101,7 +108,13 @@ def get_current_version(self) -> str | None:
 
     @abstractmethod
     def refresh(self) -> None:
-        """Retrieve the latest version of the files in the bundle."""
+        """
+        Retrieve the latest version of the files in the bundle.
+
+        This method must be safe to call concurrently from different threads or processes.
+        If the underlying implementation is not safe, the `lock` context manager can be used to
+        ensure that only one thread or process is initializing the bundle at a time.
+        """
 
     def view_url(self, version: str | None = None) -> str | None:
         """
@@ -112,3 +125,22 @@ def view_url(self, version: str | None = None) -> str | None:
         :param version: Version to view
         :return: URL to view the bundle
         """
+
+    @contextmanager
+    def lock(self):
+        if self._locked:
+            yield
+            return
+
+        lock_dir_path = self._dag_bundle_root_storage_path / "locks"
+        lock_dir_path.mkdir(parents=True, exist_ok=True)
+        lock_file_path = lock_dir_path / f"{self.name}.lock"
+        with open(lock_file_path, "w") as lock_file:
+            # Exclusive lock - blocks until it is available
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            try:
+                self._locked = True
+                yield
+            finally:
+                fcntl.flock(lock_file, fcntl.LOCK_UN)
+                self._locked = False
diff --git a/airflow/dag_processing/bundles/git.py b/airflow/dag_processing/bundles/git.py
index 60da6a678ed53..ea6833e66bae1 100644
--- a/airflow/dag_processing/bundles/git.py
+++ b/airflow/dag_processing/bundles/git.py
@@ -154,19 +154,20 @@ def __init__(
             self.log.warning("Could not create GitHook for connection %s : %s", self.git_conn_id, e)
 
     def _initialize(self):
-        with self.hook.configure_hook_env():
-            self._clone_bare_repo_if_required()
-            self._ensure_version_in_bare_repo()
-
-        self._clone_repo_if_required()
-        self.repo.git.checkout(self.tracking_ref)
-        if self.version:
-            if not self._has_version(self.repo, self.version):
-                self.repo.remotes.origin.fetch()
-            self.repo.head.set_reference(self.repo.commit(self.version))
-            self.repo.head.reset(index=True, working_tree=True)
-        else:
-            self.refresh()
+        with self.lock():
+            with self.hook.configure_hook_env():
+                self._clone_bare_repo_if_required()
+                self._ensure_version_in_bare_repo()
+
+            self._clone_repo_if_required()
+            self.repo.git.checkout(self.tracking_ref)
+            if self.version:
+                if not self._has_version(self.repo, self.version):
+                    self.repo.remotes.origin.fetch()
+                self.repo.head.set_reference(self.repo.commit(self.version))
+                self.repo.head.reset(index=True, working_tree=True)
+            else:
+                self.refresh()
 
     def initialize(self) -> None:
         if not self.repo_url:
@@ -259,9 +260,10 @@ def refresh(self) -> None:
         if self.version:
             raise AirflowException("Refreshing a specific version is not supported")
 
-        with self.hook.configure_hook_env():
-            self._fetch_bare_repo()
-            self.repo.remotes.origin.pull()
+        with self.lock():
+            with self.hook.configure_hook_env():
+                self._fetch_bare_repo()
+                self.repo.remotes.origin.pull()
 
     @staticmethod
     def _convert_git_ssh_url_to_https(url: str) -> str:
diff --git a/tests/dag_processing/test_dag_bundles.py b/tests/dag_processing/test_dag_bundles.py
index a9a48139ba9eb..087c0a5206667 100644
--- a/tests/dag_processing/test_dag_bundles.py
+++ b/tests/dag_processing/test_dag_bundles.py
@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import fcntl
 import os
 import re
 import tempfile
@@ -52,22 +53,82 @@ def test_default_dag_storage_path():
         assert bundle._dag_bundle_root_storage_path == Path(tempfile.gettempdir(), "airflow", "dag_bundles")
 
 
-def test_dag_bundle_root_storage_path():
-    class BasicBundle(BaseDagBundle):
-        def refresh(self):
-            pass
+class BasicBundle(BaseDagBundle):
+    def refresh(self):
+        pass
+
+    def get_current_version(self):
+        pass
 
-        def get_current_version(self):
-            pass
+    def path(self):
+        pass
 
-        def path(self):
-            pass
 
+def test_dag_bundle_root_storage_path():
     with conf_vars({("dag_processor", "dag_bundle_storage_path"): None}):
         bundle = BasicBundle(name="test")
         assert bundle._dag_bundle_root_storage_path == Path(tempfile.gettempdir(), "airflow", "dag_bundles")
 
 
+def test_lock_acquisition():
+    """Test that the lock context manager sets _locked and locks a lock file."""
+    bundle = BasicBundle(name="locktest")
+    lock_dir = bundle._dag_bundle_root_storage_path / "locks"
+    lock_file = lock_dir / f"{bundle.name}.lock"
+
+    assert not bundle._locked
+
+    with bundle.lock():
+        assert bundle._locked
+        assert lock_file.exists()
+
+        # Check lock file is now locked
+        with open(lock_file, "w") as f:
+            try:
+                # Try to acquire an exclusive lock in non-blocking mode.
+                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                locked = False
+            except OSError:
+                locked = True
+            assert locked
+
+    # After, _locked is False and file unlock has been called.
+    assert bundle._locked is False
+    with open(lock_file, "w") as f:
+        try:
+            fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            unlocked = True
+            fcntl.flock(f, fcntl.LOCK_UN)  # Release the lock immediately.
+        except OSError:
+            unlocked = False
+        assert unlocked
+
+
+def test_lock_exception_handling():
+    """Test that exceptions within the lock context manager still release the lock."""
+    bundle = BasicBundle(name="locktest")
+    lock_dir = bundle._dag_bundle_root_storage_path / "locks"
+    lock_file = lock_dir / f"{bundle.name}.lock"
+
+    try:
+        with bundle.lock():
+            assert bundle._locked
+            raise Exception("...")
+    except Exception:
+        pass
+
+    # lock file should be unlocked
+    assert not bundle._locked
+    with open(lock_file, "w") as f:
+        try:
+            fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            acquired = True
+            fcntl.flock(f, fcntl.LOCK_UN)
+        except OSError:
+            acquired = False
+        assert acquired
+
+
 class TestLocalDagBundle:
     def test_path(self):
         bundle = LocalDagBundle(name="test", path="/hello")
@@ -581,3 +642,12 @@ def test_repo_url_access_missing_connection_doesnt_error(self, mock_log):
         )
         assert bundle.repo_url == "some_repo_url"
         assert "Could not create GitHook for connection" in mock_log.warning.call_args[0][0]
+
+    @mock.patch("airflow.dag_processing.bundles.git.GitHook")
+    def test_lock_used(self, mock_githook, git_repo):
+        repo_path, repo = git_repo
+        mock_githook.return_value.repo_url = repo_path
+        bundle = GitDagBundle(name="test", tracking_ref=GIT_DEFAULT_BRANCH)
+        with mock.patch("airflow.dag_processing.bundles.git.GitDagBundle.lock") as mock_lock:
+            bundle.initialize()
+            assert mock_lock.call_count == 2  # both initialize and refresh


Code review comments:
diff_hunk: @@ -112,3 +125,22 @@ def view_url(self, version: str | None = None) -> str | None:
         :param version: Version to view
         :return: URL to view the bundle
         """
+
+    @contextmanager
+    def lock(self):
+        if self._locked:
+            yield
+            return
+
+        lock_dir_path = self._dag_bundle_root_storage_path / "locks"
path: airflow/dag_processing/bundles/base.py
text: @user1:
maybe `_locks`? to avoid having to "reserve" the word?

@author:
Yeah, fair! I'll do it after this test run.

diff_hunk: @@ -112,3 +125,22 @@ def view_url(self, version: str | None = None) -> str | None:
         :param version: Version to view
         :return: URL to view the bundle
         """
+
+    @contextmanager
+    def lock(self):
+        if self._locked:
path: airflow/dag_processing/bundles/base.py
text: @user1:
is there a need to tolerate nested locking? should we maybe raise instead?

@author:
For now, yes. And probably should, yes. Gits initialize calls refresh though. And I'm trying to keep it simple. Happy to refactor in a followup.