lfsx_server/storage/
mod.rs

1mod codec;
2mod dedupe;
3mod rewrite;
4mod staging;
5mod sweep;
6
7#[cfg(test)]
8mod tests;
9
10use std::collections::HashMap;
11use std::path::{Path, PathBuf};
12use std::sync::atomic::{AtomicU64, Ordering};
13use std::time::Instant;
14
15use tokio::sync::Mutex;
16
17use futures_util::{Stream, StreamExt};
18use sha2::{Digest, Sha256};
19use tokio::fs;
20use tokio::io::AsyncWriteExt;
21
22use crate::error::Error;
23use crate::namespace::Namespace;
24
25pub use dedupe::DedupeReport;
26use dedupe::shares_bytes_with;
27pub use rewrite::CompressReport;
28
29enum Sink {
30    Raw(fs::File),
31    Framed(Box<codec::Writer>),
32}
33
34impl Sink {
35    async fn write(&mut self, chunk: &[u8]) -> Result<(), Error> {
36        match self {
37            Self::Raw(file) => Ok(file.write_all(chunk).await?),
38            Self::Framed(writer) => writer.push(chunk).await,
39        }
40    }
41
42    async fn finish(self) -> Result<(), Error> {
43        match self {
44            Self::Raw(mut file) => {
45                file.flush().await?;
46                Ok(file.sync_all().await?)
47            }
48            Self::Framed(writer) => writer.finish().await,
49        }
50    }
51}
52
53// What a download reads from, whether or not the bytes on disk are the object.
54pub enum Object {
55    Raw { file: fs::File, size: u64 },
56    Framed(codec::Framed),
57}
58
59impl Object {
60    pub fn size(&self) -> u64 {
61        match self {
62            Self::Raw { size, .. } => *size,
63            Self::Framed(framed) => framed.plaintext(),
64        }
65    }
66
67    pub async fn stream(
68        self,
69        start: u64,
70        length: u64,
71    ) -> Result<futures_util::stream::BoxStream<'static, Result<axum::body::Bytes, Error>>, Error>
72    {
73        use futures_util::StreamExt;
74        use tokio::io::AsyncSeekExt;
75
76        match self {
77            Self::Raw { mut file, .. } => {
78                file.seek(std::io::SeekFrom::Start(start)).await?;
79                let reader =
80                    tokio_util::io::ReaderStream::new(tokio::io::AsyncReadExt::take(file, length));
81
82                Ok(reader.map(|chunk| chunk.map_err(Error::from)).boxed())
83            }
84            Self::Framed(framed) => Ok(framed.stream(start, length).boxed()),
85        }
86    }
87}
88pub use staging::{Reclaimed, reclaim};
89pub use sweep::SweepReport;
90
91// What is left of a repository's budget for one transfer. It travels with the
92// upload because a client that skips negotiation may also skip declaring a
93// size, and a budget checked once against a number the client chose is not a
94// budget.
95#[derive(Debug, Clone, Copy)]
96pub struct Budget {
97    pub used: u64,
98    pub limit: u64,
99}
100
101impl Budget {
102    pub fn exceeded_by(&self, arriving: u64) -> bool {
103        self.used + arriving > self.limit
104    }
105
106    pub fn refusal(&self) -> Error {
107        Error::OverQuota {
108            used: self.used,
109            limit: self.limit,
110        }
111    }
112}
113
114pub struct LocalStore {
115    root: PathBuf,
116    counter: AtomicU64,
117    usage: Mutex<Option<(Instant, u64, u64)>>,
118    per_namespace: Mutex<HashMap<String, (Instant, u64, u64)>>,
119    scans: AtomicU64,
120    max_object_size: Option<u64>,
121    compression: Option<i32>,
122}
123
124impl LocalStore {
125    pub fn new(root: impl Into<PathBuf>) -> Self {
126        Self {
127            root: root.into(),
128            counter: AtomicU64::new(0),
129            usage: Mutex::new(None),
130            per_namespace: Mutex::new(HashMap::new()),
131            scans: AtomicU64::new(0),
132            max_object_size: None,
133            compression: None,
134        }
135    }
136
137    pub fn with_compression(mut self, level: Option<i32>) -> Self {
138        self.compression = level;
139        self
140    }
141
142    pub fn with_max_object_size(mut self, limit: Option<u64>) -> Self {
143        self.max_object_size = limit;
144        self
145    }
146
147    pub fn validate_oid(oid: &str) -> Result<(), Error> {
148        let well_formed = oid.len() == 64
149            && oid
150                .bytes()
151                .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
152
153        well_formed.then_some(()).ok_or(Error::MalformedOid)
154    }
155
156    fn object_path(&self, ns: &Namespace, oid: &str) -> PathBuf {
157        self.root
158            .join(ns.org())
159            .join(ns.repo())
160            .join(&oid[0..2])
161            .join(&oid[2..4])
162            .join(oid)
163    }
164
165    fn content_path(&self, oid: &str) -> PathBuf {
166        self.root
167            .join(".content")
168            .join(&oid[0..2])
169            .join(&oid[2..4])
170            .join(oid)
171    }
172
173    pub fn scans(&self) -> u64 {
174        self.scans.load(Ordering::Relaxed)
175    }
176
177    pub async fn writable(&self) -> Result<(), Error> {
178        fs::create_dir_all(&self.root).await?;
179
180        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
181        let probe = self.root.join(format!(".readiness.{ticket}"));
182
183        fs::write(&probe, b"").await?;
184        fs::remove_file(&probe).await?;
185
186        Ok(())
187    }
188
189    pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
190        Self::validate_oid(oid).is_ok() && fs::metadata(self.object_path(ns, oid)).await.is_ok()
191    }
192
193    pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<Object, Error> {
194        Self::validate_oid(oid)?;
195        let path = self.object_path(ns, oid);
196        let file = fs::File::open(&path).await.map_err(|_| Error::NotFound)?;
197        let on_disk = file.metadata().await?.len();
198
199        match codec::Framed::open(file, on_disk).await? {
200            Some(framed) => Ok(Object::Framed(framed)),
201            None => Ok(Object::Raw {
202                file: fs::File::open(&path).await.map_err(|_| Error::NotFound)?,
203                size: on_disk,
204            }),
205        }
206    }
207
208    pub async fn write<S, E>(
209        &self,
210        ns: &Namespace,
211        oid: &str,
212        expected_size: Option<u64>,
213        budget: Option<Budget>,
214        mut chunks: S,
215    ) -> Result<u64, Error>
216    where
217        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
218        E: std::error::Error + Send + Sync + 'static,
219    {
220        Self::validate_oid(oid)?;
221
222        if let Some(limit) = self.max_object_size
223            && expected_size.is_some_and(|declared| declared > limit)
224        {
225            return Err(Error::TooLarge { limit });
226        }
227
228        let path = self.object_path(ns, oid);
229        let parent = path.parent().expect("object paths always have a parent");
230        fs::create_dir_all(parent).await?;
231
232        // A retried transfer of an object this repository already holds costs it
233        // no room, so it must not count against the budget a second time.
234        let fresh = fs::metadata(&path).await.is_err();
235
236        let staged = self.staging_path(parent, oid);
237        let outcome = self.stream_to(&staged, budget, &mut chunks).await;
238
239        match outcome {
240            Ok((digest, written)) => {
241                self.finish(&staged, &path, oid, expected_size, &digest, written)
242                    .await?;
243
244                if fresh {
245                    self.stored(ns, written).await;
246                }
247
248                Ok(written)
249            }
250            Err(error) => {
251                let _ = fs::remove_file(&staged).await;
252                Err(error)
253            }
254        }
255    }
256
257    fn staging_path(&self, parent: &Path, oid: &str) -> PathBuf {
258        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
259        parent.join(format!("{oid}.{ticket}.part"))
260    }
261
262    async fn stream_to<S, E>(
263        &self,
264        staged: &Path,
265        budget: Option<Budget>,
266        chunks: &mut S,
267    ) -> Result<(String, u64), Error>
268    where
269        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
270        E: std::error::Error + Send + Sync + 'static,
271    {
272        let file = fs::File::create(staged).await?;
273        // The digest, the declared size and the budget are all counted on the
274        // plaintext going past, whatever the bytes look like once they land —
275        // so compression is a different sink, not a different path.
276        let mut sink = match self.compression {
277            Some(level) => Sink::Framed(Box::new(codec::Writer::open(file, level).await?)),
278            None => Sink::Raw(file),
279        };
280        let mut hasher = Sha256::new();
281        let mut written = 0u64;
282
283        while let Some(chunk) = chunks.next().await {
284            let chunk = chunk.map_err(std::io::Error::other)?;
285            hasher.update(&chunk);
286            written += chunk.len() as u64;
287
288            // The declared size is a claim by the client, so the ceiling has to
289            // hold against a body that ignores it. Stopping at the chunk that
290            // crosses the line is the point: reading to the end to find out how
291            // big it was would be the outage this limit exists to prevent.
292            if let Some(limit) = self.max_object_size.filter(|limit| written > *limit) {
293                return Err(Error::TooLarge { limit });
294            }
295
296            if let Some(budget) = budget.filter(|budget| budget.exceeded_by(written)) {
297                return Err(budget.refusal());
298            }
299
300            sink.write(&chunk).await?;
301        }
302
303        sink.finish().await?;
304
305        Ok((hex::encode(hasher.finalize()), written))
306    }
307
308    async fn finish(
309        &self,
310        staged: &Path,
311        final_path: &Path,
312        oid: &str,
313        expected_size: Option<u64>,
314        digest: &str,
315        written: u64,
316    ) -> Result<(), Error> {
317        if let Some(declared) = expected_size.filter(|declared| *declared != written) {
318            let _ = fs::remove_file(staged).await;
319            return Err(Error::SizeMismatch {
320                declared,
321                actual: written,
322            });
323        }
324
325        if digest != oid {
326            let _ = fs::remove_file(staged).await;
327            return Err(Error::OidMismatch {
328                declared: oid.to_owned(),
329                actual: digest.to_owned(),
330            });
331        }
332
333        self.link_or_move(staged, final_path, oid).await
334    }
335
336    // One copy of the bytes under .content, and a hard link per repository that
337    // holds them. Two projects sharing an asset pack cost the disk once, and the
338    // link count is the reference count — the filesystem does the bookkeeping, so
339    // nothing can leak a repository's contents to another and nothing needs a
340    // migration: objects already sitting at their repository path keep working as
341    // ordinary files with one link.
342    async fn link_or_move(&self, staged: &Path, final_path: &Path, oid: &str) -> Result<(), Error> {
343        let content = self.content_path(oid);
344        let parent = content.parent().expect("content paths have a parent");
345        fs::create_dir_all(parent).await?;
346
347        if fs::metadata(&content).await.is_err() {
348            fs::rename(staged, &content).await?;
349        }
350
351        match self.link(&content, final_path).await {
352            // The content was collected between finding it and linking to it:
353            // a concurrent retain on another repository dropped its last other
354            // reference. The staged copy is still here precisely for this, so
355            // put it back and link again rather than failing a push that did
356            // nothing wrong.
357            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
358                fs::rename(staged, &content).await?;
359                self.link(&content, final_path).await?;
360            }
361            outcome => outcome?,
362        }
363
364        let _ = fs::remove_file(staged).await;
365        Ok(())
366    }
367
368    async fn link(&self, content: &Path, final_path: &Path) -> Result<(), std::io::Error> {
369        let from = content.to_path_buf();
370        let to = final_path.to_path_buf();
371        let linked = tokio::task::spawn_blocking(move || std::fs::hard_link(&from, &to))
372            .await
373            .map_err(std::io::Error::other)?;
374
375        match linked {
376            Ok(()) => Ok(()),
377            Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
378            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Err(error),
379            // A filesystem without hard links, or one crossing a device
380            // boundary: fall back to a full copy so the transfer still
381            // succeeds. The disk pays for it, the client never notices.
382            Err(_) => fs::copy(content, final_path).await.map(|_| ()),
383        }
384    }
385}
lfsx_server/storage/mod.rs

lfsx_server/storage/
mod.rs
