feat: add HEAD prefetch cache, and some cleanup

.gitignore

@@ -5,3 +5,5 @@
 .pre-commit-config.flake.yaml
 result/
 
+__pycache__/
+

Cargo.toml

@@ -21,6 +21,7 @@ bytes = "1.11"
 figment = { version = "0.10", features = ["toml"] }
 futures = "0.3"
 futures-util = "0.3"
+moka = { version = "0.12.15", features = ["future"] }
 percent-encoding = "2.1"
 prometheus = "0.14"
 regex = "1"

flake.nix

@@ -49,16 +49,38 @@
           ...
         }:
         let
-          inherit (pkgs) stdenv pkgsStatic;
+          inherit (pkgs)
+            stdenv
+            pkgsStatic
+            cacert
+            rust-jemalloc-sys
+            ;
           craneLib = (inputs.crane.mkLib pkgs).overrideToolchain (p: p.rustToolchain);
 
-          craneAttrs = import ./nix/crane.nix { inherit craneLib pkgs lib; };
-          inherit (craneAttrs)
-            src
-            commonArgs
-            cargoArtifacts
-            mergeCraneArgs
-            ;
+          src =
+            let
+              root = ./.;
+            in
+            lib.fileset.toSource {
+              inherit root;
+              fileset = lib.fileset.unions [
+                (craneLib.fileset.commonCargoSources root)
+                (lib.fileset.maybeMissing (root + "/tests"))
+              ];
+            };
+
+          # NOTE: `buildInputs` and sometimes `nativeBuildInputs`
+          # should be explicitly overridden for cross compilation
+          commonArgs = {
+            inherit src;
+            strictDeps = true;
+            nativeBuildInputs = [ cacert ];
+            buildInputs = [ rust-jemalloc-sys ];
+            doCheck = false; # Test separately with cargo-nextest
+          };
+          # Build *just* the cargo dependencies, so we can reuse
+          # all of that work (e.g. via cachix) when running in CI
+          cargoArtifacts = craneLib.buildDepsOnly commonArgs;
 
           defaultTarget = stdenv.hostPlatform.config;
           muslTarget =
@@ -77,18 +99,20 @@
               inherit cargoArtifacts;
               CARGO_PROFILE = "dev";
               CARGO_BUILD_TARGET = defaultTarget;
-              buildInputs = [ pkgs.rust-jemalloc-sys ];
             }
           );
 
           my-crate-musl = craneLib.buildPackage (
-            mergeCraneArgs commonArgs {
+            commonArgs
+            // {
               nativeBuildInputs = [
                 # Required by aws-lc-sys
+                cacert
                 stdenv.cc
                 pkgsStatic.stdenv.cc
               ];
               buildInputs = [ pkgsStatic.rust-jemalloc-sys ];
+              doCheck = true; # always run checkPhase for release artifact
               CARGO_PROFILE = "release";
               CARGO_BUILD_TARGET = muslTarget;
               CARGO_BUILD_FLAGS = "-C target-feature=+crt-static";

src/artifacts.rs

@@ -67,7 +67,7 @@ async fn remove_buffer_file(path: &Path) {
 async fn download_to_memory(response: Response, max_size: u64) -> Result<UploadPayload> {
     let body = response.bytes().await?;
     if body.len() as u64 > max_size {
-        return Err(Error::TooLarge(()));
+        return Err(Error::TooLarge);
     }
     Ok(UploadPayload::Memory(body))
 }
@@ -99,7 +99,7 @@ async fn download_to_file(
             })?;
             content_length += chunk.len() as u64;
             if content_length > max_size {
-                return Err(Error::TooLarge(()));
+                return Err(Error::TooLarge);
             }
             file.write_all(&chunk).await?;
         }
@@ -237,12 +237,12 @@ impl DownloadCtx<'static> {
             std::time::Duration::from_secs(config.download_timeout),
             task_fut,
         );
-        if let Err(err) = task_fut.await.unwrap_or(Err(Error::Timeout(()))) {
+        if let Err(err) = task_fut.await.unwrap_or(Err(Error::Timeout)) {
             warn!(error=?err, ttl=task_new.retry_limit, "failed to download task");
             task_new.retry_limit -= 1;
             self.metrics.failed_download_counter.inc();
 
-            if !matches!(err, Error::HTTPError(_)) && !matches!(err, Error::TooLarge(_)) {
+            if !matches!(err, Error::Http(_)) && !matches!(err, Error::TooLarge) {
                 self.fail_tx.send(task_new).unwrap();
                 self.metrics.task_in_queue.inc();
             }
@@ -302,7 +302,7 @@ async fn cache_task(task: Task, client: Client, config: &Config) -> Result<()> {
                 Ok(stream) => stream,
                 Err(err) => {
                     remove_buffer_file(&path).await;
-                    return Err(Error::CustomError(format!(
+                    return Err(Error::Custom(format!(
                         "failed to open buffered artifact {}: {:?}",
                         path.display(),
                         err
@@ -325,14 +325,14 @@ async fn download_payload(client: &Client, url: Url, config: &Config) -> Result<
     let response = client.get(url).send().await?;
     let status = response.status();
     if !status.is_success() {
-        return Err(Error::HTTPError(status));
+        return Err(Error::Http(status));
     }
 
     let max_size = max_download_size(config);
     let file_threshold = file_threshold(config);
 
     match response.content_length() {
-        Some(content_length) if content_length > max_size => Err(Error::TooLarge(())),
+        Some(content_length) if content_length > max_size => Err(Error::TooLarge),
         Some(content_length) if content_length > file_threshold => {
             info!("stream mode: file backend");
             download_to_file(response, config, max_size).await
@@ -489,6 +489,6 @@ mod tests {
             &config,
         )
         .await;
-        assert!(matches!(result, Err(Error::TooLarge(()))));
+        assert!(matches!(result, Err(Error::TooLarge)));
     }
 }

src/browse.rs

@@ -66,7 +66,7 @@ pub async fn list(
             .send(),
     )
     .await
-    .map_err(|_| Error::Timeout(()))?;
+    .map_err(|_| Error::Timeout)?;
 
     if result.is_ok() {
         return Ok(Redirect::Permanent(format!("{}{}", real_endpoint, mirror_clone_list)).into());
@@ -85,7 +85,7 @@ pub async fn list(
             .send(),
     )
     .await
-    .map_err(|_| Error::Timeout(()))??;
+    .map_err(|_| Error::Timeout)??;
 
     let mut body = r#"<tr>
             <td><a href="..?mirror_intel_list">..</a></td>

src/common.rs

@@ -15,6 +15,7 @@ use std::sync::Arc;
 use tokio::sync::mpsc::Sender;
 use url::Url;
 
+use crate::s3_cache::PrefetchCache;
 use crate::{Error, Result};
 
 /// A cache task.
@@ -57,7 +58,7 @@ impl Task {
             self.storage,
             percent_decode(self.path.as_bytes())
                 .decode_utf8()
-                .map_err(|_| Error::DecodePathError(()))?
+                .map_err(|_| Error::DecodePath)?
         ))
     }
 
@@ -149,12 +150,16 @@ pub struct IntelMission {
     pub tx: Option<Sender<Task>>,
     /// Reqwest client.
     pub client: Client,
+    /// Reqwest client for HEAD prefetch probes.
+    pub prefetch_client: Client,
     /// Prometheus metrics.
     pub metrics: Arc<Metrics>,
     /// S3 client.
     ///
     /// This is an anonymous client.
     pub s3_client: Arc<S3Client>,
+    /// Positive HEAD prefetch cache for `RouteAction::Cache`.
+    pub prefetch_cache: Arc<PrefetchCache>,
 }
 
 /// An upstream endpoint override rule.
@@ -273,6 +278,8 @@ pub struct Config {
     pub buffer_path: PathBuf,
     /// Worker tasks to serve requests.
     pub workers: Option<usize>,
+    /// TTL in seconds for positive S3/upstream HEAD prefetch entries.
+    pub head_prefetch_cache_ttl_secs: Option<u64>,
 }
 
 /// An empty redirect response to a given URL.
@@ -457,6 +464,7 @@ mod tests {
                 },
                 buffer_path: "/mnt/cache/".into(),
                 workers: None,
+                head_prefetch_cache_ttl_secs: None,
             };
             assert_eq!(config, expected);
             Ok(())