path: root/users/edef/fetchroots/src/main.rs

//! Fetch all[^1] GC roots from releases.nixos.org into a `roots.parquet` file.
//!
//! The resulting Parquet has three columns:
//!
//!  * `key` (`String`): the release, eg `nixos/22.11-small/nixos-22.11.513.563dc6476b8`
//!  * `timestamp` (`DateTime`): the timestamp of the GC roots file for this release
//!  * `store_path_hash` (`List[Binary]`): hash part of the store paths rooted by this release
//!
//! [^1]: some roots are truly ancient, and aren't compatible with Nix 1.x

use anyhow::Result;
use std::{
    collections::BTreeMap,
    fs::File,
    io::{BufRead, Read},
    sync::Arc,
    time::SystemTime,
};

use aws_config::Region;
use aws_sdk_s3::operation::get_object::builders::GetObjectFluentBuilder;
use bytes::{Buf, Bytes};
use bytes_utils::SegmentedBuf;
use chrono::{DateTime, Utc};
use nix_compat::nixbase32;
use polars::prelude::*;
use tokio::{
    sync::Semaphore,
    task::{block_in_place, JoinSet},
};

#[derive(Debug)]
struct Meta {
    format: Format,
    e_tag: String,
    last_modified: DateTime<Utc>,
}

#[tokio::main]
async fn main() {
    let sdk_config = aws_config::load_defaults(aws_config::BehaviorVersion::v2023_11_09())
        .await
        .into_builder()
        .region(Region::from_static("eu-west-1"))
        .build();

    let s3 = aws_sdk_s3::Client::new(&sdk_config);

    let mut keys: BTreeMap<String, Meta> = {
        let pages = s3
            .list_objects_v2()
            .bucket("nix-releases")
            .into_paginator()
            .send()
            .try_collect()
            .await
            .unwrap();

        let objects = pages.into_iter().flat_map(|page| {
            assert_eq!(page.prefix().unwrap_or_default(), "");
            assert!(page.common_prefixes.is_none());
            page.contents.unwrap_or_default()
        });

        let mut prev_key = String::new();
        objects
            .filter_map(|obj| {
                let key = obj.key().unwrap();

                assert!(&*prev_key < key);
                key.clone_into(&mut prev_key);

                let (key, tail) = key.rsplit_once('/')?;
                // Our preference order happens to match lexicographical order,
                // and listings are returned in lexicographical order.
                let format = match tail {
                    "MANIFEST" => Format::Manifest,
                    "MANIFEST.bz2" => Format::ManifestBz,
                    "store-paths.xz" => Format::StorePathsXz,
                    _ => return None,
                };

                Some((
                    key.to_owned(),
                    Meta {
                        format,
                        e_tag: obj.e_tag.unwrap(),
                        last_modified: SystemTime::try_from(obj.last_modified.unwrap())
                            .unwrap()
                            .into(),
                    },
                ))
            })
            .collect()
    };

    // These releases are so old they don't even use nixbase32 store paths.
    for key in [
        "nix/nix-0.6",
        "nix/nix-0.6.1",
        "nix/nix-0.7",
        "nix/nix-0.8",
        "nixpkgs/nixpkgs-0.5",
        "nixpkgs/nixpkgs-0.5.1",
        "nixpkgs/nixpkgs-0.6",
        "nixpkgs/nixpkgs-0.7",
        "nixpkgs/nixpkgs-0.8",
        "nixpkgs/nixpkgs-0.9",
        "nixpkgs/nixpkgs-0.10",
        "nixpkgs/nixpkgs-0.11",
    ] {
        assert!(keys.remove(key).is_some());
    }

    let mut js = JoinSet::new();
    let sem = Arc::new(Semaphore::new(16));

    let bar = indicatif::ProgressBar::new(keys.len() as u64);
    for (root, meta) in keys {
        let sem = sem.clone();
        let s3 = s3.clone();

        js.spawn(async move {
            let _permit = sem.acquire().await.unwrap();

            let body = get_object(
                s3.get_object()
                    .bucket("nix-releases")
                    .key(format!("{root}/{}", meta.format.as_str()))
                    .if_match(meta.e_tag),
            )
            .await
            .unwrap()
            .reader();

            let ph_array = block_in_place(|| meta.format.to_ph_array(body).rechunk());
            df! {
                "key" => [root],
                "timestamp" => [meta.last_modified.naive_utc()],
                "store_path_hash" => ph_array.into_series().implode().unwrap()
            }
            .unwrap()
        });
    }

    let mut writer = ParquetWriter::new(File::create("roots.parquet").unwrap())
        .batched(&Schema::from_iter([
            Field::new("key", DataType::String),
            Field::new(
                "timestamp",
                DataType::Datetime(TimeUnit::Milliseconds, None),
            ),
            Field::new(
                "store_path_hash",
                DataType::List(Box::new(DataType::Binary)),
            ),
        ]))
        .unwrap();

    while let Some(df) = js.join_next().await.transpose().unwrap() {
        block_in_place(|| writer.write_batch(&df)).unwrap();
        bar.inc(1);
    }

    writer.finish().unwrap();
}

#[derive(Debug)]
enum Format {
    Manifest,
    ManifestBz,
    StorePathsXz,
}

impl Format {
    fn as_str(&self) -> &'static str {
        match self {
            Format::Manifest => "MANIFEST",
            Format::ManifestBz => "MANIFEST.bz2",
            Format::StorePathsXz => "store-paths.xz",
        }
    }

    fn to_ph_array(&self, mut body: impl BufRead) -> BinaryChunked {
        match self {
            Format::Manifest | Format::ManifestBz => {
                let mut buf = String::new();
                match self {
                    Format::Manifest => {
                        body.read_to_string(&mut buf).unwrap();
                    }
                    Format::ManifestBz => {
                        bzip2::bufread::BzDecoder::new(body)
                            .read_to_string(&mut buf)
                            .unwrap();
                    }
                    _ => unreachable!(),
                }

                let buf = buf
                    .strip_prefix("version {\n  ManifestVersion: 3\n}\n")
                    .unwrap();

                BinaryChunked::from_iter_values(
                    "store_path_hash",
                    buf.split_terminator("}\n").map(|chunk| -> [u8; 20] {
                        let chunk = chunk.strip_prefix("patch ").unwrap_or(chunk);
                        let line = chunk.strip_prefix("{\n  StorePath: /nix/store/").unwrap();
                        nixbase32::decode_fixed(&line[..32]).unwrap()
                    }),
                )
            }
            Format::StorePathsXz => {
                let mut buf = String::new();
                xz2::bufread::XzDecoder::new(body)
                    .read_to_string(&mut buf)
                    .unwrap();

                BinaryChunked::from_iter_values(
                    "store_path_hash",
                    buf.split_terminator('\n').map(|line| -> [u8; 20] {
                        let line = line.strip_prefix("/nix/store/").unwrap();
                        nixbase32::decode_fixed(&line[..32]).unwrap()
                    }),
                )
            }
        }
    }
}

async fn get_object(request: GetObjectFluentBuilder) -> Result<SegmentedBuf<Bytes>> {
    // if we don't constrain the ETag, we might experience read skew
    assert!(request.get_if_match().is_some(), "if_match must be set");

    let mut buf: SegmentedBuf<Bytes> = SegmentedBuf::new();
    let mut resp = request.clone().send().await?;
    let content_length: usize = resp.content_length.unwrap().try_into().unwrap();

    loop {
        while let Ok(Some(chunk)) = resp.body.try_next().await {
            buf.push(chunk);
        }

        if buf.remaining() >= content_length {
            assert_eq!(buf.remaining(), content_length, "got excess bytes");
            break Ok(buf);
        }

        resp = request
            .clone()
            .range(format!("bytes={}-", buf.remaining()))
            .send()
            .await?;

        assert_ne!(resp.content_range, None);
    }
}