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 { .. } => Err(Error::Unsupported(
316 "collection is not implemented for a bucket yet",
317 )),
318 }
319 }
320
321 pub async fn dedupe(&self, ns: &Namespace, dry_run: bool) -> Result<DedupeReport, Error> {
322 match &self.backend {
323 Backend::Local(store) => {
324 let report = store.dedupe(ns, dry_run).await;
325
326 self.usage.forget(ns).await;
330
331 report
332 }
333 Backend::Bucket { .. } => Err(Error::Unsupported(
337 "a bucket stores each object once already, so there is nothing to deduplicate",
338 )),
339 }
340 }
341
342 pub async fn compress(&self, ns: &Namespace, dry_run: bool) -> Result<CompressReport, Error> {
343 match &self.backend {
344 Backend::Local(store) => {
345 let report = store.compress(ns, dry_run).await;
346
347 self.usage.forget(ns).await;
351
352 report
353 }
354 Backend::Bucket { .. } => Err(Error::Unsupported(
358 "rewriting objects already in a bucket is not implemented",
359 )),
360 }
361 }
362
363 pub async fn verify(&self, ns: &Namespace) -> Result<VerifyReport, Error> {
364 match &self.backend {
365 Backend::Local(store) => store.verify(ns).await,
366 Backend::Bucket { .. } => Err(Error::Unsupported(
367 "verification is not implemented for a bucket yet",
368 )),
369 }
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use futures_util::StreamExt;
376
377 use super::*;
378 use crate::storage::s3::tests::{bucket, store};
379
380 fn namespace() -> Namespace {
381 Namespace::new("FerrLabs", "Blastlands").unwrap()
382 }
383
384 fn bucket_store(root: &tempfile::TempDir, endpoint: &str) -> Store {
385 Store::bucket(store(endpoint), LocalStore::new(root.path()))
386 }
387
388 async fn read_back(store: &Store, ns: &Namespace, oid: &str) -> Vec<u8> {
389 let object = store.open(ns, oid).await.unwrap();
390 let size = object.size();
391 let mut chunks = object.stream(0, size).await.unwrap();
392 let mut out = Vec::new();
393
394 while let Some(chunk) = chunks.next().await {
395 out.extend_from_slice(&chunk.unwrap());
396 }
397
398 out
399 }
400
401 #[tokio::test]
402 async fn an_upload_lands_in_the_bucket_and_reads_back_through_the_same_seam() {
403 let root = tempfile::tempdir().unwrap();
404 let (endpoint, _objects) = bucket().await;
405 let store = bucket_store(&root, &endpoint);
406 let payload = b"an asset that never touches this disk for long".repeat(32);
407 let oid = hex::encode(sha2::Sha256::digest(&payload));
408
409 let written = store
410 .write(
411 &namespace(),
412 &oid,
413 Some(payload.len() as u64),
414 None,
415 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
416 payload.clone(),
417 ))]),
418 )
419 .await
420 .unwrap();
421
422 assert_eq!(written, payload.len() as u64);
423 assert!(store.exists(&namespace(), &oid).await);
424
425 assert_eq!(read_back(&store, &namespace(), &oid).await, payload);
426 }
427
428 #[tokio::test]
433 async fn a_bucket_holds_the_object_even_when_the_server_was_told_to_compress() {
434 let root = tempfile::tempdir().unwrap();
435 let (endpoint, _objects) = bucket().await;
436 let store = Store::bucket(
437 store(&endpoint),
438 LocalStore::new(root.path()).with_compression(Some(3)),
439 );
440 let payload = b"a mesh that gives up most of its ground to zstd ".repeat(4096);
441 let oid = hex::encode(sha2::Sha256::digest(&payload));
442
443 store
444 .write(
445 &namespace(),
446 &oid,
447 Some(payload.len() as u64),
448 None,
449 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
450 payload.clone(),
451 ))]),
452 )
453 .await
454 .unwrap();
455
456 let restored = read_back(&store, &namespace(), &oid).await;
457
458 assert_eq!(
459 hex::encode(sha2::Sha256::digest(&restored)),
460 oid,
461 "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",
462 restored.len()
463 );
464 assert_eq!(restored, payload);
465 }
466
467 #[tokio::test]
468 async fn the_staging_file_does_not_outlive_the_upload() {
469 let root = tempfile::tempdir().unwrap();
470 let (endpoint, _objects) = bucket().await;
471 let store = bucket_store(&root, &endpoint);
472 let payload = b"an asset passing through".to_vec();
473 let oid = hex::encode(sha2::Sha256::digest(&payload));
474
475 store
476 .write(
477 &namespace(),
478 &oid,
479 None,
480 None,
481 futures_util::stream::iter([Ok::<_, std::io::Error>(axum::body::Bytes::from(
482 payload,
483 ))]),
484 )
485 .await
486 .unwrap();
487
488 let leftovers = crate::storage::tests::staging_files(root.path());
489 assert!(
490 leftovers.is_empty(),
491 "local disk is a write buffer here, and one that is never emptied is a disk that \
492 fills: {leftovers:?}"
493 );
494 }
495
496 #[tokio::test]
497 async fn a_bucket_reports_no_capacity_rather_than_an_empty_one() {
498 let root = tempfile::tempdir().unwrap();
499 let (endpoint, _objects) = bucket().await;
500
501 assert!(
502 bucket_store(&root, &endpoint).capacity().await.is_none(),
503 "zero would be read as an empty store by every dashboard that averages it"
504 );
505 assert!(
506 Store::local(LocalStore::new(root.path()))
507 .capacity()
508 .await
509 .is_some()
510 );
511 }
512
513 #[tokio::test]
514 async fn the_maintenance_commands_say_they_do_not_apply_rather_than_lying() {
515 let root = tempfile::tempdir().unwrap();
516 let (endpoint, _objects) = bucket().await;
517 let store = bucket_store(&root, &endpoint);
518 let ns = namespace();
519
520 for outcome in [
521 store.dedupe(&ns, true).await.err(),
522 store.compress(&ns, true).await.err(),
523 store.verify(&ns).await.err(),
524 store
525 .sweep(&ns, &std::collections::HashSet::new(), Duration::ZERO, true)
526 .await
527 .err(),
528 ] {
529 assert!(
530 matches!(outcome, Some(Error::Unsupported(_))),
531 "an operator running collection against a bucket has to be told it did nothing, \
532 not handed an empty report that reads like success: {outcome:?}"
533 );
534 }
535 }
536}