Make it possible to load self-profile data from local filesystem

Dockerfile

@@ -56,4 +56,6 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
 COPY --from=build /target/release/postgres-to-sqlite /usr/local/bin/rustc-perf-postgres-to-sqlite
 COPY --from=build /target/release/site /usr/local/bin/rustc-perf-site
 
+ENV SELF_PROFILE_STORAGE_S3=1
+
 CMD rustc-perf-site

collector/src/bin/collector.rs

@@ -1041,9 +1041,7 @@ fn main_result() -> anyhow::Result<i32> {
                 backends,
                 iterations: Some(iterations),
                 self_profile_storage: if self_profile.self_profile {
-                    Some(Box::new(LocalSelfProfileStorage::new(Path::new(
-                        "self-profile-storage",
-                    ))))
+                    Some(Box::new(LocalSelfProfileStorage::default_path()))
                 } else {
                     None
                 },

collector/src/lib.rs

@@ -17,7 +17,9 @@ pub mod utils;
 
 use crate::compile::benchmark::{Benchmark, BenchmarkName};
 use crate::runtime::{BenchmarkGroup, BenchmarkSuite};
-pub use crate::self_profile::{LocalSelfProfileStorage, S3SelfProfileStorage, SelfProfileStorage};
+pub use crate::self_profile::{
+    LocalSelfProfileStorage, S3SelfProfileStorage, SelfProfileId, SelfProfileStorage,
+};
 use database::selector::CompileTestCase;
 use database::{ArtifactId, ArtifactIdNumber, Connection};
 use hashbrown::HashSet;

collector/src/self_profile.rs

@@ -1,16 +1,19 @@
 use crate::compile::benchmark::BenchmarkName;
 use crate::compile::execute::SelfProfileFiles;
+use analyzeme::ProfilingData;
 use anyhow::Context;
 use database::{ArtifactIdNumber, CollectionId, Profile, Scenario};
+use reqwest::StatusCode;
 use std::future::Future;
-use std::io::Read;
+use std::io::{Cursor, Read};
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
+use std::time::Instant;
 
 // TODO: Record codegen backend in the self profile name
 // TODO: remove collection ID from self-profile ID
 /// Uniquely identifies a self-profile archive.
-#[derive(Debug)]
+#[derive(Debug, Hash, PartialEq, Eq, Clone)]
 pub struct SelfProfileId {
     pub artifact_id_number: ArtifactIdNumber,
     pub collection: CollectionId,
@@ -36,6 +39,32 @@ pub trait SelfProfileStorage {
         id: SelfProfileId,
         files: SelfProfileFiles,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send>>;
+
+    /// Load self-profile data with the given ID.
+    /// Returns `None` if data for the ID was not found.
+    fn load(
+        &self,
+        id: SelfProfileId,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Option<analyzeme::ProfilingData>>> + Send>>
+    {
+        let fut = self.load_raw(id);
+        Box::pin(async move {
+            let Some(data) = fut.await? else {
+                return Ok(None);
+            };
+            Ok(Some(ProfilingData::from_paged_buffer(data, None).map_err(
+                |e| anyhow::anyhow!("Cannot parse self-profile data: {e}"),
+            )?))
+        })
+    }
+
+    /// Load the raw byte data of the self-profile with the given ID.
+    /// Returns `None` if data for the ID was not found.
+    #[allow(clippy::type_complexity)]
+    fn load_raw(
+        &self,
+        id: SelfProfileId,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>>;
 }
 
 pub struct LocalSelfProfileStorage {
@@ -48,6 +77,18 @@ impl LocalSelfProfileStorage {
             directory: dir.to_owned(),
         }
     }
+
+    /// Create local self-profile storage located at a local path on disk.
+    pub fn default_path() -> Self {
+        Self::new(Path::new("self-profile-storage"))
+    }
+
+    fn path(&self, id: &SelfProfileId) -> PathBuf {
+        let prefix = id.file_prefix();
+        self.directory
+            .join(prefix)
+            .join(format!("self-profile-{}.mm_profdata.sz", id.collection))
+    }
 }
 
 impl SelfProfileStorage for LocalSelfProfileStorage {
@@ -56,11 +97,7 @@ impl SelfProfileStorage for LocalSelfProfileStorage {
         id: SelfProfileId,
         files: SelfProfileFiles,
     ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send>> {
-        let prefix = id.file_prefix();
-        let path = self
-            .directory
-            .join(prefix)
-            .join(format!("self-profile-{}.mm_profdata.sz", id.collection));
+        let path = self.path(&id);
         Box::pin(async move {
             tokio::fs::create_dir_all(path.parent().unwrap()).await?;
             match files {
@@ -72,6 +109,22 @@ impl SelfProfileStorage for LocalSelfProfileStorage {
             Ok(())
         })
     }
+
+    fn load_raw(
+        &self,
+        id: SelfProfileId,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>> {
+        let path = self.path(&id);
+        Box::pin(async move {
+            if !path.is_file() {
+                return Ok(None);
+            }
+            let data = tokio::fs::read(&path).await.with_context(|| {
+                anyhow::anyhow!("Cannot read self-profile data from {}", path.display())
+            })?;
+            Ok(Some(data))
+        })
+    }
 }
 
 #[derive(Default)]
@@ -83,6 +136,10 @@ impl S3SelfProfileStorage {
     }
 }
 
+fn s3_self_profile_filename(id: &SelfProfileId) -> String {
+    format!("self-profile-{}.mm_profdata.sz", id.collection)
+}
+
 impl SelfProfileStorage for S3SelfProfileStorage {
     fn store(
         &self,
@@ -108,10 +165,12 @@ impl SelfProfileStorage for S3SelfProfileStorage {
                         .await
                         .context("cannot write compressed self-profile data")?;
 
-                    format!("self-profile-{}.mm_profdata.sz", id.collection)
+                    s3_self_profile_filename(&id)
                 }
             };
 
+            log::info!("Uploading self-profile to {}/{filename}", prefix.display());
+
             let output = tokio::process::Command::new("aws")
                 .arg("s3")
                 .arg("cp")
@@ -138,4 +197,60 @@ impl SelfProfileStorage for S3SelfProfileStorage {
             Ok(())
         })
     }
+
+    fn load_raw(
+        &self,
+        id: SelfProfileId,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send>> {
+        let path = id.file_prefix().join(s3_self_profile_filename(&id));
+        let url = format!(
+            "https://perf-data.rust-lang.org/{}",
+            path.to_str().expect("Non UTF-8 path used for self-profile")
+        );
+        Box::pin(async move {
+            log::trace!("Downloading {url}");
+            let start = Instant::now();
+            let resp = match reqwest::get(&url).await {
+                Ok(r) => r,
+                Err(e) => return Err(anyhow::anyhow!("{e:?}")),
+            };
+
+            if !resp.status().is_success() {
+                // Hitting an unknown path is returned as forbidden
+                if resp.status() == StatusCode::FORBIDDEN {
+                    return Ok(None);
+                }
+                return Err(anyhow::anyhow!(
+                    "Upstream status {:?} is not successful.\nurl={url}",
+                    resp.status(),
+                ));
+            }
+
+            let compressed = match resp.bytes().await {
+                Ok(b) => b,
+                Err(e) => {
+                    return Err(anyhow::anyhow!("Could not download from upstream: {e:?}"));
+                }
+            };
+
+            log::trace!(
+                "downloaded {} bytes in {:?}",
+                compressed.len(),
+                start.elapsed()
+            );
+
+            // The decompression is blocking, so we should not do it in the async task directly
+            let data = tokio::task::spawn_blocking(move || {
+                let mut data = Vec::new();
+                match snap::read::FrameDecoder::new(Cursor::new(compressed)).read_to_end(&mut data)
+                {
+                    Ok(_) => Ok(data),
+                    Err(e) => Err(anyhow::anyhow!("Could not decode self-profile data: {e:?}")),
+                }
+            })
+            .await??;
+
+            Ok(Some(data))
+        })
+    }
 }

database/src/lib.rs

@@ -812,7 +812,7 @@ pub struct Step {
     pub expected: Duration,
 }
 
-#[derive(Debug, PartialEq, Eq, Copy, Clone)]
+#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
 pub struct CollectionId(i32);
 
 impl fmt::Display for CollectionId {
@@ -821,6 +821,15 @@ impl fmt::Display for CollectionId {
     }
 }
 
+impl CollectionId {
+    pub fn from_inner(cid: i32) -> Self {
+        Self(cid)
+    }
+    pub fn as_inner(&self) -> i32 {
+        self.0
+    }
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub struct CompileBenchmark {
     pub name: String,

database/src/pool.rs

@@ -166,7 +166,7 @@ pub trait Connection: Send + Sync {
         crate_: &str,
         profile: &str,
         cache: &str,
-    ) -> Vec<(ArtifactIdNumber, i32)>;
+    ) -> Vec<(ArtifactIdNumber, CollectionId)>;
 
     /// Removes all data associated with the given artifact.
     async fn purge_artifact(&self, aid: &ArtifactId);

database/src/pool/postgres.rs

@@ -1284,7 +1284,7 @@ where
         crate_: &str,
         profile: &str,
         scenario: &str,
-    ) -> Vec<(ArtifactIdNumber, i32)> {
+    ) -> Vec<(ArtifactIdNumber, CollectionId)> {
         self.conn()
             .query(
                 "
@@ -1307,7 +1307,12 @@ where
             .await
             .unwrap()
             .into_iter()
-            .map(|r| (ArtifactIdNumber(r.get::<_, i32>(0) as u32), r.get(1)))
+            .map(|r| {
+                (
+                    ArtifactIdNumber(r.get::<_, i32>(0) as u32),
+                    CollectionId(r.get(1)),
+                )
+            })
             .collect()
     }
 

database/src/pool/sqlite.rs

@@ -1041,7 +1041,7 @@ impl Connection for SqliteConnection {
         crate_: &str,
         profile: &str,
         scenario: &str,
-    ) -> Vec<(ArtifactIdNumber, i32)> {
+    ) -> Vec<(ArtifactIdNumber, CollectionId)> {
         self.raw_ref()
             .prepare(
                 "select aid, cid from raw_self_profile where
@@ -1061,7 +1061,12 @@ impl Connection for SqliteConnection {
                         ArtifactId::Tag(a) => a,
                     }
                 ],
-                |r| Ok((ArtifactIdNumber(r.get::<_, i32>(0)? as u32), r.get(1)?)),
+                |r| {
+                    Ok((
+                        ArtifactIdNumber(r.get::<_, i32>(0)? as u32),
+                        CollectionId(r.get(1)?),
+                    ))
+                },
             )
             .unwrap()
             .collect::<Result<_, _>>()

site/frontend/src/components/perfetto-link.vue

@@ -13,8 +13,9 @@ const props = defineProps<{
 const link = computed(() => {
   return chromeProfileUrl(
     props.artifact.commit,
-    `${props.testCase.benchmark}-${props.testCase.profile.toLowerCase()}`,
-    props.testCase.scenario
+    `${props.testCase.benchmark}`,
+    props.testCase.scenario,
+    props.testCase.profile.toLowerCase()
   );
 });
 

site/frontend/src/pages/detailed-query/page.vue

@@ -194,11 +194,19 @@ async function loadData() {
   // Load sort state from URL
   loadSortFromUrl(params);
 
+  // To maintain backwards compatibility of URLs, parse the profile from the
+  // benchmark URL parameter
+  // We turn <benchmark>-<profile> into two separate parameters for the backend
+  const parts = benchmark.split("-");
+  const profile = parts[parts.length - 1];
+  const benchmarkSeparate = parts.slice(0, parts.length - 1).join("-");
+
   const currentSelector: Selector = {
     commit,
     base_commit: base_commit ?? null,
-    benchmark,
+    benchmark: benchmarkSeparate,
     scenario,
+    profile,
     backend,
     target,
   };