feat(fetch-content): record a profiler marker for each fetched file

src/taskgraph/run-task/fetch-content

@@ -53,10 +53,54 @@
 
 
 def log(msg):
     print(msg, file=sys.stderr)
     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."""
 
@@ -971,6 +1015,21 @@
         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.
 
@@ -984,7 +1043,7 @@
         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:
@@ -1179,81 +1238,83 @@
 
 
 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,
-            fetch_mode=args.fetch_mode,
-        )
-    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,
+                fetch_mode=args.fetch_mode,
+            )
+        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):