From 3f587156ec0b7a7edcba9e0fcdb72109e50bbddf Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Mon, 8 Jan 2024 11:13:59 +0100
Subject: [PATCH] Improve install instrumentation (#829)

Add tracing spans to different phases of the wheel installation.
---
 crates/install-wheel-rs/src/linker.rs    | 6 +++---
 crates/install-wheel-rs/src/wheel.rs     | 2 ++
 crates/puffin-installer/src/installer.rs | 2 ++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/crates/install-wheel-rs/src/linker.rs b/crates/install-wheel-rs/src/linker.rs
index 162a6da72..33099714c 100644
--- a/crates/install-wheel-rs/src/linker.rs
+++ b/crates/install-wheel-rs/src/linker.rs
@@ -6,7 +6,7 @@ use std::path::Path;
 use configparser::ini::Ini;
 use fs_err as fs;
 use fs_err::File;
-use tracing::{debug, info_span};
+use tracing::{debug, instrument};
 
 use pypi_types::DirectUrl;
 
@@ -24,6 +24,7 @@ use crate::{read_record_file, Error, Script};
 /// <https://packaging.python.org/en/latest/specifications/binary-distribution-format/#installing-a-wheel-distribution-1-0-py32-none-any-whl>
 ///
 /// Wheel 1.0: <https://www.python.org/dev/peps/pep-0427/>
+#[instrument(skip_all, fields(wheel = %wheel.as_ref().display()))]
 pub fn install_wheel(
     location: &InstallLocation<impl AsRef<Path>>,
     wheel: impl AsRef<Path>,
@@ -52,8 +53,6 @@ pub fn install_wheel(
     let metadata = dist_info_metadata(&dist_info_prefix, &wheel)?;
     let (name, _version) = parse_metadata(&dist_info_prefix, &metadata)?;
 
-    let _my_span = info_span!("install_wheel", name);
-
     // We're going step by step though
     // https://packaging.python.org/en/latest/specifications/binary-distribution-format/#installing-a-wheel-distribution-1-0-py32-none-any-whl
     // > 1.a Parse distribution-1.0.dist-info/WHEEL.
@@ -231,6 +230,7 @@ impl Default for LinkMode {
 
 impl LinkMode {
     /// Extract a wheel by linking all of its files into site packages.
+    #[instrument(skip_all)]
     pub fn link_wheel_files(
         self,
         site_packages: impl AsRef<Path>,
diff --git a/crates/install-wheel-rs/src/wheel.rs b/crates/install-wheel-rs/src/wheel.rs
index 98ca6e599..1f5d69470 100644
--- a/crates/install-wheel-rs/src/wheel.rs
+++ b/crates/install-wheel-rs/src/wheel.rs
@@ -432,6 +432,7 @@ pub(crate) fn parse_wheel_version(wheel_text: &str) -> Result<(), Error> {
 ///
 /// 2.f Compile any installed .py to .pyc. (Uninstallers should be smart enough to remove .pyc
 /// even if it is not mentioned in RECORD.)
+#[instrument(skip_all)]
 fn bytecode_compile(
     site_packages: &Path,
     unpacked_paths: Vec<PathBuf>,
@@ -723,6 +724,7 @@ fn install_script(
 
 /// Move the files from the .data directory to the right location in the venv
 #[allow(clippy::too_many_arguments)]
+#[instrument(skip_all)]
 pub(crate) fn install_data(
     venv_root: &Path,
     site_packages: &Path,
diff --git a/crates/puffin-installer/src/installer.rs b/crates/puffin-installer/src/installer.rs
index 10529e812..2bd1bd50a 100644
--- a/crates/puffin-installer/src/installer.rs
+++ b/crates/puffin-installer/src/installer.rs
@@ -1,5 +1,6 @@
 use anyhow::{Context, Error, Result};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use tracing::instrument;
 
 use distribution_types::CachedDist;
 use puffin_interpreter::Virtualenv;
@@ -36,6 +37,7 @@ impl<'a> Installer<'a> {
     }
 
     /// Install a set of wheels into a Python virtual environment.
+    #[instrument(skip_all, fields(num_wheels = %wheels.len()))]
     pub fn install(self, wheels: &[CachedDist]) -> Result<()> {
         tokio::task::block_in_place(|| {
             wheels.par_iter().try_for_each(|wheel| {