gpu-trace-perf 1.8.2

//! Download and caching of trace files for replay.
//!
//! Callers receive a [`TraceFile`] token that provides the local path to the
//! trace, which locks the file in the cache for its lifetime.
//!
//! [`UploadForbidden`] is the typed error returned when an upload is rejected
//! with HTTP 403 Forbidden; callers can downcast to detect this case and apply
//! a fallback upload policy.
//!
//! A `disk_limit` bounds how much space files in the download cache may occupy.
//! [`TraceDownloader::get`] is an async future that yields until there is room,
//! so the caller's async loop can interleave downloads and replays efficiently
//! without blocking a thread.  It will use random replacement on files not
//! currently locked by a TraceFile, as necessary, to make space.
//!
//! If the cache directory exists, it enumerates the files at the beginning of
//! the run to populate the list of cached downloads available.  An optional
//! `local_base`, if given, is an immutable lookaside: files found there are
//! served directly; files not found there fall back to downloading.

use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use anyhow::{Context, Result};
use tokio::sync::Notify;
use walkdir::WalkDir;

/// Format a `SystemTime` as an RFC 1123 HTTP date string (e.g. for `Expires`).
fn http_date(t: std::time::SystemTime) -> String {
    chrono::DateTime::<chrono::Utc>::from(t)
        .format("%a, %d %b %Y %H:%M:%S GMT")
        .to_string()
}

struct CacheEntry {
    path: PathBuf,
    size: u64,
    lock_count: u32,
}

struct DownloaderInner {
    entries: HashMap<String, CacheEntry>,
    total_bytes: u64,
}

struct DownloaderState {
    inner: Mutex<DownloaderInner>,
    notify: Notify,
}

/// Error returned by [`TraceDownloader::upload_if_absent`] when the server
/// responds with HTTP 403 Forbidden.  The caller may choose to retry the
/// upload with a different URL (e.g. fall back from the stable snapshot URL
/// to a per-job URL) or treat it as a hard failure.
#[derive(Debug)]
pub struct UploadForbidden;

impl std::fmt::Display for UploadForbidden {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "HTTP 403 Forbidden")
    }
}

impl std::error::Error for UploadForbidden {}

pub struct TraceDownloader {
    /// Directory that downloaded traces are stored.  Files stored here are
    /// cached (enumerated and reused at the start of the run), up to
    /// disk_limit.
    download_dir: PathBuf,

    /// Maximum bytes that may be held in the cache at once.
    /// 0 means unlimited.
    disk_limit: u64,

    /// Bearer token used for requests to the S3 server.
    jwt: Option<String>,

    /// Local directory of immutable files used in preference over downloading.
    local_base: Option<PathBuf>,

    client: reqwest::Client,

    /// Client configured with no automatic redirect following, used for
    /// trace downloads so that we can manually strip the Authorization header
    /// on redirect (needed for caching proxies that redirect to pre-signed
    /// S3 URLs — see piglit's chase_redirects()).
    download_client: reqwest::Client,

    state: Arc<DownloaderState>,
}

/// A downloaded (or locally-sourced) trace file.
///
/// Holds a path to the trace on the local filesystem.  When dropped, the file
/// is unlocked in the cache (allowing it to be evicted later to make room for
/// new downloads).  Local files are left untouched on drop.
pub struct TraceFile {
    path: PathBuf,
    /// `Some((key, state))` for cached files; `None` for local files.
    cache_ref: Option<(String, Arc<DownloaderState>)>,
}

impl TraceFile {
    pub fn path(&self) -> &Path {
        &self.path
    }
}

impl Drop for TraceFile {
    fn drop(&mut self) {
        if let Some((key, state)) = &self.cache_ref {
            let mut inner = state.inner.lock().unwrap();
            if let Some(entry) = inner.entries.get_mut(key) {
                entry.lock_count = entry.lock_count.saturating_sub(1);
            }
            drop(inner);
            state.notify.notify_waiters();
        }
    }
}

/// Makes a HEAD request (with optional Bearer auth) and returns the response
/// headers on success, or `Err` on network failures or non-2xx status.
async fn head_headers(
    client: &reqwest::Client,
    jwt: Option<&str>,
    url: &str,
) -> Result<reqwest::header::HeaderMap> {
    let mut req = client.head(url);
    if let Some(token) = jwt {
        req = req.header("Authorization", format!("Bearer {token}"));
    }
    let response = req.send().await.with_context(|| format!("HEAD {url}"))?;
    if !response.status().is_success() {
        anyhow::bail!("HEAD {url}: HTTP {}", response.status());
    }
    Ok(response.headers().clone())
}

/// Strips S3 ETag quoting and weak-ETag prefix, lowercases the result.
fn normalize_etag(raw: &str) -> String {
    let s = raw.trim().strip_prefix("W/").unwrap_or(raw.trim());
    s.trim_matches('"').to_lowercase()
}

/// Returns the normalized ETag from a HEAD request, or `Ok(None)` if absent.
async fn fetch_server_etag(
    client: &reqwest::Client,
    jwt: Option<&str>,
    url: &str,
) -> Result<Option<String>> {
    let headers = head_headers(client, jwt, url).await?;
    Ok(headers
        .get(reqwest::header::ETAG)
        .and_then(|v| v.to_str().ok())
        .map(normalize_etag))
}

/// Guesses the S3 multipart chunk size from total file size and part count,
/// matching the heuristic in piglit's `download_utils.py`.
fn guess_s3_chunk_size(file_size: u64, num_parts: usize) -> u64 {
    const MB: u64 = 1024 * 1024;
    for &mb in &[5u64, 10] {
        let bytes = mb * MB;
        if file_size.div_ceil(bytes) == num_parts as u64 {
            return bytes;
        }
    }
    (file_size / num_parts as u64 / MB + 1) * MB
}

/// Computes the S3 ETag for a local file, using the same algorithm S3 used
/// (single-part MD5 or multipart MD5-of-MD5s) as indicated by `server_etag`.
fn compute_local_etag(path: &Path, server_etag: &str) -> Result<String> {
    let data = std::fs::read(path).with_context(|| format!("reading {}", path.display()))?;

    if let Some((_, n_str)) = server_etag.rsplit_once('-') {
        if let Ok(num_parts) = n_str.parse::<usize>() {
            // Multipart: MD5 of concatenated per-chunk MD5 digests.
            let chunk_size = guess_s3_chunk_size(data.len() as u64, num_parts) as usize;
            let part_digests: Vec<[u8; 16]> =
                data.chunks(chunk_size).map(|c| md5::compute(c).0).collect();
            let combined: Vec<u8> = part_digests
                .iter()
                .flat_map(|d| d.iter().copied())
                .collect();
            return Ok(format!(
                "{:x}-{}",
                md5::compute(&combined),
                part_digests.len()
            ));
        }
    }
    // Single-part: MD5 of the entire file.
    Ok(format!("{:x}", md5::compute(&data)))
}

impl TraceDownloader {
    /// Evicts unlocked cached files at random until `total_bytes < disk_limit`,
    /// then yields until a token drop creates enough room if eviction alone is
    /// insufficient (all remaining files are locked).
    async fn acquire_capacity(&self, trace_path: &str) {
        if self.disk_limit == 0 {
            return;
        }
        loop {
            // Register with the Notify *before* checking the condition to
            // avoid the lost-wakeup race: a drop() between the check and
            // the await would otherwise go unobserved.
            let notified = self.state.notify.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();

            let has_room = {
                let mut inner = self.state.inner.lock().unwrap();
                // Evict random unlocked files until we are under the limit.
                while inner.total_bytes >= self.disk_limit {
                    let unlocked: Vec<String> = inner
                        .entries
                        .iter()
                        .filter(|(_, e)| e.lock_count == 0)
                        .map(|(k, _)| k.clone())
                        .collect();
                    if unlocked.is_empty() {
                        break;
                    }
                    let key = unlocked[rand::random::<u64>() as usize % unlocked.len()].clone();
                    let entry = inner.entries.remove(&key).unwrap();
                    inner.total_bytes = inner.total_bytes.saturating_sub(entry.size);
                    if let Err(e) = std::fs::remove_file(&entry.path) {
                        log::warn!(
                            "Failed to remove evicted cache file {}: {e}",
                            entry.path.display()
                        );
                    }
                }
                inner.total_bytes < self.disk_limit
            };

            if has_room {
                return;
            }
            log::debug!("Waiting for disk space before downloading {trace_path}");
            notified.await;
        }
    }

    pub fn new(
        download_dir: PathBuf,
        disk_limit: u64,
        jwt_path: Option<&Path>,
        local_base: Option<PathBuf>,
    ) -> Result<Self> {
        let mut entries = HashMap::new();
        let mut total_bytes = 0u64;

        if local_base.is_none() {
            std::fs::create_dir_all(&download_dir)
                .with_context(|| format!("creating download dir {}", download_dir.display()))?;

            // Enumerate existing cached files, skipping incomplete downloads.
            for entry in WalkDir::new(&download_dir).min_depth(1) {
                let entry = match entry {
                    Ok(e) => e,
                    Err(_) => continue,
                };
                if !entry.file_type().is_file() {
                    continue;
                }
                let path = entry.path();
                if path.extension() == Some(OsStr::new("tmp")) {
                    continue;
                }
                let Some(name) = path
                    .strip_prefix(&download_dir)
                    .ok()
                    .and_then(|r| r.to_str())
                    .map(|s| s.to_string())
                else {
                    continue;
                };
                if let Ok(meta) = std::fs::metadata(path) {
                    let size = meta.len();
                    total_bytes += size;
                    entries.insert(
                        name,
                        CacheEntry {
                            path: path.to_path_buf(),
                            size,
                            lock_count: 0,
                        },
                    );
                }
            }
        }

        let jwt = match jwt_path {
            Some(path) => {
                let content = std::fs::read_to_string(path)
                    .with_context(|| format!("reading JWT file {}", path.display()))?;
                Some(content.trim().to_string())
            }
            None => None,
        };

        Ok(TraceDownloader {
            download_dir,
            disk_limit,
            jwt,
            local_base,
            client: reqwest::Client::new(),
            download_client: reqwest::Client::builder()
                .redirect(reqwest::redirect::Policy::none())
                .build()
                .context("building download client")?,
            state: Arc::new(DownloaderState {
                inner: Mutex::new(DownloaderInner {
                    entries,
                    total_bytes,
                }),
                notify: Notify::new(),
            }),
        })
    }

    /// Returns the root directory that should be used to compute a trace's
    /// display name.  When `--traces-db` is in use and the file lives under
    /// the local lookaside, returns the lookaside root; otherwise returns the
    /// download cache directory.
    pub fn root_for_file(&self, path: &Path) -> &Path {
        if let Some(local_base) = &self.local_base {
            if path.starts_with(local_base) {
                return local_base;
            }
        }
        &self.download_dir
    }

    /// Fetches the bytes at `url`, attaching the JWT bearer token if present.
    ///
    /// Returns an error on network failure or a non-2xx response.
    pub async fn fetch_bytes(&self, url: &str) -> Result<Vec<u8>> {
        let mut req = self.client.get(url);
        if let Some(token) = &self.jwt {
            req = req.header("Authorization", format!("Bearer {token}"));
        }
        let resp = req.send().await.with_context(|| format!("GET {url}"))?;
        if !resp.status().is_success() {
            anyhow::bail!("HTTP {} from {url}", resp.status());
        }
        Ok(resp
            .bytes()
            .await
            .with_context(|| format!("reading body from {url}"))?
            .to_vec())
    }

    /// Uploads `data` to `upload_base_url/<checksum>.png` unless the
    /// object already exists (checked via HEAD).  Returns the URL of the object.
    ///
    /// If `expires` is `Some`, an `Expires` header is added to the PUT so the
    /// s3-proxy can set a TTL on the object (matching ci_fairy.py's `--expires`
    /// behaviour).
    ///
    /// Skips the upload entirely in local mode (`--traces-db` was provided) and
    /// just returns the constructed URL.
    pub async fn upload_if_absent(
        &self,
        upload_base_url: &str,
        checksum: &str,
        data: &[u8],
        expires: Option<std::time::Duration>,
    ) -> Result<String> {
        let url = format!("{}/{}.png", upload_base_url.trim_end_matches('/'), checksum);

        // Check whether the image already exists.
        let mut head_req = self.client.head(&url);
        if let Some(token) = &self.jwt {
            head_req = head_req.header("Authorization", format!("Bearer {token}"));
        }
        let already_exists = match head_req.send().await {
            Ok(r) => r.status().is_success(),
            Err(_) => false,
        };

        if already_exists {
            log::debug!("Skipping upload of {url} (already exists)");
            return Ok(url);
        }

        log::debug!("Uploading {url}");
        // The s3-proxy expects a multipart PUT to the directory URL with the
        // file as a named part (matching ci_fairy.py's s3cp implementation).
        let dir_url = format!("{}/", upload_base_url.trim_end_matches('/'));
        let checksum_png = format!("{checksum}.png");
        let part = reqwest::multipart::Part::bytes(data.to_vec()).file_name(checksum_png);
        let form = reqwest::multipart::Form::new().part("file", part);
        let mut put_req = self
            .client
            .put(&dir_url)
            .header("x-amz-acl", "public-read-write")
            .multipart(form);
        if let Some(token) = &self.jwt {
            put_req = put_req.header("Authorization", format!("Bearer {token}"));
        }
        if let Some(ttl) = expires {
            put_req = put_req.header("Expires", http_date(std::time::SystemTime::now() + ttl));
        }
        let response = put_req
            .send()
            .await
            .with_context(|| format!("PUT {dir_url}"))?;
        if !response.status().is_success() {
            let status = response.status();
            if status == reqwest::StatusCode::FORBIDDEN {
                return Err(anyhow::Error::new(UploadForbidden));
            }
            let body = response
                .text()
                .await
                .unwrap_or_else(|_| "(unreadable body)".to_string());
            anyhow::bail!("HTTP {} uploading {}: {}", status, url, body);
        }

        Ok(url)
    }

    /// Returns a [`TraceFile`] for the given trace, serving from the local
    /// cache if available or downloading from the remote URL otherwise.
    ///
    /// Cached files are served without any network check; corruption is
    /// detected lazily by [`invalidate_if_corrupted`] after a replay failure.
    ///
    /// Before downloading, evicts unlocked cached files at random to make
    /// room, then yields (async) if all cached files are locked and the cache
    /// is still full.
    ///
    /// If `--traces-db` was provided, the file is looked up there first and
    /// returned with no network access and no cleanup on drop.  If the file is
    /// not found in `local_base`, the normal download path is used as a
    /// fallback.
    pub async fn get(&self, download_url: &str, trace_path: &str) -> Result<TraceFile> {
        if let Some(local_base) = &self.local_base {
            let path = local_base.join(trace_path);
            // Accept `path` directly, or the ANGLE subtest case where
            // `path.parent()` is the binary file (not a directory).
            let exists = path.exists()
                || path
                    .parent()
                    .and_then(|p| p.metadata().ok())
                    .is_some_and(|m| !m.is_dir());
            if exists {
                log::debug!("Using local trace: {}", path.display());
                return Ok(TraceFile {
                    path,
                    cache_ref: None,
                });
            }
            log::debug!(
                "Local trace not found in traces-db, falling back to download: {}",
                path.display()
            );
        }

        let filename = trace_path.trim_start_matches('/').to_string();

        // Serve from cache if available.
        {
            let mut inner = self.state.inner.lock().unwrap();
            if let Some(entry) = inner.entries.get_mut(&filename) {
                entry.lock_count += 1;
                let path = entry.path.clone();
                log::debug!("Serving {} from cache: {}", trace_path, path.display());
                return Ok(TraceFile {
                    path,
                    cache_ref: Some((filename, Arc::clone(&self.state))),
                });
            }
        }

        let url = format!(
            "{}/{}",
            download_url.trim_end_matches('/'),
            trace_path.trim_start_matches('/')
        );

        // Need to download — wait until there is room in the cache.
        self.acquire_capacity(trace_path).await;

        // Re-check: a concurrent get() for the same file may have completed
        // the download while we were waiting for capacity.
        {
            let mut inner = self.state.inner.lock().unwrap();
            if let Some(entry) = inner.entries.get_mut(&filename) {
                entry.lock_count += 1;
                log::debug!("Serving {} from cache after capacity wait", trace_path);
                return Ok(TraceFile {
                    path: entry.path.clone(),
                    cache_ref: Some((filename, Arc::clone(&self.state))),
                });
            }
        }

        let target_path = self.download_dir.join(&filename);
        if let Some(parent) = target_path.parent() {
            std::fs::create_dir_all(parent)
                .with_context(|| format!("creating directory {}", parent.display()))?;
        }
        log::debug!("Downloading {} to {}", trace_path, target_path.display());
        download_file(
            &self.download_client,
            self.jwt.as_deref(),
            &url,
            &target_path,
        )
        .await
        .with_context(|| format!("downloading {trace_path}"))?;
        log::debug!("Downloading {} completed", trace_path);

        let size = std::fs::metadata(&target_path)
            .with_context(|| format!("stat of {}", target_path.display()))?
            .len();

        let mut inner = self.state.inner.lock().unwrap();
        // A concurrent get() may have completed the same download while we
        // were downloading.  Latch onto their entry to avoid double-counting.
        if let Some(entry) = inner.entries.get_mut(&filename) {
            entry.lock_count += 1;
            return Ok(TraceFile {
                path: entry.path.clone(),
                cache_ref: Some((filename, Arc::clone(&self.state))),
            });
        }
        inner.total_bytes += size;
        inner.entries.insert(
            filename.clone(),
            CacheEntry {
                path: target_path.clone(),
                size,
                lock_count: 1,
            },
        );

        Ok(TraceFile {
            path: target_path,
            cache_ref: Some((filename, Arc::clone(&self.state))),
        })
    }

    /// Checks whether the cached copy of `trace_path` matches the server's
    /// ETag and removes the entry from the cache if the check fails or cannot
    /// be performed, so the next [`get`] re-downloads it.
    ///
    /// Returns `true` if the entry was invalidated, `false` if the ETag
    /// matched (file confirmed good).
    ///
    /// Call this after a trace replay fails to detect both truncation and
    /// byte-for-byte corruption.
    pub async fn invalidate_if_corrupted(
        &self,
        download_url: &str,
        trace_path: &str,
    ) -> Result<bool> {
        let filename = trace_path.trim_start_matches('/').to_string();
        // Note: entries is only the cached files, not the local_base files,
        // which are immutable.
        let local_path = {
            let inner = self.state.inner.lock().unwrap();
            inner.entries.get(&filename).map(|e| e.path.clone())
        };
        let local_path = match local_path {
            Some(p) => p,
            None => return Ok(false),
        };

        let url = format!(
            "{}/{}",
            download_url.trim_end_matches('/'),
            trace_path.trim_start_matches('/')
        );

        let server_etag = match fetch_server_etag(&self.client, self.jwt.as_deref(), &url).await? {
            Some(etag) => etag,
            None => {
                // Can't confirm the file is valid — treat it as corrupted.
                log::warn!(
                    "{trace_path}: server returned no ETag, evicting unverifiable cache entry"
                );
                let mut inner = self.state.inner.lock().unwrap();
                if let Some(entry) = inner.entries.remove(&filename) {
                    inner.total_bytes = inner.total_bytes.saturating_sub(entry.size);
                }
                if let Err(e) = std::fs::remove_file(&local_path) {
                    log::warn!(
                        "Failed to remove unverifiable cache file {}: {e}",
                        local_path.display()
                    );
                }
                drop(inner);
                self.state.notify.notify_waiters();
                return Ok(true);
            }
        };

        let local_etag = compute_local_etag(&local_path, &server_etag)
            .with_context(|| format!("computing ETag for {}", local_path.display()))?;

        if local_etag == server_etag {
            log::debug!("{trace_path}: ETag matches, cache is valid");
            return Ok(false);
        }

        log::warn!(
            "{trace_path}: ETag mismatch (local {local_etag}, server {server_etag}), \
             invalidating corrupted cache entry"
        );

        let mut inner = self.state.inner.lock().unwrap();
        if let Some(entry) = inner.entries.remove(&filename) {
            inner.total_bytes = inner.total_bytes.saturating_sub(entry.size);
        }
        if let Err(e) = std::fs::remove_file(&local_path) {
            log::warn!(
                "Failed to remove invalidated cache file {}: {e}",
                local_path.display()
            );
        }
        drop(inner);
        self.state.notify.notify_waiters();

        Ok(true)
    }
}

/// Sends a GET request to `url`, manually following redirects while stripping
/// the Authorization header on redirect.
///
/// When using a passthrough caching proxy, the hostname in the URI does not
/// change from reqwest's perspective (all requests go to the proxy), so
/// reqwest's default redirect handling does not strip Authorization.  The
/// proxy redirects to a pre-signed S3 URL; if Authorization is still present
/// on that request the S3 backend returns HTTP 400.  We strip it on every
/// redirect to match piglit's `chase_redirects()` behaviour.
///
/// `client` must be built with `redirect::Policy::none()`.
async fn chase_and_get(
    client: &reqwest::Client,
    jwt: Option<&str>,
    url: &str,
) -> Result<reqwest::Response> {
    const MAX_REDIRECTS: usize = 10;
    let mut current_url = url.to_string();
    let mut send_auth = true;

    for _ in 0..=MAX_REDIRECTS {
        let mut req = client.get(&current_url);
        if send_auth {
            if let Some(token) = jwt {
                req = req.header("Authorization", format!("Bearer {token}"));
            }
        }
        let resp = req
            .send()
            .await
            .with_context(|| format!("GET {current_url}"))?;
        if !resp.status().is_redirection() {
            return Ok(resp);
        }
        let location = resp
            .headers()
            .get(reqwest::header::LOCATION)
            .and_then(|v| v.to_str().ok())
            .ok_or_else(|| anyhow::anyhow!("redirect from {current_url} had no Location header"))?
            .to_string();
        // Resolve relative Location URLs against the current URL.
        let base = reqwest::Url::parse(&current_url)
            .with_context(|| format!("parsing current URL {current_url}"))?;
        let next = base
            .join(&location)
            .with_context(|| format!("resolving Location {location:?} against {current_url}"))?;
        send_auth = false;
        current_url = next.to_string();
    }
    anyhow::bail!("too many redirects from {url}");
}

/// Downloads `url` to `target_path` via a unique temp file in the same
/// directory to avoid leaving partial files on crash and to allow concurrent
/// downloads of the same target without interleaved writes.
///
/// `client` must be built with `redirect::Policy::none()`; redirects are
/// followed via [`chase_and_get`] which strips Authorization on redirect.
async fn download_file(
    client: &reqwest::Client,
    jwt: Option<&str>,
    url: &str,
    target_path: &Path,
) -> Result<()> {
    let response = chase_and_get(client, jwt, url)
        .await
        .with_context(|| format!("sending GET {url}"))?;

    let status = response.status();
    let content_type = response
        .headers()
        .get(reqwest::header::CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string();
    let final_url = response.url().clone();

    if !status.is_success() {
        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
            anyhow::bail!(
                "HTTP {status} downloading {url} (authentication failure - is --jwt required?)"
            );
        }
        anyhow::bail!("HTTP {} downloading {}", status, url);
    }

    // Detect auth redirects: a server that redirects unauthenticated requests
    // to a login page will return 200 OK with an HTML body.  Trace files are
    // never HTML, so this always indicates a configuration problem.
    if content_type.contains("text/html") {
        anyhow::bail!(
            "download from {url} returned HTML instead of a trace file \
             (authentication redirect to {final_url}? try --jwt)"
        );
    }

    let bytes = response
        .bytes()
        .await
        .with_context(|| format!("reading response body from {url}"))?;

    let parent = target_path.parent().unwrap_or(Path::new("."));
    let mut tmp = tempfile::NamedTempFile::new_in(parent)
        .with_context(|| format!("creating temp file in {}", parent.display()))?;
    std::io::Write::write_all(&mut tmp, &bytes)
        .with_context(|| format!("writing to temp file for {}", target_path.display()))?;
    tmp.persist(target_path)
        .map_err(|e| anyhow::anyhow!("renaming temp file to {}: {}", target_path.display(), e))?;

    Ok(())
}

#[cfg(test)]
impl TraceDownloader {
    /// Test helper: adds `path` to the cache as a locked entry of `size` bytes
    /// immediately (no waiting).  Use this to set up the "limit full" state.
    fn make_test_token(&self, path: PathBuf, size: u64) -> TraceFile {
        let key = path.file_name().unwrap().to_str().unwrap().to_string();
        let mut inner = self.state.inner.lock().unwrap();
        inner.entries.insert(
            key.clone(),
            CacheEntry {
                path: path.clone(),
                size,
                lock_count: 1,
            },
        );
        inner.total_bytes += size;
        TraceFile {
            path,
            cache_ref: Some((key, Arc::clone(&self.state))),
        }
    }

    /// Test helper: waits for capacity (via the real `acquire_capacity`) then
    /// adds `path` to the cache as a locked entry instead of downloading.
    async fn wait_and_make_token(&self, path: PathBuf) -> TraceFile {
        let key = path.file_name().unwrap().to_str().unwrap().to_string();
        let size = std::fs::metadata(&path).unwrap().len();
        self.acquire_capacity(&path.display().to_string()).await;
        let mut inner = self.state.inner.lock().unwrap();
        inner.entries.insert(
            key.clone(),
            CacheEntry {
                path: path.clone(),
                size,
                lock_count: 1,
            },
        );
        inner.total_bytes += size;
        TraceFile {
            path,
            cache_ref: Some((key, Arc::clone(&self.state))),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{
        Arc,
        atomic::{AtomicBool, AtomicU32, Ordering},
    };

    use bytes::Bytes;
    use http_body_util::Full;
    use hyper::body::Incoming;
    use hyper::{Method, Request, Response};
    use hyper_util::rt::{TokioExecutor, TokioIo};
    use hyper_util::server::conn::auto::Builder as ConnBuilder;
    use std::convert::Infallible;
    use tokio::net::TcpListener;

    use super::*;

    /// Starts a minimal HTTP server that always responds with the given `status`
    /// and optional `content_type` header.  Returns the `http://127.0.0.1:<port>`
    /// base URL.  The server is driven by the current tokio runtime.
    async fn start_static_server(status: u16, content_type: Option<&'static str>) -> String {
        let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        std_listener.set_nonblocking(true).unwrap();
        let port = std_listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let listener = TcpListener::from_std(std_listener).unwrap();
            let http = Arc::new(ConnBuilder::new(TokioExecutor::new()));
            loop {
                let (stream, _) = match listener.accept().await {
                    Ok(x) => x,
                    Err(_) => break,
                };
                let http = http.clone();
                tokio::spawn(async move {
                    let svc =
                        hyper::service::service_fn(move |_req: Request<Incoming>| async move {
                            let mut builder = Response::builder().status(status);
                            if let Some(ct) = content_type {
                                builder = builder.header("content-type", ct);
                            }
                            Ok::<_, Infallible>(builder.body(Full::new(Bytes::new())).unwrap())
                        });
                    let _ = http.serve_connection(TokioIo::new(stream), svc).await;
                });
            }
        });

        format!("http://127.0.0.1:{port}")
    }

    /// Starts an HTTP server that serves `body` for every request.
    /// For HEAD requests the Content-Length and ETag reflect the body but no
    /// body is sent.  `get_count` is incremented for each GET (not HEAD).
    async fn start_mock_file_server(body: Bytes, get_count: Arc<AtomicU32>) -> String {
        let etag = format!("{:x}", md5::compute(&body));
        let std_listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        std_listener.set_nonblocking(true).unwrap();
        let port = std_listener.local_addr().unwrap().port();

        tokio::spawn(async move {
            let listener = TcpListener::from_std(std_listener).unwrap();
            let http = Arc::new(ConnBuilder::new(TokioExecutor::new()));
            loop {
                let (stream, _) = match listener.accept().await {
                    Ok(x) => x,
                    Err(_) => break,
                };
                let http = Arc::clone(&http);
                let body = body.clone();
                let etag = etag.clone();
                let get_count = Arc::clone(&get_count);
                tokio::spawn(async move {
                    let svc = hyper::service::service_fn(move |req: Request<Incoming>| {
                        let body = body.clone();
                        let etag = etag.clone();
                        let get_count = Arc::clone(&get_count);
                        async move {
                            let builder = Response::builder()
                                .status(200)
                                .header("content-type", "application/octet-stream")
                                .header("content-length", body.len())
                                .header("etag", format!("\"{etag}\""));
                            let response_body = if req.method() == Method::HEAD {
                                Bytes::new()
                            } else {
                                get_count.fetch_add(1, Ordering::SeqCst);
                                body
                            };
                            Ok::<_, Infallible>(builder.body(Full::new(response_body)).unwrap())
                        }
                    });
                    let _ = http.serve_connection(TokioIo::new(stream), svc).await;
                });
            }
        });

        format!("http://127.0.0.1:{port}")
    }

    /// A 401 response from the server should produce a clear error that mentions
    /// both the HTTP status and `--jwt`.
    #[tokio::test]
    async fn download_401_reports_auth_error() {
        let tmpdir = tempfile::tempdir().unwrap();
        let base_url = start_static_server(401, None).await;
        let downloader = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();

        let msg = match downloader.get(&base_url, "trace.mock-trace").await {
            Ok(_) => panic!("expected an error on HTTP 401"),
            Err(e) => format!("{e:#}"),
        };
        assert!(msg.contains("401"), "error should mention HTTP 401: {msg}");
        assert!(msg.contains("--jwt"), "error should suggest --jwt: {msg}");
    }

    /// A 200 OK response whose Content-Type is text/html indicates an auth
    /// redirect (the server sent a login page).  The error should mention HTML
    /// and suggest --jwt.
    #[tokio::test]
    async fn download_html_response_reports_auth_error() {
        let tmpdir = tempfile::tempdir().unwrap();
        let base_url = start_static_server(200, Some("text/html; charset=utf-8")).await;
        let downloader = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();

        let msg = match downloader.get(&base_url, "trace.mock-trace").await {
            Ok(_) => panic!("expected an error when server returns HTML"),
            Err(e) => format!("{e:#}"),
        };
        assert!(
            msg.to_lowercase().contains("html"),
            "error should mention HTML response: {msg}"
        );
        assert!(msg.contains("--jwt"), "error should suggest --jwt: {msg}");
    }

    /// When the cache is full and all cached files are locked, `get` (via
    /// `wait_and_make_token`) must block until a token is dropped.  On drop,
    /// the unlocked file is evicted to make room and the file stays on disk
    /// until it is evicted, not until the token is dropped.
    #[tokio::test]
    async fn disk_limit_blocks_until_token_dropped() {
        let tmpdir = tempfile::tempdir().unwrap();
        let downloader =
            Arc::new(TraceDownloader::new(tmpdir.path().to_path_buf(), 100, None, None).unwrap());

        // Fill the disk limit with a 100-byte token (no blocking — limit not yet reached).
        let held = tmpdir.path().join("held");
        std::fs::write(&held, vec![0u8; 100]).unwrap();
        let token1 = downloader.make_test_token(held.clone(), 100);

        // A second file to acquire once space is freed.
        let waiting = tmpdir.path().join("waiting");
        std::fs::write(&waiting, vec![0u8; 10]).unwrap();

        // Spawn a task that should block until token1 is dropped.
        let d2 = Arc::clone(&downloader);
        let w2 = waiting.clone();
        let unblocked = Arc::new(AtomicBool::new(false));
        let unblocked_bg = Arc::clone(&unblocked);
        let handle = tokio::spawn(async move {
            let token = d2.wait_and_make_token(w2).await;
            unblocked_bg.store(true, Ordering::SeqCst);
            token
        });

        // Give the task time to reach the wait.
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        assert!(
            !unblocked.load(Ordering::SeqCst),
            "task should be blocked on disk limit"
        );

        // Drop token1 — held becomes unlocked and can be evicted by the
        // waiting task to make room.
        drop(token1);

        let token2 = handle.await.expect("background task panicked");
        assert!(
            unblocked.load(Ordering::SeqCst),
            "task should have unblocked after token drop"
        );

        // The waiting task evicted "held" to make room for "waiting".
        assert!(
            !held.exists(),
            "held file should be evicted when it is unlocked and space is needed"
        );
        assert!(
            waiting.exists(),
            "waiting file should still exist while token2 is held"
        );

        // Dropping token2 unlocks the file but does NOT delete it — it stays
        // in the cache for potential reuse.
        drop(token2);
        assert!(
            waiting.exists(),
            "waiting file should remain in the cache after token2 is dropped"
        );
    }

    /// A file already in the cache is served on repeat requests without
    /// issuing any request to the server.
    #[tokio::test]
    async fn cached_file_not_redownloaded_in_same_run() {
        let tmpdir = tempfile::tempdir().unwrap();
        let trace_path = "subdir/app.mock-trace";
        let get_count = Arc::new(AtomicU32::new(0));
        let base_url =
            start_mock_file_server(Bytes::from_static(b"trace body"), Arc::clone(&get_count)).await;

        let downloader = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();

        // First get: triggers a download.
        let t1 = downloader.get(&base_url, trace_path).await.unwrap();
        assert_eq!(
            get_count.load(Ordering::SeqCst),
            1,
            "first get should download"
        );
        drop(t1);

        // Second get: served from cache, no HEAD, no GET.
        let _t2 = downloader.get(&base_url, trace_path).await.unwrap();
        assert_eq!(
            get_count.load(Ordering::SeqCst),
            1,
            "second get in the same run should be served from cache without a new GET"
        );
    }

    /// A file enumerated from a previous run is served from cache without
    /// any network request.
    #[tokio::test]
    async fn cross_run_cache_reused_when_valid() {
        let tmpdir = tempfile::tempdir().unwrap();
        let trace_path = "subdir/app.mock-trace";
        let get_count = Arc::new(AtomicU32::new(0));
        let base_url =
            start_mock_file_server(Bytes::from_static(b"trace body"), Arc::clone(&get_count)).await;

        // First run: download the file.
        let d1 = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();
        let t1 = d1.get(&base_url, trace_path).await.unwrap();
        assert_eq!(get_count.load(Ordering::SeqCst), 1);
        drop(t1);
        drop(d1);

        // Second run: fresh downloader enumerates the cache.
        // Served directly from cache — no GET.
        let d2 = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();
        let _t2 = d2.get(&base_url, trace_path).await.unwrap();
        assert_eq!(
            get_count.load(Ordering::SeqCst),
            1,
            "valid cross-run cache should be served without re-downloading"
        );
    }

    /// A corrupted cached file is detected by `invalidate_if_corrupted` via
    /// ETag mismatch and evicted from the cache.
    #[tokio::test]
    async fn content_corrupted_cache_invalidated() {
        let tmpdir = tempfile::tempdir().unwrap();
        let trace_path = "subdir/app.mock-trace";
        let trace_body: &[u8] = b"trace body contents";
        let get_count = Arc::new(AtomicU32::new(0));
        let base_url =
            start_mock_file_server(Bytes::from_static(trace_body), Arc::clone(&get_count)).await;

        // First run: download the file.
        let d1 = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();
        let t1 = d1.get(&base_url, trace_path).await.unwrap();
        let cached_path = t1.path().to_path_buf();
        drop(t1);
        drop(d1);

        // Corrupt the content while keeping the same file size (flip a byte).
        let mut data = std::fs::read(&cached_path).unwrap();
        data[0] ^= 0xff;
        std::fs::write(&cached_path, &data).unwrap();

        // Second run: enumerate the corrupted file.  get() serves it from
        // cache without any network check.
        let d2 = TraceDownloader::new(tmpdir.path().to_path_buf(), 0, None, None).unwrap();
        let t2 = d2.get(&base_url, trace_path).await.unwrap();
        assert_eq!(
            get_count.load(Ordering::SeqCst),
            1,
            "corruption not yet detected: get() serves from cache without checking"
        );

        // After a replay failure the caller invokes invalidate_if_corrupted.
        // The ETag mismatch is detected and the entry is evicted.
        let invalidated = d2
            .invalidate_if_corrupted(&base_url, trace_path)
            .await
            .unwrap();
        assert!(invalidated, "ETag mismatch should trigger invalidation");

        // Drop the token from the failed replay; entry is already gone.
        drop(t2);

        // Next get() re-downloads the correct file.
        let t3 = d2.get(&base_url, trace_path).await.unwrap();
        assert_eq!(
            get_count.load(Ordering::SeqCst),
            2,
            "re-download after invalidation"
        );
        assert_eq!(
            std::fs::read(t3.path()).unwrap(),
            trace_body,
            "re-downloaded file should have correct content"
        );
    }
}