diff --git a/src/taskgraph/run-task/fetch-content b/src/taskgraph/run-task/fetch-content index 787d9eec2..337d807fb 100755 --- a/src/taskgraph/run-task/fetch-content +++ b/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( diff --git a/test/test_scripts_fetch_content.py b/test/test_scripts_fetch_content.py index 658e0c8c2..a971c22c5 100644 --- a/test/test_scripts_fetch_content.py +++ b/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