chore(pegboard): send artifact image size from workflow instead of fetching with HEAD

Author: NathanFlurry

packages/edge/infra/client/manager/src/image_download_handler.rs

@@ -99,28 +99,28 @@ impl ImageDownloadHandler {
 				let mut conn = ctx.sql().await?;
 				let mut tx = conn.begin().await?;
 
-				let ((cache_count, images_dir_size), image_download_size) = tokio::try_join!(
-					async {
-						// Get total size of images directory. Note that it doesn't matter if this doesn't
-						// match the actual fs size because it should either be exactly at or below actual fs
-						// size. Also calculating fs size manually is expensive.
-						sqlx::query_as::<_, (i64, i64)>(indoc!(
-							"
-							SELECT COUNT(size), COALESCE(SUM(size), 0) FROM images_cache
-							",
-						))
-						.fetch_one(&mut *tx)
-						.await
-						.map_err(Into::<anyhow::Error>::into)
-					},
-					// NOTE: The image size here is somewhat misleading because its only the size of the
-					// downloaded archive and not the total disk usage after it is unpacked. However, this is
-					// good enough
-					self.fetch_image_download_size(ctx, image_config),
-				)?;
+				// Get total size of images directory. Note that it doesn't matter if this doesn't
+				// match the actual fs size because it should either be exactly at or below actual fs
+				// size. Also calculating fs size manually is expensive.
+				let (cache_count, images_dir_size) = sqlx::query_as::<_, (i64, i64)>(indoc!(
+					"
+					SELECT COUNT(size), COALESCE(SUM(size), 0) FROM images_cache
+					",
+				))
+				.fetch_one(&mut *tx)
+				.await
+				.map_err(Into::<anyhow::Error>::into)?;
 
 				// Prune images
-				let (removed_count, removed_bytes) = if images_dir_size as u64 + image_download_size
+				//
+				// HACK: The artifact_size_bytes here is somewhat misleading because its only the size of the
+				// downloaded archive and not the total disk usage after it is unpacked. However, this is size
+				// is recalculated later once decompressed, so this will only ever exceed the cache
+				// size limit in edge cases by `actual size - compressed size`. In this situation,
+				// that extra difference is already reserved on the file system by the actor
+				// itself.
+				let (removed_count, removed_bytes) = if images_dir_size as u64
+					+ image_config.artifact_size_bytes
 					> ctx.config().images.max_cache_size()
 				{
 					// Fetch as many images as it takes to clear up enough space for this new image.
@@ -157,7 +157,7 @@ impl ImageDownloadHandler {
 					.bind(image_config.id)
 					.bind(
 						(images_dir_size as u64)
-							.saturating_add(image_download_size)
+							.saturating_add(image_config.artifact_size_bytes)
 							.saturating_sub(ctx.config().images.max_cache_size()) as i64,
 					)
 					.fetch_all(&mut *tx)
@@ -202,7 +202,7 @@ impl ImageDownloadHandler {
 
 				metrics::IMAGE_CACHE_COUNT.set(cache_count + 1 - removed_count);
 				metrics::IMAGE_CACHE_SIZE
-					.set(images_dir_size + image_download_size as i64 - removed_bytes);
+					.set(images_dir_size + image_config.artifact_size_bytes as i64 - removed_bytes);
 
 				sqlx::query(indoc!(
 					"
@@ -517,53 +517,6 @@ impl ImageDownloadHandler {
 
 		Ok(addresses)
 	}
-
-	/// Attempts to fetch HEAD for the image download url and determine the image's download size.
-	async fn fetch_image_download_size(
-		&self,
-		ctx: &Ctx,
-		image_config: &protocol::Image,
-	) -> Result<u64> {
-		let addresses = self.get_image_addresses(ctx, image_config).await?;
-
-		let mut iter = addresses.into_iter();
-		while let Some(artifact_url) = iter.next() {
-			// Log the full URL we're attempting to download from
-			tracing::info!(image_id=?image_config.id, %artifact_url, "attempting to download image");
-
-			match reqwest::Client::new()
-				.head(&artifact_url)
-				.send()
-				.await
-				.and_then(|res| res.error_for_status())
-			{
-				Ok(res) => {
-					tracing::info!(image_id=?image_config.id, %artifact_url, "successfully fetched image HEAD");
-
-					// Read Content-Length header from response
-					let image_size = res
-						.headers()
-						.get(reqwest::header::CONTENT_LENGTH)
-						.context("no Content-Length header")?
-						.to_str()?
-						.parse::<u64>()
-						.context("invalid Content-Length header")?;
-
-					return Ok(image_size);
-				}
-				Err(err) => {
-					tracing::warn!(
-						image_id=?image_config.id,
-						%artifact_url,
-						%err,
-						"failed to fetch image HEAD"
-					);
-				}
-			}
-		}
-
-		bail!("artifact url could not be resolved");
-	}
 }
 
 /// Parses total bytes read from tar output.