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 struct Store {
17 backend: Backend,
18 usage: super::usage::Usage,
19}
20
21enum Backend {
22 Local(LocalStore),
23 Bucket {
29 bucket: Box<S3Store>,
30 staging: LocalStore,
31 },
32}
33
34impl Store {
35 pub fn local(store: LocalStore) -> Self {
36 Self::over(Backend::Local(store))
37 }
38
39 pub fn bucket(bucket: S3Store, staging: LocalStore) -> Self {
45 Self::over(Backend::Bucket {
46 bucket: Box::new(bucket),
47 staging,
48 })
49 }
50
51 fn over(backend: Backend) -> Self {
52 Self {
53 backend,
54 usage: super::usage::Usage::default(),
55 }
56 }
57
58 fn staging(&self) -> &LocalStore {
59 match &self.backend {
60 Backend::Local(store) => store,
61 Backend::Bucket { staging, .. } => staging,
62 }
63 }
64
65 pub async fn reclaim(&self, older_than: std::time::Duration) -> super::Reclaimed {
69 let mut reclaimed = self.staging().reclaim_staging(older_than).await;
70
71 if let Backend::Bucket { bucket, .. } = &self.backend {
72 match bucket.reclaim_incoming(older_than).await {
73 Ok(theirs) => {
74 reclaimed.files += theirs.files;
75 reclaimed.bytes += theirs.bytes;
76 }
77 Err(error) => {
78 tracing::warn!(%error, "abandoned uploads in the bucket could not be reclaimed");
79 }
80 }
81 }
82
83 reclaimed
84 }
85
86 pub async fn writable(&self) -> Result<(), Error> {
94 self.staging().writable().await.map_err(|error| {
95 Error::Storage(std::io::Error::other(format!(
96 "the staging volume is not writable: {error}"
97 )))
98 })?;
99
100 if let Backend::Bucket { bucket, .. } = &self.backend {
101 bucket.reachable().await?;
102 }
103
104 Ok(())
105 }
106
107 pub fn scans(&self) -> u64 {
108 self.staging().scans()
109 }
110
111 pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
112 match &self.backend {
113 Backend::Local(store) => store.exists(ns, oid).await,
114 Backend::Bucket { bucket, .. } => bucket.exists(ns, oid).await,
115 }
116 }
117
118 pub fn redirect(&self, oid: &str) -> Option<String> {
126 match &self.backend {
127 Backend::Local(_) => None,
128 Backend::Bucket { bucket, .. } => bucket.presigned_download(oid),
129 }
130 }
131
132 pub fn presigned_upload(&self, ns: &Namespace, oid: &str) -> Option<super::s3::Presigned> {
136 match &self.backend {
137 Backend::Local(_) => None,
138 Backend::Bucket { staging, .. } if staging.encrypts() => None,
145 Backend::Bucket { bucket, .. } => bucket.presigned_upload(ns, oid),
146 }
147 }
148
149 pub async fn uploaded_size(&self, ns: &Namespace, oid: &str) -> Result<Option<u64>, Error> {
153 match &self.backend {
154 Backend::Local(_) => Ok(None),
155 Backend::Bucket { bucket, .. } => Ok(bucket.uploaded_size(ns, oid).await.ok()),
156 }
157 }
158
159 pub async fn adopt(&self, ns: &Namespace, oid: &str, arrived: u64) -> Result<(), Error> {
163 let outcome = match &self.backend {
164 Backend::Local(_) => Err(Error::Unsupported(
165 "objects are written through this server, so there is nothing to adopt",
166 )),
167 Backend::Bucket { bucket, .. } => bucket.adopt(ns, oid).await,
168 };
169
170 if outcome.is_ok() {
173 self.usage.stored(ns, arrived).await;
174 }
175
176 outcome
177 }
178
179 pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<Object, Error> {
180 match &self.backend {
181 Backend::Local(store) => store.open(ns, oid).await,
182 Backend::Bucket { bucket, staging } => {
183 if !bucket.exists(ns, oid).await {
186 return Err(Error::NotFound);
187 }
188
189 let size = bucket.size_of(oid).await?;
190 let reader = super::codec::Reader::Bucket {
191 bucket: (**bucket).clone(),
192 oid: oid.to_owned(),
193 };
194
195 match super::codec::Framed::open(
196 reader,
197 size,
198 staging.keyring().map(AsRef::as_ref),
199 oid,
200 )
201 .await?
202 {
203 Some(framed) => Ok(Object::Framed(framed)),
204 None => Ok(Object::Remote {
207 bucket: (**bucket).clone(),
208 oid: oid.to_owned(),
209 size,
210 }),
211 }
212 }
213 }
214 }
215
216 pub async fn write<S, E>(
217 &self,
218 ns: &Namespace,
219 oid: &str,
220 expected_size: Option<u64>,
221 budget: Option<Budget>,
222 chunks: S,
223 ) -> Result<u64, Error>
224 where
225 S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
226 E: std::error::Error + Send + Sync + 'static,
227 {
228 let written = match &self.backend {
229 Backend::Local(store) => store.write(ns, oid, expected_size, budget, chunks).await?,
230 Backend::Bucket { bucket, staging } => {
231 let fresh = !bucket.exists(ns, oid).await;
237
238 let staged = staging
239 .stage(ns, oid, expected_size, budget, chunks)
240 .await?;
241 let outcome = bucket.store(ns, oid, &staged.path).await;
242
243 let _ = tokio::fs::remove_file(&staged.path).await;
246 outcome?;
247
248 super::Written {
249 bytes: staged.written,
250 fresh,
251 }
252 }
253 };
254
255 if written.fresh {
260 self.usage.stored(ns, written.bytes).await;
261 }
262
263 Ok(written.bytes)
264 }
265
266 pub async fn capacity(&self) -> Option<(u64, u64)> {
272 match &self.backend {
273 Backend::Local(store) => Some(store.usage().await),
274 Backend::Bucket { .. } => None,
275 }
276 }
277
278 pub async fn usage_of(&self, ns: &Namespace) -> (u64, u64) {
283 if let Some(cached) = self.usage.cached(ns).await {
284 return cached;
285 }
286
287 let measured = match &self.backend {
288 Backend::Local(store) => store.measure_of(ns).await,
289 Backend::Bucket { bucket, .. } => bucket.usage_of(ns).await,
290 };
291
292 self.usage.remember(ns, measured.0, measured.1).await;
293
294 measured
295 }
296
297 pub async fn sweep(
298 &self,
299 ns: &Namespace,
300 retained: &std::collections::HashSet<String>,
301 grace: std::time::Duration,
302 dry_run: bool,
303 ) -> Result<SweepReport, Error> {
304 match &self.backend {
305 Backend::Local(store) => {
306 let report = store.sweep(ns, retained, grace, dry_run).await;
307
308 self.usage.forget(ns).await;
312
313 report
314 }
315 Backend::Bucket { bucket, .. } => {
316 let report = bucket.sweep(ns, retained, grace, dry_run).await;
317
318 if report.is_ok() && !dry_run {
319 self.usage.forget(ns).await;
320 }
321
322 report
323 }
324 }
325 }
326
327 pub async fn dedupe(&self, ns: &Namespace, dry_run: bool) -> Result<DedupeReport, Error> {
328 match &self.backend {
329 Backend::Local(store) => {
330 let report = store.dedupe(ns, dry_run).await;
331
332 self.usage.forget(ns).await;
336
337 report
338 }
339 Backend::Bucket { .. } => Err(Error::Unsupported(
343 "a bucket stores each object once already, so there is nothing to deduplicate",
344 )),
345 }
346 }
347
348 pub async fn compress(&self, ns: &Namespace, dry_run: bool) -> Result<CompressReport, Error> {
349 match &self.backend {
350 Backend::Local(store) => {
351 let report = store.compress(ns, dry_run).await;
352
353 self.usage.forget(ns).await;
357
358 report
359 }
360 Backend::Bucket { .. } => Err(Error::Unsupported(
364 "rewriting objects already in a bucket is not implemented",
365 )),
366 }
367 }
368
369 pub async fn verify(&self, ns: &Namespace) -> Result<VerifyReport, Error> {
370 match &self.backend {
371 Backend::Local(store) => store.verify(ns).await,
372 Backend::Bucket { .. } => Err(Error::Unsupported(
373 "verification is not implemented for a bucket yet",
374 )),
375 }
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use futures_util::StreamExt;
382
383 use super::*;
384 use crate::storage::s3::tests::{bucket, store};
385
386 fn namespace() -> Namespace {
387 Namespace::new("FerrLabs", "Blastlands").unwrap()
388 }
389
390 fn bucket_store(root: &tempfile::TempDir, endpoint: &str) -> Store {
391 Store::bucket(store(endpoint), LocalStore::new(root.path()))
392 }
393
394 async fn read_back(store: &Store, ns: &Namespace, oid: &str) -> Vec<u8> {
395 let object = store.open(ns, oid).await.unwrap();
396 let size = object.size();
397 let mut chunks = object.stream(0, size).await.unwrap();
398 let mut out = Vec::new();
399
400 while let Some(chunk) = chunks.next().await {
401 out.extend_from_slice(&chunk.unwrap());
402 }
403
404 out
405 }
406
407 #[tokio::test]
408 async fn an_upload_lands_in_the_bucket_and_reads_back_through_the_same_seam() {
409 let root = tempfile::tempdir().unwrap();
410 let (endpoint, _objects) = bucket().await;
411 let store = bucket_store(&root, &endpoint);
412 let payload = b"an asset that never touches this disk for long".repeat(32);
413 let oid = hex::encode(sha2::Sha256::digest(&payload));
414
415 let written = store
416 .write(
417 &namespace(),
418 &oid,
419 Some(payload.len() as u64),
420 None,
421 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
422 payload.clone(),
423 ))]),
424 )
425 .await
426 .unwrap();
427
428 assert_eq!(written, payload.len() as u64);
429 assert!(store.exists(&namespace(), &oid).await);
430
431 assert_eq!(read_back(&store, &namespace(), &oid).await, payload);
432 }
433
434 #[tokio::test]
439 async fn a_bucket_holds_the_object_even_when_the_server_was_told_to_compress() {
440 let root = tempfile::tempdir().unwrap();
441 let (endpoint, _objects) = bucket().await;
442 let store = Store::bucket(
443 store(&endpoint),
444 LocalStore::new(root.path()).with_compression(Some(3)),
445 );
446 let payload = b"a mesh that gives up most of its ground to zstd ".repeat(4096);
447 let oid = hex::encode(sha2::Sha256::digest(&payload));
448
449 store
450 .write(
451 &namespace(),
452 &oid,
453 Some(payload.len() as u64),
454 None,
455 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
456 payload.clone(),
457 ))]),
458 )
459 .await
460 .unwrap();
461
462 let restored = read_back(&store, &namespace(), &oid).await;
463
464 assert_eq!(
465 hex::encode(sha2::Sha256::digest(&restored)),
466 oid,
467 "the client asked for the object named by this digest and has no way to know the server framed it on the way past: {} bytes came back",
468 restored.len()
469 );
470 assert_eq!(restored, payload);
471 }
472
473 #[tokio::test]
474 async fn the_staging_file_does_not_outlive_the_upload() {
475 let root = tempfile::tempdir().unwrap();
476 let (endpoint, _objects) = bucket().await;
477 let store = bucket_store(&root, &endpoint);
478 let payload = b"an asset passing through".to_vec();
479 let oid = hex::encode(sha2::Sha256::digest(&payload));
480
481 store
482 .write(
483 &namespace(),
484 &oid,
485 None,
486 None,
487 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
488 payload,
489 ))]),
490 )
491 .await
492 .unwrap();
493
494 let leftovers = crate::storage::tests::staging_files(root.path());
495 assert!(
496 leftovers.is_empty(),
497 "local disk is a write buffer here, and one that is never emptied is a disk that \
498 fills: {leftovers:?}"
499 );
500 }
501
502 #[tokio::test]
503 async fn a_bucket_reports_no_capacity_rather_than_an_empty_one() {
504 let root = tempfile::tempdir().unwrap();
505 let (endpoint, _objects) = bucket().await;
506
507 assert!(
508 bucket_store(&root, &endpoint).capacity().await.is_none(),
509 "zero would be read as an empty store by every dashboard that averages it"
510 );
511 assert!(
512 Store::local(LocalStore::new(root.path()))
513 .capacity()
514 .await
515 .is_some()
516 );
517 }
518
519 #[tokio::test]
520 async fn the_maintenance_commands_say_they_do_not_apply_rather_than_lying() {
521 let root = tempfile::tempdir().unwrap();
522 let (endpoint, _objects) = bucket().await;
523 let store = bucket_store(&root, &endpoint);
524 let ns = namespace();
525
526 for outcome in [
527 store.dedupe(&ns, true).await.err(),
528 store.compress(&ns, true).await.err(),
529 store.verify(&ns).await.err(),
530 ] {
531 assert!(
532 matches!(outcome, Some(Error::Unsupported(_))),
533 "an operator running one of these against a bucket has to be told it did nothing, not handed an empty report that reads like success: {outcome:?}"
534 );
535 }
536
537 assert!(
539 store
540 .sweep(&ns, &std::collections::HashSet::new(), Duration::ZERO, true)
541 .await
542 .is_ok(),
543 "collection is implemented for a bucket and must not answer Unsupported"
544 );
545 }
546}