1use futures_util::Stream;
2
3use super::s3::S3Store;
4use super::{Budget, CompressReport, DedupeReport, LocalStore, Object, SweepReport, VerifyReport};
5use crate::error::Error;
6use crate::namespace::Namespace;
7#[cfg(test)]
8use sha2::Digest;
9#[cfg(test)]
10use std::time::Duration;
11
12pub enum Store {
17 Local(LocalStore),
18 Bucket {
24 bucket: Box<S3Store>,
25 staging: LocalStore,
26 },
27}
28
29impl Store {
30 fn staging(&self) -> &LocalStore {
31 match self {
32 Self::Local(store) => store,
33 Self::Bucket { staging, .. } => staging,
34 }
35 }
36
37 pub async fn writable(&self) -> Result<(), Error> {
38 self.staging().writable().await
39 }
40
41 pub fn scans(&self) -> u64 {
42 self.staging().scans()
43 }
44
45 pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
46 match self {
47 Self::Local(store) => store.exists(ns, oid).await,
48 Self::Bucket { bucket, .. } => bucket.exists(ns, oid).await,
49 }
50 }
51
52 pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<Object, Error> {
53 match self {
54 Self::Local(store) => store.open(ns, oid).await,
55 Self::Bucket { bucket, .. } => {
56 if !bucket.exists(ns, oid).await {
57 return Err(Error::NotFound);
58 }
59
60 Ok(Object::Remote {
61 bucket: (**bucket).clone(),
62 oid: oid.to_owned(),
63 size: bucket.size_of(oid).await?,
64 })
65 }
66 }
67 }
68
69 pub async fn write<S, E>(
70 &self,
71 ns: &Namespace,
72 oid: &str,
73 expected_size: Option<u64>,
74 budget: Option<Budget>,
75 chunks: S,
76 ) -> Result<u64, Error>
77 where
78 S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
79 E: std::error::Error + Send + Sync + 'static,
80 {
81 match self {
82 Self::Local(store) => store.write(ns, oid, expected_size, budget, chunks).await,
83 Self::Bucket { bucket, staging } => {
84 let staged = staging
85 .stage(ns, oid, expected_size, budget, chunks)
86 .await?;
87 let outcome = bucket.store(ns, oid, &staged.path).await;
88
89 let _ = tokio::fs::remove_file(&staged.path).await;
92 outcome?;
93
94 Ok(staged.written)
95 }
96 }
97 }
98
99 pub async fn capacity(&self) -> Option<(u64, u64)> {
105 match self {
106 Self::Local(store) => Some(store.usage().await),
107 Self::Bucket { .. } => None,
108 }
109 }
110
111 pub async fn usage_of(&self, ns: &Namespace) -> (u64, u64) {
112 match self {
113 Self::Local(store) => store.usage_of(ns).await,
114 Self::Bucket { bucket, .. } => bucket.usage_of(ns).await,
115 }
116 }
117
118 pub async fn sweep(
119 &self,
120 ns: &Namespace,
121 retained: &std::collections::HashSet<String>,
122 grace: std::time::Duration,
123 dry_run: bool,
124 ) -> Result<SweepReport, Error> {
125 match self {
126 Self::Local(store) => store.sweep(ns, retained, grace, dry_run).await,
127 Self::Bucket { .. } => Err(Error::Unsupported(
128 "collection is not implemented for a bucket yet",
129 )),
130 }
131 }
132
133 pub async fn dedupe(&self, ns: &Namespace, dry_run: bool) -> Result<DedupeReport, Error> {
134 match self {
135 Self::Local(store) => store.dedupe(ns, dry_run).await,
136 Self::Bucket { .. } => Err(Error::Unsupported(
140 "a bucket stores each object once already, so there is nothing to deduplicate",
141 )),
142 }
143 }
144
145 pub async fn compress(&self, ns: &Namespace, dry_run: bool) -> Result<CompressReport, Error> {
146 match self {
147 Self::Local(store) => store.compress(ns, dry_run).await,
148 Self::Bucket { .. } => Err(Error::Unsupported(
149 "compression is not implemented for a bucket yet",
150 )),
151 }
152 }
153
154 pub async fn verify(&self, ns: &Namespace) -> Result<VerifyReport, Error> {
155 match self {
156 Self::Local(store) => store.verify(ns).await,
157 Self::Bucket { .. } => Err(Error::Unsupported(
158 "verification is not implemented for a bucket yet",
159 )),
160 }
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use futures_util::StreamExt;
167
168 use super::*;
169 use crate::storage::s3::tests::{bucket, store};
170
171 fn namespace() -> Namespace {
172 Namespace::new("FerrLabs", "Blastlands").unwrap()
173 }
174
175 async fn bucket_store(root: &tempfile::TempDir, endpoint: &str) -> Store {
176 Store::Bucket {
177 bucket: Box::new(store(endpoint)),
178 staging: LocalStore::new(root.path()),
179 }
180 }
181
182 #[tokio::test]
183 async fn an_upload_lands_in_the_bucket_and_reads_back_through_the_same_seam() {
184 let root = tempfile::tempdir().unwrap();
185 let (endpoint, _objects) = bucket().await;
186 let store = bucket_store(&root, &endpoint).await;
187 let payload = b"an asset that never touches this disk for long".repeat(32);
188 let oid = hex::encode(sha2::Sha256::digest(&payload));
189
190 let written = store
191 .write(
192 &namespace(),
193 &oid,
194 Some(payload.len() as u64),
195 None,
196 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
197 payload.clone(),
198 ))]),
199 )
200 .await
201 .unwrap();
202
203 assert_eq!(written, payload.len() as u64);
204 assert!(store.exists(&namespace(), &oid).await);
205
206 let object = store.open(&namespace(), &oid).await.unwrap();
207 let size = object.size();
208 let mut chunks = object.stream(0, size).await.unwrap();
209 let mut out = Vec::new();
210 while let Some(chunk) = chunks.next().await {
211 out.extend_from_slice(&chunk.unwrap());
212 }
213
214 assert_eq!(out, payload);
215 }
216
217 #[tokio::test]
218 async fn the_staging_file_does_not_outlive_the_upload() {
219 let root = tempfile::tempdir().unwrap();
220 let (endpoint, _objects) = bucket().await;
221 let store = bucket_store(&root, &endpoint).await;
222 let payload = b"an asset passing through".to_vec();
223 let oid = hex::encode(sha2::Sha256::digest(&payload));
224
225 store
226 .write(
227 &namespace(),
228 &oid,
229 None,
230 None,
231 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
232 payload,
233 ))]),
234 )
235 .await
236 .unwrap();
237
238 let leftovers = crate::storage::tests::staging_files(root.path());
239 assert!(
240 leftovers.is_empty(),
241 "local disk is a write buffer here, and one that is never emptied is a disk that \
242 fills: {leftovers:?}"
243 );
244 }
245
246 #[tokio::test]
247 async fn a_bucket_reports_no_capacity_rather_than_an_empty_one() {
248 let root = tempfile::tempdir().unwrap();
249 let (endpoint, _objects) = bucket().await;
250
251 assert!(
252 bucket_store(&root, &endpoint)
253 .await
254 .capacity()
255 .await
256 .is_none(),
257 "zero would be read as an empty store by every dashboard that averages it"
258 );
259 assert!(
260 Store::Local(LocalStore::new(root.path()))
261 .capacity()
262 .await
263 .is_some()
264 );
265 }
266
267 #[tokio::test]
268 async fn the_maintenance_commands_say_they_do_not_apply_rather_than_lying() {
269 let root = tempfile::tempdir().unwrap();
270 let (endpoint, _objects) = bucket().await;
271 let store = bucket_store(&root, &endpoint).await;
272 let ns = namespace();
273
274 for outcome in [
275 store.dedupe(&ns, true).await.err(),
276 store.compress(&ns, true).await.err(),
277 store.verify(&ns).await.err(),
278 store
279 .sweep(&ns, &std::collections::HashSet::new(), Duration::ZERO, true)
280 .await
281 .err(),
282 ] {
283 assert!(
284 matches!(outcome, Some(Error::Unsupported(_))),
285 "an operator running collection against a bucket has to be told it did nothing, \
286 not handed an empty report that reads like success: {outcome:?}"
287 );
288 }
289 }
290}