Skip to main content

lfsx_server/storage/
mod.rs

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