This repository has no description
7.9 kB
236 lines
1//! `closure OUT`: the closure of the store paths on stdin, from the caches'
2//! narinfos only.
3
4use std::collections::{BTreeMap, BTreeSet, HashMap};
5use std::io::{self, BufRead, Write};
6use std::path::Path;
7use std::sync::atomic::{AtomicUsize, Ordering};
8use std::sync::Mutex;
9
10use serde::{Deserialize, Serialize};
11
12use crate::http::{self, Http};
13use crate::stats::Stats;
14use crate::STORE;
15
16/// Per store path, what the lazy layer needs to fetch it (OUT/paths.json).
17#[derive(Serialize, Deserialize, Clone, Debug)]
18#[serde(rename_all = "camelCase")]
19pub struct PathInfo {
20 pub cache: String,
21 pub url: String,
22 pub compression: String,
23 pub nar_hash: String,
24 pub nar_size: u64,
25 pub file_size: u64,
26}
27
28pub fn caches() -> Vec<String> {
29 std::env::var("LAZY_STORE_CACHES")
30 .unwrap_or_else(|_| "https://cache.nixos.org https://cache.lassul.us".into())
31 .split_whitespace()
32 .map(str::to_string)
33 .collect()
34}
35
36fn parse_narinfo(text: &[u8]) -> HashMap<String, String> {
37 String::from_utf8_lossy(text)
38 .lines()
39 .filter_map(|l| l.split_once(": "))
40 .map(|(k, v)| (k.to_string(), v.to_string()))
41 .collect()
42}
43
44fn references(i: &HashMap<String, String>) -> impl Iterator<Item = &str> {
45 i.get("References")
46 .map(String::as_str)
47 .unwrap_or("")
48 .split_whitespace()
49}
50
51/// The parsed narinfo of `name` from the first cache that has it, and its
52/// text.
53fn narinfo(
54 http: &Http,
55 caches: &[String],
56 name: &str,
57 stats: &Stats,
58) -> Option<(HashMap<String, String>, Vec<u8>)> {
59 let h = name.split('-').next().unwrap_or(name);
60 for c in caches {
61 match http.get(&format!("{c}/{h}.narinfo")) {
62 Ok((body, len)) => {
63 stats.add("narinfo", len as u64);
64 let mut i = parse_narinfo(&body);
65 i.insert("cache".into(), c.clone());
66 return Some((i, body));
67 }
68 Err(http::Error::NotFound) => {}
69 // still failing after the retries: try the next cache
70 Err(e) => {
71 eprintln!("giving up on {c}/{h}.narinfo: {e}");
72 stats.add("narinfo_error", 0);
73 }
74 }
75 }
76 None
77}
78
79/// Run `f` over `items` on `threads` threads, results in order.
80pub fn par_map<T: Sync, R: Send>(
81 items: &[T],
82 threads: usize,
83 f: impl Fn(&T) -> R + Sync,
84) -> Vec<R> {
85 let next = AtomicUsize::new(0);
86 let out: Mutex<Vec<Option<R>>> = Mutex::new((0..items.len()).map(|_| None).collect());
87 std::thread::scope(|s| {
88 for _ in 0..threads.min(items.len()) {
89 s.spawn(|| loop {
90 let i = next.fetch_add(1, Ordering::Relaxed);
91 if i >= items.len() {
92 break;
93 }
94 let r = f(&items[i]);
95 out.lock().unwrap()[i] = Some(r);
96 });
97 }
98 });
99 out.into_inner()
100 .unwrap()
101 .into_iter()
102 .map(Option::unwrap)
103 .collect()
104}
105
106/// Writes OUT/paths.json (per path: cache, NAR url, compression, hash),
107/// OUT/db.dump (input for `nix-store --load-db`), OUT/narinfo (a binary
108/// cache of just the narinfos) and OUT/closure.json.
109/// Paths that are in none of the caches, and everything referring to them,
110/// are left out: a path is registered only with its whole closure.
111pub fn closure(outdir: &Path) -> io::Result<()> {
112 let mut roots = BTreeSet::new();
113 for l in io::stdin().lock().lines() {
114 let l = l?;
115 let l = l.trim();
116 if !l.is_empty() {
117 roots.insert(l.strip_prefix(STORE).unwrap_or(l).to_string());
118 }
119 }
120 let caches = caches();
121 let http = Http::new(caches.iter().map(String::as_str));
122 let stats = Stats::new(None);
123 let mut infos: BTreeMap<String, HashMap<String, String>> = BTreeMap::new();
124 let mut texts: HashMap<String, Vec<u8>> = HashMap::new();
125 let mut missing: BTreeSet<String> = BTreeSet::new();
126 let mut todo: Vec<String> = roots.iter().cloned().collect();
127 while !todo.is_empty() {
128 let res = par_map(&todo, 32, |n| narinfo(&http, &caches, n, &stats));
129 let mut next = BTreeSet::new();
130 for (name, i) in todo.iter().zip(res) {
131 let Some((i, text)) = i else {
132 missing.insert(name.clone());
133 continue;
134 };
135 next.extend(references(&i).map(str::to_string));
136 infos.insert(name.clone(), i);
137 texts.insert(name.clone(), text);
138 }
139 todo = next
140 .into_iter()
141 .filter(|r| !infos.contains_key(r) && !missing.contains(r))
142 .collect();
143 }
144
145 let mut referrers: HashMap<String, Vec<String>> = HashMap::new();
146 for (n, i) in &infos {
147 for r in references(i).filter(|r| *r != n.as_str()) {
148 referrers.entry(r.to_string()).or_default().push(n.clone());
149 }
150 }
151 let mut incomplete = BTreeSet::new();
152 let mut todo: Vec<&String> = missing.iter().collect();
153 while let Some(m) = todo.pop() {
154 for r in referrers.get(m).into_iter().flatten() {
155 if incomplete.insert(r.clone()) {
156 todo.push(r);
157 }
158 }
159 }
160 infos.retain(|n, _| !incomplete.contains(n));
161
162 let mut paths = BTreeMap::new();
163 for (n, i) in &infos {
164 let get = |k: &str| {
165 i.get(k)
166 .cloned()
167 .ok_or_else(|| io::Error::other(format!("{n}: narinfo without {k}")))
168 };
169 paths.insert(
170 n.clone(),
171 PathInfo {
172 cache: get("cache")?,
173 url: get("URL")?,
174 compression: i
175 .get("Compression")
176 .cloned()
177 .unwrap_or_else(|| "none".into()),
178 nar_hash: get("NarHash")?,
179 nar_size: get("NarSize")?.parse().map_err(io::Error::other)?,
180 file_size: i.get("FileSize").and_then(|s| s.parse().ok()).unwrap_or(0),
181 },
182 );
183 }
184 std::fs::write(outdir.join("paths.json"), serde_json::to_vec(&paths)?)?;
185
186 let mut db = io::BufWriter::new(std::fs::File::create(outdir.join("db.dump"))?);
187 for (n, i) in &infos {
188 let refs: Vec<&str> = references(i).collect();
189 let d = i.get("Deriver").map(String::as_str).unwrap_or("");
190 let deriver = if d.ends_with(".drv") {
191 format!("{STORE}{d}")
192 } else {
193 String::new()
194 };
195 writeln!(
196 db,
197 "{STORE}{n}\n{}\n{}\n{deriver}\n{}",
198 i["NarHash"],
199 i["NarSize"],
200 refs.len()
201 )?;
202 for r in refs {
203 writeln!(db, "{STORE}{r}")?;
204 }
205 }
206 db.flush()?;
207
208 // the narinfos as a binary cache, for their signatures: `nix-store
209 // --load-db` registers none, `nix store copy-sigs` adds them
210 let dir = outdir.join("narinfo");
211 std::fs::create_dir_all(&dir)?;
212 std::fs::write(dir.join("nix-cache-info"), "StoreDir: /nix/store\n")?;
213 for n in infos.keys() {
214 let h = n.split('-').next().unwrap_or(n);
215 std::fs::write(dir.join(format!("{h}.narinfo")), &texts[n])?;
216 }
217
218 let mut by_cache: BTreeMap<String, u64> = BTreeMap::new();
219 for p in paths.values() {
220 *by_cache.entry(p.cache.clone()).or_default() += 1;
221 }
222 let mut summary = serde_json::json!({
223 "roots": roots.len(),
224 "paths": infos.len(),
225 "missing": missing,
226 "incomplete": incomplete.len(),
227 "by_cache": by_cache,
228 "nar_size": paths.values().map(|p| p.nar_size).sum::<u64>(),
229 "file_size": paths.values().map(|p| p.file_size).sum::<u64>(),
230 });
231 summary.as_object_mut().unwrap().extend(stats.counters());
232 let s = serde_json::to_string_pretty(&summary)?;
233 std::fs::write(outdir.join("closure.json"), &s)?;
234 println!("{s}");
235 Ok(())
236}