1use std::collections::BTreeMap;
33use std::path::Path;
34use std::sync::Arc;
35
36use anyhow::{Result, anyhow};
37use znippy_common::GUNNAR_REFS_MODULE;
38use znippy_common::arrow::array::{Array, StringArray, StringBuilder, UInt64Array, UInt64Builder};
39use znippy_common::arrow::datatypes::{DataType, Field, Schema};
40use znippy_common::arrow::record_batch::RecordBatch;
41
42use crate::pushlog::{PushLog, PushLogScan, read_sealed};
43
44pub use git_storage_trait::RefUpdate;
47
48#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct RefState {
51 pub target: Option<String>,
52 pub peeled: Option<String>,
53 pub symref_target: Option<String>,
54 pub push_seq: u64,
55 pub updated_ms: u64,
56}
57
58pub fn refs_schema() -> Arc<Schema> {
59 Arc::new(Schema::new(vec![
60 Field::new("name", DataType::Utf8, false),
61 Field::new("target", DataType::Utf8, true),
62 Field::new("peeled", DataType::Utf8, true),
63 Field::new("symref_target", DataType::Utf8, true),
64 Field::new("push_seq", DataType::UInt64, false),
65 Field::new("updated_ms", DataType::UInt64, false),
66 ]))
67}
68
69pub fn build_push_batch(updates: &[RefUpdate], push_seq: u64, updated_ms: u64) -> Result<RecordBatch> {
71 let n = updates.len();
72 let mut name = StringBuilder::with_capacity(n, n * 32);
73 let mut target = StringBuilder::with_capacity(n, n * 64);
74 let mut peeled = StringBuilder::with_capacity(n, n * 64);
75 let mut symref = StringBuilder::with_capacity(n, n * 32);
76 let mut seq = UInt64Builder::with_capacity(n);
77 let mut ms = UInt64Builder::with_capacity(n);
78
79 for u in updates {
80 name.append_value(&u.name);
81 match &u.target {
82 Some(t) => target.append_value(t),
83 None => target.append_null(),
84 }
85 match &u.peeled {
86 Some(t) => peeled.append_value(t),
87 None => peeled.append_null(),
88 }
89 match &u.symref_target {
90 Some(t) => symref.append_value(t),
91 None => symref.append_null(),
92 }
93 seq.append_value(push_seq);
94 ms.append_value(updated_ms);
95 }
96
97 RecordBatch::try_new(
98 refs_schema(),
99 vec![
100 Arc::new(name.finish()),
101 Arc::new(target.finish()),
102 Arc::new(peeled.finish()),
103 Arc::new(symref.finish()),
104 Arc::new(seq.finish()),
105 Arc::new(ms.finish()),
106 ],
107 )
108 .map_err(|e| anyhow!("refs push batch: {e}"))
109}
110
111pub struct RefLog {
113 log: PushLog,
114}
115
116impl RefLog {
117 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
118 Self { log: PushLog::new(path, refs_schema()) }
119 }
120
121 pub fn next_push_seq(&self) -> Result<u64> {
125 let scan = self.log.scan()?;
126 Ok(max_push_seq(&scan.pushes).map_or(0, |m| m + 1))
127 }
128
129 pub fn push(&self, updates: &[RefUpdate]) -> Result<u64> {
131 let seq = self.next_push_seq()?;
132 let ms = std::time::SystemTime::now()
133 .duration_since(std::time::UNIX_EPOCH)
134 .map(|d| d.as_millis() as u64)
135 .unwrap_or(0);
136 let batch = build_push_batch(updates, seq, ms)?;
137 self.log.append(&batch)?;
138 Ok(seq)
139 }
140
141 pub fn append_batch(&self, batch: &RecordBatch) -> Result<u64> {
147 self.log.append(batch)
148 }
149
150 pub fn scan(&self) -> Result<PushLogScan> {
151 self.log.scan()
152 }
153
154 pub fn compact(&self) -> Result<crate::pushlog::CompactionReport> {
157 self.log.compact()
158 }
159
160 pub fn maybe_compact(
162 &self,
163 policy: crate::pushlog::CompactionPolicy,
164 ) -> Result<Option<crate::pushlog::CompactionReport>> {
165 self.log.maybe_compact(policy)
166 }
167
168 pub fn current(&self) -> Result<BTreeMap<String, RefState>> {
170 Ok(fold(&self.log.scan()?.pushes)?)
171 }
172
173 pub fn seal_section(&self) -> Result<znippy_common::ReservedSection> {
175 self.log.seal_section(GUNNAR_REFS_MODULE)
176 }
177}
178
179fn max_push_seq(batches: &[RecordBatch]) -> Option<u64> {
180 let mut max = None;
181 for b in batches {
182 let seq = b.column_by_name("push_seq")?.as_any().downcast_ref::<UInt64Array>()?;
183 for i in 0..seq.len() {
184 max = Some(max.map_or(seq.value(i), |m: u64| m.max(seq.value(i))));
185 }
186 }
187 max
188}
189
190pub fn fold(batches: &[RecordBatch]) -> Result<BTreeMap<String, RefState>> {
196 let mut rows: Vec<(u64, usize, RefState, String)> = Vec::new();
197
198 for (bi, b) in batches.iter().enumerate() {
199 let name = col::<StringArray>(b, "name")?;
200 let target = col::<StringArray>(b, "target")?;
201 let peeled = col::<StringArray>(b, "peeled")?;
202 let symref = col::<StringArray>(b, "symref_target")?;
203 let seq = col::<UInt64Array>(b, "push_seq")?;
204 let ms = col::<UInt64Array>(b, "updated_ms")?;
205
206 for i in 0..b.num_rows() {
207 rows.push((
208 seq.value(i),
209 bi,
210 RefState {
211 target: (!target.is_null(i)).then(|| target.value(i).to_string()),
212 peeled: (!peeled.is_null(i)).then(|| peeled.value(i).to_string()),
213 symref_target: (!symref.is_null(i)).then(|| symref.value(i).to_string()),
214 push_seq: seq.value(i),
215 updated_ms: ms.value(i),
216 },
217 name.value(i).to_string(),
218 ));
219 }
220 }
221
222 rows.sort_by_key(|(seq, bi, _, _)| (*seq, *bi));
223
224 let mut out: BTreeMap<String, RefState> = BTreeMap::new();
225 for (_, _, state, name) in rows {
226 if state.target.is_none() && state.symref_target.is_none() {
227 out.remove(&name);
228 } else {
229 out.insert(name, state);
230 }
231 }
232 Ok(out)
233}
234
235pub fn read_refs(archive: &Path) -> Result<Option<BTreeMap<String, RefState>>> {
238 match read_sealed(archive, GUNNAR_REFS_MODULE)? {
239 Some(batches) => Ok(Some(fold(&batches)?)),
240 None => Ok(None),
241 }
242}
243
244fn col<'a, T: Array + 'static>(b: &'a RecordBatch, name: &str) -> Result<&'a T> {
245 b.column_by_name(name)
246 .ok_or_else(|| anyhow!("refs: no `{name}` column"))?
247 .as_any()
248 .downcast_ref::<T>()
249 .ok_or_else(|| anyhow!("refs: `{name}` has an unexpected type"))
250}
251
252#[cfg(test)]
253mod tests {
254 use super::*;
255 use crate::pushlog::truncate_for_test;
256
257 fn tmpdir(tag: &str) -> std::path::PathBuf {
258 let ns = std::time::SystemTime::now()
259 .duration_since(std::time::UNIX_EPOCH)
260 .unwrap()
261 .as_nanos();
262 let d = std::env::temp_dir().join(format!("znippy_refs_{tag}_{ns}"));
263 std::fs::create_dir_all(&d).unwrap();
264 d
265 }
266
267 fn oid(c: char) -> String {
268 std::iter::repeat_n(c, 40).collect()
269 }
270
271 #[test]
272 fn last_writer_wins_and_a_null_target_deletes() {
273 let dir = tmpdir("fold");
274 let log = RefLog::new(dir.join("refs.log"));
275 log.push(&[
276 RefUpdate::set("refs/heads/main", oid('a')),
277 RefUpdate::set("refs/heads/topic", oid('b')),
278 ])
279 .unwrap();
280 log.push(&[RefUpdate::set("refs/heads/main", oid('c'))]).unwrap();
281 log.push(&[RefUpdate::delete("refs/heads/topic")]).unwrap();
282
283 let refs = log.current().unwrap();
284 assert_eq!(refs["refs/heads/main"].target, Some(oid('c')), "second push must win");
285 assert!(!refs.contains_key("refs/heads/topic"), "a null target deletes the ref");
286 assert_eq!(refs.len(), 1);
287
288 std::fs::remove_dir_all(&dir).ok();
289 }
290
291 #[test]
296 fn a_crash_mid_push_leaves_no_partial_ref_update() {
297 let dir = tmpdir("atomic");
298 let path = dir.join("refs.log");
299 let log = RefLog::new(&path);
300 log.push(&[RefUpdate::set("refs/heads/main", oid('a'))]).unwrap();
301 let before = std::fs::metadata(&path).unwrap().len();
302
303 log.push(&[
305 RefUpdate::set("refs/heads/main", oid('9')),
306 RefUpdate::set("refs/heads/a", oid('1')),
307 RefUpdate::set("refs/heads/b", oid('2')),
308 ])
309 .unwrap();
310 let after = std::fs::metadata(&path).unwrap().len();
311 let intact = std::fs::read(&path).unwrap();
312
313 for cut in (before + 1)..after {
314 std::fs::write(&path, &intact).unwrap();
315 truncate_for_test(&path, cut).unwrap();
316 let refs = log.current().unwrap();
317 assert_eq!(
318 refs.len(),
319 1,
320 "cut at {cut}: a torn push must not publish ANY of its refs (got {refs:?})"
321 );
322 assert_eq!(
323 refs["refs/heads/main"].target,
324 Some(oid('a')),
325 "cut at {cut}: main must still be the pre-push value"
326 );
327 assert!(!refs.contains_key("refs/heads/a"), "cut at {cut}: leaked a partial ref");
328 assert!(!refs.contains_key("refs/heads/b"), "cut at {cut}: leaked a partial ref");
329 }
330
331 std::fs::write(&path, &intact).unwrap();
333 let refs = log.current().unwrap();
334 assert_eq!(refs.len(), 3, "the complete push publishes all three refs");
335 assert_eq!(refs["refs/heads/main"].target, Some(oid('9')));
336
337 std::fs::remove_dir_all(&dir).ok();
338 }
339
340 #[test]
343 fn push_seq_is_recovered_from_the_log_not_from_memory() {
344 let dir = tmpdir("seq");
345 let path = dir.join("refs.log");
346 let a = RefLog::new(&path);
347 assert_eq!(a.push(&[RefUpdate::set("refs/heads/main", oid('a'))]).unwrap(), 0);
348 assert_eq!(a.push(&[RefUpdate::set("refs/heads/main", oid('b'))]).unwrap(), 1);
349
350 let b = RefLog::new(&path);
352 assert_eq!(
353 b.next_push_seq().unwrap(),
354 2,
355 "a restarted writer must continue the sequence, not restart it"
356 );
357 assert_eq!(b.push(&[RefUpdate::set("refs/heads/main", oid('c'))]).unwrap(), 2);
358 assert_eq!(b.current().unwrap()["refs/heads/main"].target, Some(oid('c')));
359
360 std::fs::remove_dir_all(&dir).ok();
361 }
362
363 #[test]
367 fn ordering_is_by_push_seq_not_by_timestamp_or_position() {
368 let newer = build_push_batch(&[RefUpdate::set("refs/heads/main", oid('c'))], 7, 1000).unwrap();
370 let older = build_push_batch(&[RefUpdate::set("refs/heads/main", oid('a'))], 3, 9999).unwrap();
371 let refs = fold(&[newer, older]).unwrap();
372 assert_eq!(
373 refs["refs/heads/main"].target,
374 Some(oid('c')),
375 "push_seq 7 must beat push_seq 3 regardless of order or clock"
376 );
377 assert_eq!(refs["refs/heads/main"].push_seq, 7);
378 }
379
380 #[test]
381 fn symbolic_and_peeled_refs_round_trip() {
382 let dir = tmpdir("sym");
383 let log = RefLog::new(dir.join("refs.log"));
384 log.push(&[
385 RefUpdate::symbolic("HEAD", "refs/heads/main"),
386 RefUpdate::set("refs/tags/v1", oid('t')).with_peeled(oid('e')),
387 ])
388 .unwrap();
389 let refs = log.current().unwrap();
390 assert_eq!(refs["HEAD"].symref_target.as_deref(), Some("refs/heads/main"));
391 assert!(refs["HEAD"].target.is_none(), "a symref has no direct target");
392 assert_eq!(refs["refs/tags/v1"].peeled, Some(oid('e')));
393 std::fs::remove_dir_all(&dir).ok();
394 }
395}