Endpoint for fetching tarballs

Author: Mark-Simulacrum

database/src/lib.rs

@@ -67,6 +67,10 @@ impl Date {
         Date(Utc.ymd(year, month, day).and_hms(h, m, s))
     }
 
+    pub fn empty() -> Date {
+        Date::ymd_hms(2000, 1, 1, 1, 1, 1)
+    }
+
     pub fn start_of_week(&self) -> Date {
         let weekday = self.0.weekday();
         // num_days_from_sunday is 0 for Sunday

database/src/pool.rs

@@ -106,6 +106,18 @@ pub trait Connection: Send + Sync {
     ///
     /// (Currently only works for try commits)
     async fn pr_of(&self, sha: &str) -> Option<u32>;
+
+    /// Returns the collection ids corresponding to the query. Usually just one.
+    ///
+    /// Currently only supported by postgres (sqlite does not store self-profile
+    /// results in the raw format).
+    async fn list_self_profile(
+        &self,
+        aid: ArtifactId,
+        crate_: &str,
+        profile: &str,
+        cache: &str,
+    ) -> Vec<(ArtifactIdNumber, i32)>;
 }
 
 #[async_trait::async_trait]

database/src/pool/postgres.rs

@@ -179,7 +179,6 @@ static MIGRATIONS: &[&str] = &[
         crate text not null references benchmark(name) on delete cascade on update cascade,
         profile text not null,
         cache text not null,
-        prefix integer not null generated always as identity,
         PRIMARY KEY(aid, cid, crate, profile, cache)
     );
     "#,
@@ -1045,4 +1044,36 @@ where
             &[&(artifact.0 as i16), &collection.0, &krate, &profile, &cache],
         ).await.unwrap();
     }
+    async fn list_self_profile(
+        &self,
+        aid: ArtifactId,
+        crate_: &str,
+        profile: &str,
+        cache: &str,
+    ) -> Vec<(ArtifactIdNumber, i32)> {
+        self.conn()
+            .query(
+                "
+            select aid, cid from raw_self_profile where
+                crate = $1
+                and profile = $2
+                and cache = $3
+                and aid = (select id from artifact where name = $4);
+        ",
+                &[
+                    &crate_,
+                    &profile,
+                    &cache,
+                    &match aid {
+                        ArtifactId::Commit(c) => c.sha,
+                        ArtifactId::Artifact(a) => a,
+                    },
+                ],
+            )
+            .await
+            .unwrap()
+            .into_iter()
+            .map(|r| (ArtifactIdNumber(r.get::<_, i16>(0) as u32), r.get(1)))
+            .collect()
+    }
 }

database/src/pool/sqlite.rs

@@ -853,4 +853,13 @@ impl Connection for SqliteConnection {
         // yet clear.
         unimplemented!("recording raw self profile files is not implemented for sqlite")
     }
+    async fn list_self_profile(
+        &self,
+        _aid: ArtifactId,
+        _crate_: &str,
+        _profile: &str,
+        _cache: &str,
+    ) -> Vec<(ArtifactIdNumber, i32)> {
+        Vec::new()
+    }
 }

site/src/api.rs

@@ -210,6 +210,25 @@ pub mod status {
     }
 }
 
+pub mod self_profile_raw {
+    use serde::{Deserialize, Serialize};
+
+    #[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
+    pub struct Request {
+        pub commit: String,
+        pub benchmark: String,
+        pub run_name: String,
+        pub cid: Option<i32>,
+    }
+
+    #[derive(Debug, Clone, Serialize)]
+    pub struct Response {
+        pub cids: Vec<i32>,
+        pub cid: i32,
+        pub url: String,
+    }
+}
+
 pub mod self_profile {
     use database::QueryLabel;
     use serde::{Deserialize, Serialize};

site/src/server.rs

@@ -37,8 +37,8 @@ type Request = http::Request<hyper::Body>;
 type Response = http::Response<hyper::Body>;
 
 pub use crate::api::{
-    self, dashboard, data, days, github, graph, info, self_profile, status, CommitResponse,
-    DateData, ServerResult, StyledBenchmarkName,
+    self, dashboard, data, days, github, graph, info, self_profile, self_profile_raw, status,
+    CommitResponse, DateData, ServerResult, StyledBenchmarkName,
 };
 use crate::db::{self, Cache, Crate, Profile};
 use crate::interpolate::Interpolated;
@@ -783,6 +783,77 @@ fn get_self_profile_data(
     Ok(profile)
 }
 
+pub async fn handle_self_profile_raw(
+    body: self_profile_raw::Request,
+    data: &InputData,
+) -> ServerResult<self_profile_raw::Response> {
+    log::info!("handle_self_profile_raw({:?})", body);
+    let mut it = body.benchmark.rsplitn(2, '-');
+    let bench_ty = it.next().ok_or(format!("no benchmark type"))?;
+    let bench_name = it.next().ok_or(format!("no benchmark name"))?;
+
+    let cache = body
+        .run_name
+        .parse::<database::Cache>()
+        .map_err(|e| format!("invalid run name: {:?}", e))?;
+
+    let conn = data.conn().await;
+
+    let aids_and_cids = conn
+        .list_self_profile(
+            ArtifactId::Commit(database::Commit {
+                sha: body.commit,
+                date: database::Date::empty(),
+            }),
+            bench_name,
+            bench_ty,
+            &body.run_name,
+        )
+        .await;
+    let (aid, first_cid) = aids_and_cids
+        .first()
+        .copied()
+        .ok_or_else(|| format!("No results for this commit"))?;
+
+    let cid = match body.cid {
+        Some(cid) => {
+            if aids_and_cids.iter().any(|(_, v)| *v == cid) {
+                cid
+            } else {
+                return Err(format!("{} is not a collection ID at this artifact", cid));
+            }
+        }
+        _ => first_cid,
+    };
+
+    let url = format!(
+        "https://perf-data.rust-lang.org/self-profile/{}/{}/{}/{}/self-profile-{}.tar.sz",
+        aid.0,
+        bench_name,
+        bench_ty,
+        cache.to_id(),
+        cid
+    );
+
+    let resp = reqwest::Client::new()
+        .head(&url)
+        .send()
+        .await
+        .map_err(|e| format!("fetching artifact: {:?}", e))?;
+    if !resp.status().is_success() {
+        return Err(format!(
+            "Artifact did not resolve successfully: {:?} received",
+            resp.status()
+        ));
+    }
+
+    Ok(self_profile_raw::Response {
+        cids: aids_and_cids.into_iter().map(|(_, cid)| cid).collect(),
+        cid,
+        url,
+    })
+}
+
 pub async fn handle_self_profile(
     body: self_profile::Request,
     data: &InputData,
@@ -1204,6 +1275,10 @@ async fn serve_req(ctx: Arc<Server>, req: Request) -> Result<Response, ServerErr
         Ok(to_response(
             handle_self_profile(body!(parse_body(&body)), &data).await,
         ))
+    } else if p == "/perf/self-profile-raw" {
+        Ok(to_response(
+            handle_self_profile_raw(body!(parse_body(&body)), &data).await,
+        ))
     } else {
         return Ok(http::Response::builder()
             .header_typed(ContentType::html())