lfsx_server/storage/
mod.rs

1mod sweep;
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Instant;
7
8use tokio::sync::Mutex;
9
10use futures_util::{Stream, StreamExt};
11use sha2::{Digest, Sha256};
12use tokio::fs;
13use tokio::io::AsyncWriteExt;
14
15use crate::error::Error;
16use crate::namespace::Namespace;
17
18pub use sweep::SweepReport;
19
20pub struct LocalStore {
21    root: PathBuf,
22    counter: AtomicU64,
23    usage: Mutex<Option<(Instant, u64, u64)>>,
24    per_namespace: Mutex<HashMap<String, (Instant, u64, u64)>>,
25    scans: AtomicU64,
26}
27
28impl LocalStore {
29    pub fn new(root: impl Into<PathBuf>) -> Self {
30        Self {
31            root: root.into(),
32            counter: AtomicU64::new(0),
33            usage: Mutex::new(None),
34            per_namespace: Mutex::new(HashMap::new()),
35            scans: AtomicU64::new(0),
36        }
37    }
38
39    pub fn validate_oid(oid: &str) -> Result<(), Error> {
40        let well_formed = oid.len() == 64
41            && oid
42                .bytes()
43                .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
44
45        well_formed.then_some(()).ok_or(Error::MalformedOid)
46    }
47
48    fn object_path(&self, ns: &Namespace, oid: &str) -> PathBuf {
49        self.root
50            .join(ns.org())
51            .join(ns.repo())
52            .join(&oid[0..2])
53            .join(&oid[2..4])
54            .join(oid)
55    }
56
57    pub fn scans(&self) -> u64 {
58        self.scans.load(Ordering::Relaxed)
59    }
60
61    pub async fn writable(&self) -> Result<(), Error> {
62        fs::create_dir_all(&self.root).await?;
63
64        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
65        let probe = self.root.join(format!(".readiness.{ticket}"));
66
67        fs::write(&probe, b"").await?;
68        fs::remove_file(&probe).await?;
69
70        Ok(())
71    }
72
73    pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
74        Self::validate_oid(oid).is_ok() && fs::metadata(self.object_path(ns, oid)).await.is_ok()
75    }
76
77    pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<(fs::File, u64), Error> {
78        Self::validate_oid(oid)?;
79        let path = self.object_path(ns, oid);
80        let file = fs::File::open(&path).await.map_err(|_| Error::NotFound)?;
81        let size = file.metadata().await?.len();
82        Ok((file, size))
83    }
84
85    pub async fn write<S, E>(
86        &self,
87        ns: &Namespace,
88        oid: &str,
89        expected_size: Option<u64>,
90        mut chunks: S,
91    ) -> Result<u64, Error>
92    where
93        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
94        E: std::error::Error + Send + Sync + 'static,
95    {
96        Self::validate_oid(oid)?;
97
98        let path = self.object_path(ns, oid);
99        let parent = path.parent().expect("object paths always have a parent");
100        fs::create_dir_all(parent).await?;
101
102        let staged = self.staging_path(parent, oid);
103        let outcome = self.stream_to(&staged, &mut chunks).await;
104
105        match outcome {
106            Ok((digest, written)) => {
107                self.finish(&staged, &path, oid, expected_size, &digest, written)
108                    .await?;
109                Ok(written)
110            }
111            Err(error) => {
112                let _ = fs::remove_file(&staged).await;
113                Err(error)
114            }
115        }
116    }
117
118    fn staging_path(&self, parent: &Path, oid: &str) -> PathBuf {
119        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
120        parent.join(format!("{oid}.{ticket}.part"))
121    }
122
123    async fn stream_to<S, E>(&self, staged: &Path, chunks: &mut S) -> Result<(String, u64), Error>
124    where
125        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
126        E: std::error::Error + Send + Sync + 'static,
127    {
128        let mut file = fs::File::create(staged).await?;
129        let mut hasher = Sha256::new();
130        let mut written = 0u64;
131
132        while let Some(chunk) = chunks.next().await {
133            let chunk = chunk.map_err(std::io::Error::other)?;
134            hasher.update(&chunk);
135            written += chunk.len() as u64;
136            file.write_all(&chunk).await?;
137        }
138
139        file.flush().await?;
140        file.sync_all().await?;
141
142        Ok((hex::encode(hasher.finalize()), written))
143    }
144
145    async fn finish(
146        &self,
147        staged: &Path,
148        final_path: &Path,
149        oid: &str,
150        expected_size: Option<u64>,
151        digest: &str,
152        written: u64,
153    ) -> Result<(), Error> {
154        if let Some(declared) = expected_size.filter(|declared| *declared != written) {
155            let _ = fs::remove_file(staged).await;
156            return Err(Error::SizeMismatch {
157                declared,
158                actual: written,
159            });
160        }
161
162        if digest != oid {
163            let _ = fs::remove_file(staged).await;
164            return Err(Error::OidMismatch {
165                declared: oid.to_owned(),
166                actual: digest.to_owned(),
167            });
168        }
169
170        fs::rename(staged, final_path).await?;
171        Ok(())
172    }
173}
lfsx_server/storage/mod.rs

lfsx_server/storage/
mod.rs
