fix: slowness in download endpoint (#1065)

1. It was meant to limit to most recent 5 versions, but it was querying
all.
2. It was doing a bunch of queries in series instead of querying all
versions in a single call.

Author: lucacasonato

api/.sqlx/query-d85e9c34ba1e0bf729603f81093c42c4f48fa8ab9ea6330062c4b5fbc7df09e7.json renamed to api/.sqlx/query-47775e53578b988ebd14f2d9bbbeb81c82a13ab9a73ac6f6d158fde2cd563188.json

@@ -1647,35 +1647,51 @@ pub async fn get_downloads_handler(
   let current = Utc::now();
   let start = current - chrono::Duration::days(90);
 
-  let total = db
-    .get_package_downloads_24h(&scope, &package, start, current)
-    .await?;
+  let total_fut = async {
+    db.get_package_downloads_24h(&scope, &package, start, current)
+      .await
+      .map_err(ApiError::from)
+  };
 
-  let recent_versions = db
-    .list_latest_unyanked_versions_for_package(&scope, &package, 5)
-    .await?;
+  let recent_versions_fut = async {
+    let recent_versions = db
+      .list_latest_unyanked_versions_for_package(&scope, &package, 5)
+      .await?;
 
-  let versions = futures::stream::iter(recent_versions.into_iter())
-    .then(|version| async {
-      let res = db
-        .get_package_version_downloads_24h(
-          &scope, &package, &version, start, current,
-        )
-        .await;
-      (version, res)
-    })
-    .collect::<Vec<_>>()
-    .await;
+    let data_points = db
+      .get_package_versions_downloads_24h(
+        &scope,
+        &package,
+        &recent_versions,
+        start,
+        current,
+      )
+      .await?;
 
-  let mut recent_versions = Vec::with_capacity(versions.len());
-  for (version, downloads) in versions {
-    let downloads = downloads?
-      .into_iter()
-      .map(ApiDownloadDataPoint::from)
-      .collect();
-    recent_versions
-      .push(ApiPackageDownloadsRecentVersion { version, downloads });
-  }
+    let mut data_points_by_version =
+      indexmap::IndexMap::<_, Vec<_>>::with_capacity(recent_versions.len());
+
+    for data_point in data_points {
+      let version = data_point.version.clone();
+      let downloads = data_points_by_version
+        .entry(version)
+        .or_insert_with(Vec::new);
+      downloads.push(ApiDownloadDataPoint::from(data_point));
+    }
+
+    Ok::<_, ApiError>(
+      data_points_by_version
+        .into_iter()
+        .map(|(version, data_points)| ApiPackageDownloadsRecentVersion {
+          version,
+          downloads: data_points,
+        })
+        .collect(),
+    )
+  };
+
+  let (total, recent_versions) =
+    futures::try_join!(total_fut, recent_versions_fut)?;
 
   Ok(ApiPackageDownloads {
     total: total.into_iter().map(ApiDownloadDataPoint::from).collect(),

api/.sqlx/query-7b88c2e5072155bf7adaad3cfad86da140bf9bea4825836a563af75da02815b3.json renamed to api/.sqlx/query-c4621e8ce78b0df7bc634ff805ebbe65af37e70a1af5e25e6cc6e2f1354cea61.json

@@ -904,6 +904,16 @@ impl From<DownloadDataPoint> for ApiDownloadDataPoint {
   }
 }
 
+impl From<VersionDownloadDataPoint> for ApiDownloadDataPoint {
+  fn from(value: VersionDownloadDataPoint) -> Self {
+    Self {
+      time_bucket: value.time_bucket,
+      kind: value.kind.into(),
+      count: value.count as u64,
+    }
+  }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum ApiDownloadKind {

api/src/api/package.rs

@@ -1858,7 +1858,7 @@ impl Database {
   }
 
   #[instrument(
-    name = "Database::list_latest_package_versions",
+    name = "Database::list_latest_unyanked_versions_for_package",
     skip(self),
     err
   )]
@@ -1874,9 +1874,11 @@ impl Database {
       FROM package_versions
       WHERE scope = $1 AND name = $2 AND version NOT LIKE '%-%' AND is_yanked = false
       ORDER BY version DESC
+      LIMIT $3
       "#,
       scope as _,
       name as _,
+      limit as i32,
     )
       .map(|r| r.version)
       .fetch_all(&self.pool)
@@ -4017,25 +4019,25 @@ impl Database {
     skip(self),
     err
   )]
-  pub async fn get_package_version_downloads_24h(
+  pub async fn get_package_versions_downloads_24h(
     &self,
     scope: &ScopeName,
     name: &PackageName,
-    version: &Version,
+    versions: &[Version],
     start: DateTime<Utc>,
     end: DateTime<Utc>,
-  ) -> Result<Vec<DownloadDataPoint>> {
+  ) -> Result<Vec<VersionDownloadDataPoint>> {
     sqlx::query_as!(
-      DownloadDataPoint,
+      VersionDownloadDataPoint,
       r#"
-      SELECT time_bucket, kind as "kind: DownloadKind", count
+      SELECT version as "version: Version", time_bucket, kind as "kind: DownloadKind", count
       FROM version_download_counts_24h
-      WHERE scope = $1 AND package = $2 AND version = $3 AND time_bucket >= $4 AND time_bucket < $5
-      ORDER BY time_bucket ASC
+      WHERE scope = $1 AND package = $2 AND version = ANY($3) AND time_bucket >= $4 AND time_bucket < $5
+      ORDER BY time_bucket ASC, version DESC
       "#,
       scope as _,
       name as _,
-      version as _,
+      versions as _,
       start,
       end,
     )

api/src/api/types.rs

@@ -656,10 +656,10 @@ mod tests {
 
     // 24 hour window
     let downloads = db
-      .get_package_version_downloads_24h(
+      .get_package_versions_downloads_24h(
         &std,
         &fs,
-        &v0_215_0,
+        &[v0_215_0.clone()],
         "2024-06-01T00:00:00Z".parse().unwrap(),
         "2024-07-31T00:00:00Z".parse().unwrap(),
       )
@@ -670,13 +670,14 @@ mod tests {
       downloads[0].time_bucket,
       "2024-07-16T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
     );
+    assert_eq!(downloads[0].version, v0_215_0);
     assert_eq!(downloads[0].count, 13);
 
     let downloads = db
-      .get_package_version_downloads_24h(
+      .get_package_versions_downloads_24h(
         &luca,
         &flag,
-        &v1_0_0,
+        &[v1_0_0.clone()],
         "2024-06-01T00:00:00Z".parse().unwrap(),
         "2024-07-31T00:00:00Z".parse().unwrap(),
       )
@@ -687,6 +688,7 @@ mod tests {
       downloads[0].time_bucket,
       "2024-07-16T00:00:00Z".parse::<DateTime<Utc>>().unwrap()
     );
+    assert_eq!(downloads[0].version, v1_0_0);
     assert_eq!(downloads[0].count, 238);
   }
 }

api/src/db/database.rs

@@ -188,9 +188,9 @@ function getSeries(
   );
 
   const dataPointsToDisplay = dataPointsWithDownloads.slice(0, 5);
-  const others = dataPointsWithDownloads.slice(5).map((dataPoints) =>
-    dataPoints.downloads
-  ).flat();
+  // const others = dataPointsWithDownloads.slice(5).map((dataPoints) =>
+  //   dataPoints.downloads
+  // ).flat();
 
   const xValues = collectX(
     dataPointsWithDownloads.map((version) => version.downloads).flat(),
@@ -202,10 +202,10 @@ function getSeries(
       name: version.version,
       data: normalize(version.downloads, xValues, aggregationPeriod),
     })),
-    {
-      name: "Other",
-      data: normalize(others, xValues, aggregationPeriod),
-      color: "#e7c50b", // jsr-yellow-500
-    },
+    // {
+    //   name: "Other",
+    //   data: normalize(others, xValues, aggregationPeriod),
+    //   color: "#e7c50b", // jsr-yellow-500
+    // },
   ];
 }