202 lines · 7 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 | /// Stores an object (no-op if present). Returns its ID. |
| 38 | pub fn put(&self, object: &Object) -> Result<Id> { |
| 39 | let id = object.id(); |
| 40 | if self.has(&id) { |
| 41 | return Ok(id); |
| 42 | } |
| 43 | self.write_frame(&encode(object, id))?; |
| 44 | Ok(id) |
| 45 | } |
| 46 |
| 47 | fn write_frame(&self, frame: &Frame) -> Result<()> { |
| 48 | let path = self.path(&frame.id); |
| 49 | if path.is_file() { |
| 50 | return Ok(()); |
| 51 | } |
| 52 | let dir = path.parent().unwrap(); |
| 53 | fs::create_dir_all(dir)?; |
| 54 | let mut bytes = Vec::with_capacity(frame.payload.len() + 10); |
| 55 | bytes.push(frame.kind as u8); |
| 56 | bytes.push(frame.codec); |
| 57 | bytes.extend_from_slice(&frame.raw_len.to_le_bytes()); |
| 58 | bytes.extend_from_slice(&frame.payload); |
| 59 | let tmp = dir.join(format!(".tmp-{}-{}", std::process::id(), rand::random::<u32>())); |
| 60 | { |
| 61 | let mut f = fs::File::create(&tmp)?; |
| 62 | f.write_all(&bytes)?; |
| 63 | } |
| 64 | fs::rename(&tmp, &path)?; |
| 65 | Ok(()) |
| 66 | } |
| 67 |
| 68 | /// The stored frame (possibly compressed), for bundles. |
| 69 | pub fn frame(&self, id: &Id) -> Result<Option<Frame>> { |
| 70 | let bytes = match fs::read(self.path(id)) { |
| 71 | Ok(b) => b, |
| 72 | Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), |
| 73 | Err(e) => return Err(e.into()), |
| 74 | }; |
| 75 | if bytes.len() < 10 { |
| 76 | return Err(anyhow!("object {} is damaged", id.short())); |
| 77 | } |
| 78 | let kind = Kind::from_u8(bytes[0]).ok_or_else(|| anyhow!("object {} is damaged", id.short()))?; |
| 79 | let codec = bytes[1]; |
| 80 | let raw_len = u64::from_le_bytes(bytes[2..10].try_into().unwrap()); |
| 81 | Ok(Some(Frame { id: *id, kind, codec, raw_len, payload: bytes[10..].to_vec() })) |
| 82 | } |
| 83 |
| 84 | pub fn read(&self, id: &Id) -> Result<Option<Object>> { |
| 85 | match self.frame(id)? { |
| 86 | Some(frame) => Ok(Some(decode(&frame)?)), |
| 87 | None => Ok(None), |
| 88 | } |
| 89 | } |
| 90 |
| 91 | /// Stores a file's content (chunked when big) and returns its ID. |
| 92 | pub fn put_file(&self, content: &[u8]) -> Result<Id> { |
| 93 | let objects = chunk::file_objects(content); |
| 94 | let id = objects[0].id(); |
| 95 | if self.has(&id) { |
| 96 | return Ok(id); |
| 97 | } |
| 98 | // Chunks first, so the list never points at missing chunks. |
| 99 | for o in objects.iter().skip(1) { |
| 100 | self.put(o)?; |
| 101 | } |
| 102 | self.put(&objects[0])?; |
| 103 | Ok(id) |
| 104 | } |
| 105 |
| 106 | pub fn read_file(&self, id: &Id) -> Result<Vec<u8>> { |
| 107 | Ok(chunk::read_file(self, id)?) |
| 108 | } |
| 109 |
| 110 | /// Every object a file ID needs (itself plus chunks). |
| 111 | pub fn file_parts(&self, id: &Id) -> Result<Vec<Id>> { |
| 112 | let mut out = vec![*id]; |
| 113 | // The first byte says whether it's a chunk list; blobs aren't read. |
| 114 | let mut kind = [0u8; 1]; |
| 115 | match fs::File::open(self.path(id)).and_then(|mut f| std::io::Read::read_exact(&mut f, &mut kind)) { |
| 116 | Ok(()) if kind[0] == Kind::Chunks as u8 => {} |
| 117 | Ok(()) => return Ok(out), |
| 118 | Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(out), |
| 119 | Err(e) => return Err(e.into()), |
| 120 | } |
| 121 | if let Some(frame) = self.frame(id)? |
| 122 | && frame.kind == Kind::Chunks |
| 123 | { |
| 124 | let o = decode(&frame)?; |
| 125 | out.extend(vcs_core::object::decode_chunks(&o.data)?.into_iter().map(|(c, _)| c)); |
| 126 | } |
| 127 | Ok(out) |
| 128 | } |
| 129 |
| 130 | /// A bundle of the given objects (those present). |
| 131 | pub fn bundle(&self, ids: &[Id]) -> Result<Vec<u8>> { |
| 132 | let mut out = Vec::new(); |
| 133 | bundle::write_header(&mut out); |
| 134 | for id in ids { |
| 135 | if let Some(frame) = self.frame(id)? { |
| 136 | bundle::write_frame(&mut out, &frame); |
| 137 | } |
| 138 | } |
| 139 | bundle::write_end(&mut out); |
| 140 | Ok(out) |
| 141 | } |
| 142 |
| 143 | /// Stores every object in a bundle after checking it. Returns how many. |
| 144 | pub fn ingest(&self, bytes: &[u8]) -> Result<usize> { |
| 145 | let frames = bundle::read(bytes)?; |
| 146 | let n = frames.len(); |
| 147 | for frame in frames { |
| 148 | if self.has(&frame.id) { |
| 149 | continue; |
| 150 | } |
| 151 | let object = decode(&frame)?; |
| 152 | bundle::verify(&frame, &object.data)?; |
| 153 | self.write_frame(&frame)?; |
| 154 | } |
| 155 | Ok(n) |
| 156 | } |
| 157 | } |
| 158 |
| 159 | impl Source for Store { |
| 160 | fn get(&self, id: &Id) -> vcs_core::Result<Option<Object>> { |
| 161 | self.read(id).map_err(|e| vcs_core::Error::Corrupt(format!("{e:#}"))) |
| 162 | } |
| 163 | } |
| 164 |
| 165 | pub fn encode(object: &Object, id: Id) -> Frame { |
| 166 | if object.data.len() >= COMPRESS_OVER |
| 167 | && let Ok(packed) = zstd::bulk::compress(&object.data, 3) |
| 168 | && packed.len() < object.data.len() |
| 169 | { |
| 170 | return Frame { id, kind: object.kind, codec: ZSTD, raw_len: object.data.len() as u64, payload: packed }; |
| 171 | } |
| 172 | Frame { id, kind: object.kind, codec: RAW, raw_len: object.data.len() as u64, payload: object.data.clone() } |
| 173 | } |
| 174 |
| 175 | pub fn decode(frame: &Frame) -> Result<Object> { |
| 176 | let data = match frame.codec { |
| 177 | RAW => frame.payload.clone(), |
| 178 | ZSTD => zstd::bulk::decompress(&frame.payload, frame.raw_len as usize).context("damaged object")?, |
| 179 | _ => return Err(anyhow!("unknown codec")), |
| 180 | }; |
| 181 | Ok(Object { kind: frame.kind, data }) |
| 182 | } |
| 183 |
| 184 | #[cfg(test)] |
| 185 | mod tests { |
| 186 | use super::*; |
| 187 |
| 188 | #[test] |
| 189 | fn put_read_bundle() { |
| 190 | let dir = tempfile::tempdir().unwrap(); |
| 191 | let store = Store::open(dir.path()).unwrap(); |
| 192 | let big = "x".repeat(10_000).into_bytes(); |
| 193 | let id = store.put_file(&big).unwrap(); |
| 194 | assert_eq!(store.read_file(&id).unwrap(), big); |
| 195 | let bundle = store.bundle(&[id]).unwrap(); |
| 196 | let other = tempfile::tempdir().unwrap(); |
| 197 | let store2 = Store::open(other.path()).unwrap(); |
| 198 | assert_eq!(store2.ingest(&bundle).unwrap(), 1); |
| 199 | assert_eq!(store2.read_file(&id).unwrap(), big); |
| 200 | } |
| 201 | } |
| 202 |