feat(fetch-content): merge concurrently extracted fetches in a fixed order

src/taskgraph/run-task/fetch-content

@@ -683,6 +683,36 @@ def repack_archive(
             )
 
 
+def merge_tree(src, dest):
+    """Move the contents of the src directory into dest and remove src.
+
+    Directories present on both sides are merged entry by entry and take
+    the mode from src. Any other path present on both sides is replaced by
+    the one from src.
+    """
+    dest.mkdir(parents=True, exist_ok=True)
+    os.chmod(src, stat.S_IMODE(os.stat(src).st_mode) | stat.S_IWUSR)
+    for entry in os.scandir(src):
+        target = dest / entry.name
+        if (
+            entry.is_dir(follow_symlinks=False)
+            and target.is_dir()
+            and not target.is_symlink()
+        ):
+            mode = stat.S_IMODE(entry.stat(follow_symlinks=False).st_mode)
+            merge_tree(pathlib.Path(entry.path), target)
+            os.chmod(target, mode)
+            continue
+        if target.is_symlink() or target.exists():
+            log(f"{target} is provided by more than one fetch, replacing it")
+            if target.is_dir() and not target.is_symlink():
+                shutil.rmtree(target)
+            else:
+                target.unlink()
+        os.replace(entry.path, target)
+    src.rmdir()
+
+
 def fetch_and_extract(url, dest_dir, extract=True, sha256=None, size=None):
     """Fetch a URL and extract it to a destination path.
 
@@ -708,15 +738,25 @@ def fetch_and_extract(url, dest_dir, extract=True, sha256=None, size=None):
 
 
 def fetch_urls(downloads):
-    """Fetch URLs pairs to a pathlib.Path."""
-    with concurrent.futures.ThreadPoolExecutor(CONCURRENCY) as e:
-        fs = []
-
-        for download in downloads:
-            fs.append(e.submit(fetch_and_extract, *download))
+    """Fetch URLs pairs to a pathlib.Path.
 
-        for f in fs:
-            f.result()
+    Each fetch is downloaded and extracted into a staging directory under
+    its destination, all in parallel. The staging directories are then
+    merged into the destinations sequentially, in the order the fetches
+    were given. Two archives never write to the same directory at the
+    same time, and when they contain the same path the later fetch wins.
+    """
+    with concurrent.futures.ThreadPoolExecutor(CONCURRENCY) as e:
+        extractions = []
+        for url, dest_dir, *rest in downloads:
+            staging_dir = pathlib.Path(tempfile.mkdtemp(prefix=".fetch.", dir=dest_dir))
+            future = e.submit(fetch_and_extract, url, staging_dir, *rest)
+            extractions.append((future, staging_dir, dest_dir))
+
+        for future, staging_dir, dest_dir in extractions:
+            future.result()
+            log(f"Merging {staging_dir} into {dest_dir}")
+            merge_tree(staging_dir, dest_dir)
 
 
 def _git_checkout_github_archive(

test/test_scripts_fetch_content.py

@@ -1,6 +1,11 @@
+import io
 import json
 import os
 import pathlib
+import shutil
+import stat
+import sys
+import tarfile
 import urllib.request
 from importlib.machinery import SourceFileLoader
 from importlib.util import module_from_spec, spec_from_loader
@@ -259,3 +264,118 @@ def test_should_repack_archive(
     ), (
         f"Failed for orig: {orig}, dest: {dest}, strip_components: {strip_components}, add_prefix: {add_prefix}, expected {expected} but received {not expected}"
     )
+
+
+def _make_tar(path, files):
+    with tarfile.open(path, "w") as tar:
+        for name, content in files.items():
+            data = content.encode("utf-8")
+            info = tarfile.TarInfo(name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+
+def test_fetch_urls_merges_staged_extractions(monkeypatch, tmp_path, fetch_content_mod):
+    archives = tmp_path / "archives"
+    archives.mkdir()
+    _make_tar(
+        archives / "common.tar",
+        {"tests/common/a.txt": "a", "tests/shared.txt": "first"},
+    )
+    _make_tar(
+        archives / "suite.tar",
+        {"tests/suite/b.txt": "b", "tests/shared.txt": "second"},
+    )
+    dest = tmp_path / "fetches"
+    dest.mkdir()
+
+    def mock_download_to_path(url, path, sha256=None, size=None):
+        shutil.copy(archives / path.name, path)
+
+    monkeypatch.setattr(fetch_content_mod, "download_to_path", mock_download_to_path)
+
+    fetch_content_mod.fetch_urls(
+        [
+            ("https://example.com/common.tar", dest, True, None),
+            ("https://example.com/suite.tar", dest, True, None),
+        ]
+    )
+
+    assert sorted(p.relative_to(dest).as_posix() for p in dest.rglob("*")) == [
+        "tests",
+        "tests/common",
+        "tests/common/a.txt",
+        "tests/shared.txt",
+        "tests/suite",
+        "tests/suite/b.txt",
+    ]
+    assert (dest / "tests" / "shared.txt").read_text() == "second"
+
+
+def test_fetch_urls_places_unextracted_files(monkeypatch, tmp_path, fetch_content_mod):
+    archives = tmp_path / "archives"
+    archives.mkdir()
+    _make_tar(archives / "tool.tar", {"tool/bin/tool": "t"})
+    dest = tmp_path / "fetches"
+    dest.mkdir()
+
+    def mock_download_to_path(url, path, sha256=None, size=None):
+        if path.name.endswith(".tar"):
+            shutil.copy(archives / path.name, path)
+        else:
+            path.write_text("plain")
+
+    monkeypatch.setattr(fetch_content_mod, "download_to_path", mock_download_to_path)
+
+    fetch_content_mod.fetch_urls(
+        [
+            ("https://example.com/tool.tar", dest, True, None),
+            ("https://example.com/plain.txt", dest, False, None),
+            ("https://example.com/notatar.txt", dest, True, None),
+        ]
+    )
+
+    assert sorted(p.name for p in dest.iterdir()) == [
+        "notatar.txt",
+        "plain.txt",
+        "tool",
+    ]
+    assert (dest / "tool" / "bin" / "tool").read_text() == "t"
+    assert (dest / "notatar.txt").read_text() == "plain"
+
+
+def test_merge_tree_replaces_conflicting_entries(tmp_path, fetch_content_mod):
+    src = tmp_path / "src"
+    dest = tmp_path / "dest"
+    (src / "dir").mkdir(parents=True)
+    (src / "dir" / "new.txt").write_text("new")
+    (src / "file").write_text("file")
+    (dest / "dir" / "kept").mkdir(parents=True)
+    (dest / "dir" / "new.txt").write_text("old")
+    (dest / "file").mkdir()
+
+    fetch_content_mod.merge_tree(src, dest)
+
+    assert not src.exists()
+    assert (dest / "dir" / "new.txt").read_text() == "new"
+    assert (dest / "dir" / "kept").is_dir()
+    assert (dest / "file").is_file()
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32" or os.getuid() == 0, reason="needs POSIX directory modes"
+)
+def test_merge_tree_readonly_dir_from_later_fetch(tmp_path, fetch_content_mod):
+    src = tmp_path / "src"
+    dest = tmp_path / "dest"
+    (src / "tests").mkdir(parents=True)
+    (src / "tests" / "b.txt").write_text("b")
+    (dest / "tests").mkdir(parents=True)
+    (dest / "tests" / "a.txt").write_text("a")
+    (src / "tests").chmod(0o555)
+
+    fetch_content_mod.merge_tree(src, dest)
+
+    assert (dest / "tests" / "a.txt").read_text() == "a"
+    assert (dest / "tests" / "b.txt").read_text() == "b"
+    assert stat.S_IMODE((dest / "tests").stat().st_mode) == 0o555