242 lines · 8 KBtext/plain; charset=utf-8
Onset: version control where nobody waits and nothing gets lost. The on command, its core and onset-server.
| 6 | use std::io::Write; |
| 7 | use std::path::{Path, PathBuf}; |
| 8 |
| 9 | use anyhow::{Context, Result, anyhow}; |
| 10 | use vcs_core::bundle::{self, Frame, RAW, ZSTD}; |
| 11 | use vcs_core::object::{Kind, Object}; |
| 12 | use vcs_core::{Id, Source, chunk}; |
| 13 |
| 14 | /// Objects at least this big are compressed on disk and on the wire. |
| 15 | const COMPRESS_OVER: usize = 512; |
| 16 |
| 17 | #[derive(Clone)] |
| 18 | pub struct Store { |
| 19 | root: PathBuf, |
| 20 | } |
| 21 |
| 22 | impl Store { |
| 23 | pub fn open(root: &Path) -> Result<Store> { |
| 24 | fs::create_dir_all(root).with_context(|| format!("creating {}", root.display()))?; |
| 25 | Ok(Store { root: root.to_path_buf() }) |
| 26 | } |
| 27 |
| 28 | fn path(&self, id: &Id) -> PathBuf { |
| 29 | let hex = id.hex(); |
| 30 | self.root.join(&hex[..2]).join(&hex[2..]) |
| 31 | } |
| 32 |
| 33 | pub fn has(&self, id: &Id) -> bool { |
| 34 | self.path(id).is_file() |
| 35 | } |
| 36 |
| 37 | /// How many bytes the object takes here (0 when it's missing). |
| 38 | pub fn stored_len(&self, id: &Id) -> u64 { |
| 39 | fs::metadata(self.path(id)).map(|m| m.len()).unwrap_or(0) |
| 40 | } |
| 41 |
| 42 | /// Stores an object (no-op if present). Returns its ID. |
| 43 | pub fn put(&self, object: &Object) -> Result<Id> { |
| 44 | let id = object.id(); |
| 45 | if self.has(&id) { |
| 46 | return Ok(id); |
| 47 | } |
| 48 | self.write_frame(&encode(object, id))?; |
| 49 | Ok(id) |
| 50 | } |
| 51 |
| 52 | fn write_frame(&self, frame: &Frame) -> Result<()> { |
| 53 | let path = self.path(&frame.id); |
| 54 | if path.is_file() { |
| 55 | return Ok(()); |
| 56 | } |
| 57 | let dir = path.parent().unwrap(); |
| 58 | fs::create_dir_all(dir)?; |
| 59 | let mut bytes = Vec::with_capacity(frame.payload.len() + 10); |
| 60 | bytes.push(frame.kind as u8); |
| 61 | bytes.push(frame.codec); |
| 62 | bytes.extend_from_slice(&frame.raw_len.to_le_bytes()); |
| 63 | bytes.extend_from_slice(&frame.payload); |
| 64 | let tmp = dir.join(format!(".tmp-{}-{}", std::process::id(), rand::random::<u32>())); |
| 65 | { |
| 66 | let mut f = fs::File::create(&tmp)?; |
| 67 | f.write_all(&bytes)?; |
| 68 | } |
| 69 | fs::rename(&tmp, &path)?; |
| 70 | Ok(()) |
| 71 | } |
| 72 |
| 73 | /// The stored frame (possibly compressed), for bundles. |
| 74 | pub fn frame(&self, id: &Id) -> Result<Option<Frame>> { |
| 75 | let bytes = match fs::read(self.path(id)) { |
| 76 | Ok(b) => b, |
| 77 | Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), |
| 78 | Err(e) => return Err(e.into()), |
| 79 | }; |
| 80 | if bytes.len() < 10 { |
| 81 | return Err(anyhow!("object {} is damaged", id.short())); |
| 82 | } |
| 83 | let kind = Kind::from_u8(bytes[0]).ok_or_else(|| anyhow!("object {} is damaged", id.short()))?; |
| 84 | let codec = bytes[1]; |
| 85 | let raw_len = u64::from_le_bytes(bytes[2..10].try_into().unwrap()); |
| 86 | Ok(Some(Frame { id: *id, kind, codec, raw_len, payload: bytes[10..].to_vec() })) |
| 87 | } |
| 88 |
| 89 | pub fn read(&self, id: &Id) -> Result<Option<Object>> { |
| 90 | match self.frame(id)? { |
| 91 | Some(frame) => Ok(Some(decode(&frame)?)), |
| 92 | None => Ok(None), |
| 93 | } |
| 94 | } |
| 95 |
| 96 | /// Reads every object kept here and checks it against its ID, in |
| 97 | /// parallel. Answers how many there are and the IDs of damaged ones. |
| 98 | pub fn verify_all(&self) -> Result<(usize, Vec<Id>)> { |
| 99 | use rayon::prelude::*; |
| 100 | let mut ids = Vec::new(); |
| 101 | for dir in fs::read_dir(&self.root)? { |
| 102 | let dir = dir?; |
| 103 | let Some(prefix) = dir.file_name().to_str().filter(|p| p.len() == 2).map(str::to_string) else { continue }; |
| 104 | if !dir.file_type()?.is_dir() { |
| 105 | continue; |
| 106 | } |
| 107 | for f in fs::read_dir(dir.path())? { |
| 108 | let name = f?.file_name(); |
| 109 | if let Some(id) = name.to_str().and_then(|rest| Id::parse(&format!("{prefix}{rest}"))) { |
| 110 | ids.push(id); |
| 111 | } |
| 112 | } |
| 113 | } |
| 114 | let damaged: Vec<Id> = ids |
| 115 | .par_iter() |
| 116 | .filter(|id| match self.frame(id) { |
| 117 | Ok(Some(frame)) => decode(&frame).map_or(true, |o| bundle::verify(&frame, &o.data).is_err()), |
| 118 | _ => true, |
| 119 | }) |
| 120 | .copied() |
| 121 | .collect(); |
| 122 | Ok((ids.len(), damaged)) |
| 123 | } |
| 124 |
| 125 | /// Stores a file's content (chunked when big) and returns its ID. |
| 126 | pub fn put_file(&self, content: &[u8]) -> Result<Id> { |
| 127 | let objects = chunk::file_objects(content); |
| 128 | let id = objects[0].id(); |
| 129 | if self.has(&id) { |
| 130 | return Ok(id); |
| 131 | } |
| 132 | // Chunks first, so the list never points at missing chunks. |
| 133 | for o in objects.iter().skip(1) { |
| 134 | self.put(o)?; |
| 135 | } |
| 136 | self.put(&objects[0])?; |
| 137 | Ok(id) |
| 138 | } |
| 139 |
| 140 | pub fn read_file(&self, id: &Id) -> Result<Vec<u8>> { |
| 141 | Ok(chunk::read_file(self, id)?) |
| 142 | } |
| 143 |
| 144 | /// Every object a file ID needs (itself plus chunks). |
| 145 | pub fn file_parts(&self, id: &Id) -> Result<Vec<Id>> { |
| 146 | let mut out = vec![*id]; |
| 147 | // The first byte says whether it's a chunk list; blobs aren't read. |
| 148 | let mut kind = [0u8; 1]; |
| 149 | match fs::File::open(self.path(id)).and_then(|mut f| std::io::Read::read_exact(&mut f, &mut kind)) { |
| 150 | Ok(()) if kind[0] == Kind::Chunks as u8 => {} |
| 151 | Ok(()) => return Ok(out), |
| 152 | Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(out), |
| 153 | Err(e) => return Err(e.into()), |
| 154 | } |
| 155 | if let Some(frame) = self.frame(id)? |
| 156 | && frame.kind == Kind::Chunks |
| 157 | { |
| 158 | let o = decode(&frame)?; |
| 159 | out.extend(vcs_core::object::decode_chunks(&o.data)?.into_iter().map(|(c, _)| c)); |
| 160 | } |
| 161 | Ok(out) |
| 162 | } |
| 163 |
| 164 | /// A bundle of the given objects (those present). |
| 165 | pub fn bundle(&self, ids: &[Id]) -> Result<Vec<u8>> { |
| 166 | let mut out = Vec::new(); |
| 167 | bundle::write_header(&mut out); |
| 168 | for id in ids { |
| 169 | if let Some(frame) = self.frame(id)? { |
| 170 | bundle::write_frame(&mut out, &frame); |
| 171 | } |
| 172 | } |
| 173 | bundle::write_end(&mut out); |
| 174 | Ok(out) |
| 175 | } |
| 176 |
| 177 | /// Stores every object in a bundle after checking it. Returns how many. |
| 178 | pub fn ingest(&self, bytes: &[u8]) -> Result<usize> { |
| 179 | let frames = bundle::read(bytes)?; |
| 180 | let n = frames.len(); |
| 181 | for frame in frames { |
| 182 | if self.has(&frame.id) { |
| 183 | continue; |
| 184 | } |
| 185 | let object = decode(&frame)?; |
| 186 | bundle::verify(&frame, &object.data)?; |
| 187 | self.write_frame(&frame)?; |
| 188 | } |
| 189 | Ok(n) |
| 190 | } |
| 191 | } |
| 192 |
| 193 | impl Source for Store { |
| 194 | fn get(&self, id: &Id) -> vcs_core::Result<Option<Object>> { |
| 195 | self.read(id).map_err(|e| vcs_core::Error::Corrupt(format!("{e:#}"))) |
| 196 | } |
| 197 | } |
| 198 |
| 199 | pub fn encode(object: &Object, id: Id) -> Frame { |
| 200 | if object.data.len() >= COMPRESS_OVER |
| 201 | && let Ok(packed) = zstd::bulk::compress(&object.data, 3) |
| 202 | && packed.len() < object.data.len() |
| 203 | { |
| 204 | return Frame { id, kind: object.kind, codec: ZSTD, raw_len: object.data.len() as u64, payload: packed }; |
| 205 | } |
| 206 | Frame { id, kind: object.kind, codec: RAW, raw_len: object.data.len() as u64, payload: object.data.clone() } |
| 207 | } |
| 208 |
| 209 | pub fn decode(frame: &Frame) -> Result<Object> { |
| 210 | let data = match frame.codec { |
| 211 | RAW => frame.payload.clone(), |
| 212 | ZSTD => { |
| 213 | // A frame's claimed size isn't trusted: reading stops just past it. |
| 214 | use std::io::Read; |
| 215 | let mut out = Vec::with_capacity((frame.raw_len as usize).min(frame.payload.len().saturating_mul(4)).min(1 << 20)); |
| 216 | zstd::stream::read::Decoder::new(&frame.payload[..])?.take(frame.raw_len + 1).read_to_end(&mut out).context("damaged object")?; |
| 217 | out |
| 218 | } |
| 219 | _ => return Err(anyhow!("unknown codec")), |
| 220 | }; |
| 221 | Ok(Object { kind: frame.kind, data }) |
| 222 | } |
| 223 |
| 224 | #[cfg(test)] |
| 225 | mod tests { |
| 226 | use super::*; |
| 227 |
| 228 | #[test] |
| 229 | fn put_read_bundle() { |
| 230 | let dir = tempfile::tempdir().unwrap(); |
| 231 | let store = Store::open(dir.path()).unwrap(); |
| 232 | let big = "x".repeat(10_000).into_bytes(); |
| 233 | let id = store.put_file(&big).unwrap(); |
| 234 | assert_eq!(store.read_file(&id).unwrap(), big); |
| 235 | let bundle = store.bundle(&[id]).unwrap(); |
| 236 | let other = tempfile::tempdir().unwrap(); |
| 237 | let store2 = Store::open(other.path()).unwrap(); |
| 238 | assert_eq!(store2.ingest(&bundle).unwrap(), 1); |
| 239 | assert_eq!(store2.read_file(&id).unwrap(), big); |
| 240 | } |
| 241 | } |
| 242 |