Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 121 additions & 60 deletions src/taskgraph/run-task/fetch-content
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,50 @@ def log(msg):
sys.stderr.flush()


def open_for_append(path):
"""Open path for writes that are each appended atomically, even when
several processes append to it concurrently."""
if sys.platform != "win32":
return os.open(path, os.O_WRONLY | os.O_APPEND)

import _winapi
import msvcrt

# O_APPEND seeks to the end before each write on Windows, which isn't
# atomic. Writes through a handle that may only append go to the end of
# the file atomically.
FILE_APPEND_DATA = 0x4
FILE_SHARE_READ_WRITE_DELETE = 0x7
handle = _winapi.CreateFile(
path,
FILE_APPEND_DATA,
FILE_SHARE_READ_WRITE_DELETE,
0,
_winapi.OPEN_EXISTING,
0,
0,
)
return msvcrt.open_osfhandle(handle, 0)


def record_marker(name, start, end, text):
"""Add a marker to the profile of the resource monitor run-task runs, if
any, using the forwarding protocol of mozsystemmonitor."""
path = os.environ.get("MOZ_RESOURCE_MONITOR_FORWARD")
if not path:
return
# 2 is the "Tasks" category.
record = ["marker", name, start, end, {"type": "Text", "text": text}, 2]
try:
fd = open_for_append(path)
try:
os.write(fd, (json.dumps(record) + "\n").encode("utf-8"))
finally:
os.close(fd)
except OSError:
pass


class IntegrityError(Exception):
"""Represents an integrity error when downloading a URL."""

Expand Down Expand Up @@ -971,6 +1015,21 @@ def fetch_and_extract(url, dest_dir, extract=True, sha256=None, size=None):
pass


@contextlib.contextmanager
def fetch_marker(url):
"""Record the duration of a fetch, including failures."""
start = time.monotonic()
try:
yield
finally:
record_marker("fetch-content", start, time.monotonic(), url)


def fetch_and_extract_with_marker(url, *args):
with fetch_marker(url):
fetch_and_extract(url, *args)


def fetch_urls(downloads):
"""Fetch URLs pairs to a pathlib.Path.

Expand All @@ -984,7 +1043,7 @@ def fetch_urls(downloads):
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)
future = e.submit(fetch_and_extract_with_marker, url, staging_dir, *rest)
extractions.append((future, staging_dir, dest_dir))

for future, staging_dir, dest_dir in extractions:
Expand Down Expand Up @@ -1154,80 +1213,82 @@ def git_checkout_archive(


def command_git_checkout_archive(args):
dest = pathlib.Path(args.dest)
with fetch_marker(args.repo):
dest = pathlib.Path(args.dest)

try:
git_checkout_archive(
dest,
args.repo,
args.commit,
prefix=args.path_prefix,
ssh_key=args.ssh_key_secret,
include_dot_git=args.include_dot_git,
)
except Exception:
try:
dest.unlink()
except FileNotFoundError:
pass
git_checkout_archive(
dest,
args.repo,
args.commit,
prefix=args.path_prefix,
ssh_key=args.ssh_key_secret,
include_dot_git=args.include_dot_git,
)
except Exception:
try:
dest.unlink()
except FileNotFoundError:
pass

raise
raise


def command_static_url(args):
gpg_sig_url = args.gpg_sig_url
gpg_env_key = args.gpg_key_env

if bool(gpg_sig_url) != bool(gpg_env_key):
print("--gpg-sig-url and --gpg-key-env must both be defined")
return 1
with fetch_marker(args.url):
gpg_sig_url = args.gpg_sig_url
gpg_env_key = args.gpg_key_env

if gpg_sig_url:
gpg_signature = b"".join(stream_download(gpg_sig_url))
gpg_key = os.environb[gpg_env_key.encode("ascii")]
if bool(gpg_sig_url) != bool(gpg_env_key):
print("--gpg-sig-url and --gpg-key-env must both be defined")
return 1

dest = pathlib.Path(args.dest)
dest.parent.mkdir(parents=True, exist_ok=True)
if gpg_sig_url:
gpg_signature = b"".join(stream_download(gpg_sig_url))
gpg_key = os.environb[gpg_env_key.encode("ascii")]

basename = urllib.parse.urlparse(args.url).path.split("/")[-1]
if basename.endswith("".join(dest.suffixes)):
dl_dest = dest
else:
dl_dest = dest.parent / basename
dest = pathlib.Path(args.dest)
dest.parent.mkdir(parents=True, exist_ok=True)

try:
download_to_path(
args.url, dl_dest, sha256=args.sha256, size=args.size, headers=args.headers
)

if gpg_sig_url:
gpg_verify_path(dl_dest, gpg_key, gpg_signature)
basename = urllib.parse.urlparse(args.url).path.split("/")[-1]
if basename.endswith("".join(dest.suffixes)):
dl_dest = dest
else:
dl_dest = dest.parent / basename

if args.force_archive or should_repack_archive(
dl_dest, dest, args.strip_components, args.add_prefix
):
repack_archive(
dl_dest,
dest,
args.strip_components,
args.add_prefix,
args.force_archive,
try:
download_to_path(
args.url, dl_dest, sha256=args.sha256, size=args.size, headers=args.headers
)
elif dl_dest != dest:
log(f"Renaming {dl_dest} to {dest}")
dl_dest.rename(dest)

except Exception:
try:
dl_dest.unlink()
except FileNotFoundError:
pass
if gpg_sig_url:
gpg_verify_path(dl_dest, gpg_key, gpg_signature)

if args.force_archive or should_repack_archive(
dl_dest, dest, args.strip_components, args.add_prefix
):
repack_archive(
dl_dest,
dest,
args.strip_components,
args.add_prefix,
args.force_archive,
)
elif dl_dest != dest:
log(f"Renaming {dl_dest} to {dest}")
dl_dest.rename(dest)

raise
except Exception:
try:
dl_dest.unlink()
except FileNotFoundError:
pass

if dl_dest != dest and dl_dest.exists():
log(f"Removing {dl_dest}")
dl_dest.unlink()
raise

if dl_dest != dest and dl_dest.exists():
log(f"Removing {dl_dest}")
dl_dest.unlink()


def api(root_url, service, version, path):
Expand Down
Loading