1use std::future::Future;
9use std::sync::Arc;
10
11use objectstore_types::metadata::Metadata;
12use objectstore_types::range::{ByteRange, ContentRange};
13
14use crate::backend::common::Backend;
15use crate::backend::counting::CountingBackend;
16use crate::concurrency::ConcurrencyLimiter;
17use crate::error::Result;
18use crate::id::{ObjectContext, ObjectId};
19use crate::multipart::{
20 AbortMultipartResponse, CompleteMultipartResponse, CompletedPart, InitiateMultipartResponse,
21 ListPartsResponse, PartNumber, UploadId, UploadPartResponse,
22};
23use crate::stream::{ClientStream, PayloadStream};
24use crate::streaming::StreamExecutor;
25
26pub type GetResponse = Option<(Metadata, Option<ContentRange>, PayloadStream)>;
28pub type MetadataResponse = Option<Metadata>;
30pub type InsertResponse = ObjectId;
32pub type DeleteResponse = ();
34
35pub const DEFAULT_CONCURRENCY_LIMIT: u32 = 500;
40
41#[derive(Clone, Debug)]
73pub struct StorageService {
74 inner: Arc<dyn Backend>,
75 concurrency: ConcurrencyLimiter,
76}
77
78impl StorageService {
79 pub fn new(backend: Box<dyn Backend>) -> Self {
86 Self {
87 inner: Arc::new(CountingBackend::new(backend)),
88 concurrency: ConcurrencyLimiter::new(DEFAULT_CONCURRENCY_LIMIT),
89 }
90 }
91
92 pub fn with_concurrency(mut self, limiter: ConcurrencyLimiter) -> Self {
98 self.concurrency = limiter;
99 self
100 }
101
102 pub fn concurrency_limiter(&self) -> &ConcurrencyLimiter {
104 &self.concurrency
105 }
106
107 pub fn tasks_running(&self) -> u32 {
109 self.concurrency.used_permits()
110 }
111
112 pub fn tasks_limit(&self) -> u32 {
114 self.concurrency.total_permits()
115 }
116
117 pub fn stream(&self) -> StreamExecutor {
124 StreamExecutor::new(Arc::clone(&self.inner), self.concurrency.clone())
125 }
126
127 pub fn start(&self) {
140 let concurrency = self.concurrency.clone();
141 objectstore_metrics::gauge!("service.concurrency.limit" = concurrency.total_permits());
142 objectstore_metrics::gauge!("service.concurrency.queue_limit" = concurrency.total_queue());
143 objectstore_metrics::gauge!("service.concurrency.bulk_limit" = concurrency.total_bulk());
144
145 tokio::spawn(async move {
146 concurrency
147 .run_emitter(|stats| async move {
148 objectstore_metrics::gauge!("service.concurrency.in_use" = stats.in_use);
149 objectstore_metrics::gauge!("service.concurrency.queued" = stats.queued);
150 objectstore_metrics::gauge!(
151 "service.concurrency.bulk_in_use" = stats.bulk_in_use
152 );
153 })
154 .await;
155 });
156 }
157
158 async fn spawn<T, F>(&self, operation: &'static str, f: F) -> Result<T>
176 where
177 T: Send + 'static,
178 F: Future<Output = Result<T>> + Send + 'static,
179 {
180 let timer = objectstore_metrics::timer!("service.concurrency.wait");
181 let permit = self.concurrency.acquire().await.inspect_err(|_| {
182 objectstore_metrics::count!("service.concurrency.rejected", class = "normal");
183 objectstore_log::warn!("Request rejected: service at capacity");
184 })?;
185
186 timer.record();
187 crate::concurrency::spawn_metered(operation, permit, f).await
188 }
189
190 pub async fn insert_object(
202 &self,
203 context: ObjectContext,
204 key: Option<String>,
205 metadata: Metadata,
206 stream: ClientStream,
207 ) -> Result<InsertResponse> {
208 metadata.validate()?;
209 let id = ObjectId::optional(context, key);
210 let inner = Arc::clone(&self.inner);
211 self.spawn("insert", async move {
212 inner.put_object(&id, &metadata, stream).await?;
213 Ok(id)
214 })
215 .await
216 }
217
218 pub async fn get_metadata(&self, id: ObjectId) -> Result<MetadataResponse> {
220 let inner = Arc::clone(&self.inner);
221 self.spawn("get_metadata", async move { inner.get_metadata(&id).await })
222 .await
223 }
224
225 pub async fn get_object(&self, id: ObjectId, range: Option<ByteRange>) -> Result<GetResponse> {
227 let inner = Arc::clone(&self.inner);
228 self.spawn("get", async move { inner.get_object(&id, range).await })
229 .await
230 }
231
232 pub async fn delete_object(&self, id: ObjectId) -> Result<DeleteResponse> {
240 let inner = Arc::clone(&self.inner);
241 self.spawn("delete", async move { inner.delete_object(&id).await })
242 .await
243 }
244
245 pub async fn join(&self) {
251 self.inner.join().await;
252 }
253
254 pub async fn initiate_multipart(
258 &self,
259 id: ObjectId,
260 metadata: Metadata,
261 ) -> Result<InitiateMultipartResponse> {
262 metadata.validate()?;
263 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
265 self.spawn("initiate_multipart", async move {
266 inner
267 .as_multipart_upload_backend()?
268 .initiate_multipart(&id, &metadata)
269 .await
270 })
271 .await
272 }
273
274 pub async fn upload_part(
283 &self,
284 id: ObjectId,
285 upload_id: UploadId,
286 part_number: PartNumber,
287 content_length: u64,
288 content_md5: Option<String>,
289 body: ClientStream,
290 ) -> Result<UploadPartResponse> {
291 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
293 self.spawn("upload_part", async move {
294 inner
295 .as_multipart_upload_backend()?
296 .upload_part(
297 &id,
298 &upload_id,
299 part_number,
300 content_length,
301 content_md5.as_deref(),
302 body,
303 )
304 .await
305 })
306 .await
307 }
308
309 pub async fn list_parts(
311 &self,
312 id: ObjectId,
313 upload_id: UploadId,
314 max_parts: Option<u32>,
315 part_number_marker: Option<PartNumber>,
316 ) -> Result<ListPartsResponse> {
317 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
319 self.spawn("list_parts", async move {
320 inner
321 .as_multipart_upload_backend()?
322 .list_parts(&id, &upload_id, max_parts, part_number_marker)
323 .await
324 })
325 .await
326 }
327
328 pub async fn abort_multipart(
330 &self,
331 id: ObjectId,
332 upload_id: UploadId,
333 ) -> Result<AbortMultipartResponse> {
334 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
336 self.spawn("abort_multipart", async move {
337 inner
338 .as_multipart_upload_backend()?
339 .abort_multipart(&id, &upload_id)
340 .await
341 })
342 .await
343 }
344
345 pub async fn complete_multipart(
347 &self,
348 id: ObjectId,
349 upload_id: UploadId,
350 parts: Vec<CompletedPart>,
351 ) -> Result<CompleteMultipartResponse> {
352 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
354 self.spawn("complete_multipart", async move {
355 inner
356 .as_multipart_upload_backend()?
357 .complete_multipart(&id, &upload_id, parts)
358 .await
359 })
360 .await
361 }
362}
363
364#[cfg(test)]
365mod tests {
366 use std::sync::Arc;
367 use std::time::Duration;
368
369 use bytes::BytesMut;
370 use futures_util::TryStreamExt;
371 use objectstore_types::metadata::Metadata;
372 use objectstore_types::range::ByteRange;
373 use objectstore_types::scope::{Scope, Scopes};
374
375 use super::*;
376 use crate::backend::bigtable::{BigTableBackend, BigTableConfig};
377 use crate::backend::changelog::NoopChangeLog;
378 use crate::backend::common::{HighVolumeBackend, PutResponse, TieredWrite};
379 use crate::backend::gcs::{GcsBackend, GcsConfig};
380 use crate::backend::in_memory::InMemoryBackend;
381 use crate::backend::testing::{Hooks, TestBackend};
382 use crate::backend::tiered::TieredStorage;
383 use crate::error::Error;
384 use crate::stream::{self, ClientStream};
385
386 fn make_context() -> ObjectContext {
387 ObjectContext {
388 usecase: "testing".into(),
389 scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
390 }
391 }
392
393 fn make_service() -> StorageService {
394 StorageService::new(Box::new(InMemoryBackend::new("in-memory")))
395 }
396
397 #[tokio::test]
398 async fn insert_without_key_generates_unique_id() {
399 let service = make_service();
400
401 let id = service
402 .insert_object(
403 make_context(),
404 None,
405 Metadata::default(),
406 stream::single("auto-keyed"),
407 )
408 .await
409 .unwrap();
410
411 assert!(uuid::Uuid::parse_str(id.key()).is_ok());
412 }
413
414 #[tokio::test]
415 async fn stores_files() {
416 let service = make_service();
417
418 let key = service
419 .insert_object(
420 make_context(),
421 Some("testing".into()),
422 Metadata::default(),
423 stream::single("oh hai!"),
424 )
425 .await
426 .unwrap();
427
428 let (_metadata, _, stream) = service.get_object(key, None).await.unwrap().unwrap();
429 let file_contents: BytesMut = stream.try_collect().await.unwrap();
430
431 assert_eq!(file_contents.as_ref(), b"oh hai!");
432 }
433
434 #[tokio::test]
435 async fn works_with_gcs() {
436 let config = GcsConfig {
437 endpoint: Some("http://localhost:8087".into()),
438 bucket: "test-bucket".into(), };
440
441 let backend = GcsBackend::new(config).await.unwrap();
442 let service = StorageService::new(Box::new(backend));
443
444 let key = service
445 .insert_object(
446 make_context(),
447 Some("testing".into()),
448 Metadata::default(),
449 stream::single("oh hai!"),
450 )
451 .await
452 .unwrap();
453
454 let (_metadata, _, stream) = service.get_object(key, None).await.unwrap().unwrap();
455 let file_contents: BytesMut = stream.try_collect().await.unwrap();
456
457 assert_eq!(file_contents.as_ref(), b"oh hai!");
458 }
459
460 #[tokio::test]
461 async fn tombstone_redirect_and_delete() {
462 let bigtable_config = BigTableConfig {
463 endpoint: Some("localhost:8086".into()),
464 project_id: "testing".into(),
465 instance_name: "objectstore".into(),
466 table_name: "objectstore".into(),
467 connections: None,
468 };
469 let gcs_config = GcsConfig {
470 endpoint: Some("http://localhost:8087".into()),
471 bucket: "test-bucket".into(),
472 };
473
474 let high_volume = Box::new(BigTableBackend::new(bigtable_config).await.unwrap());
475 let long_term = Box::new(GcsBackend::new(gcs_config.clone()).await.unwrap());
476 let backend = TieredStorage::new(high_volume, long_term, Box::new(NoopChangeLog));
477 let service = StorageService::new(Box::new(backend));
478
479 let gcs_backend = GcsBackend::new(gcs_config.clone()).await.unwrap();
481
482 let payload_len = 2 * 1024 * 1024;
485 let payload = vec![0xAB; payload_len]; let id = service
487 .insert_object(
488 make_context(),
489 Some("delete-cleanup-test".into()),
490 Metadata::default(),
491 stream::single(payload),
492 )
493 .await
494 .unwrap();
495
496 let (_, _, stream) = service.get_object(id.clone(), None).await.unwrap().unwrap();
498 let body: BytesMut = stream.try_collect().await.unwrap();
499 assert_eq!(body.len(), payload_len);
500
501 service.delete_object(id.clone()).await.unwrap();
503
504 let after_delete = service.get_object(id.clone(), None).await.unwrap();
506 assert!(after_delete.is_none(), "tombstone not deleted");
507
508 let orphan = gcs_backend.get_object(&id, None).await.unwrap();
510 assert!(orphan.is_none(), "object leaked");
511 }
512
513 #[tokio::test]
516 async fn basic_spawn_insert_and_get() {
517 let service = make_service();
518
519 let id = service
520 .insert_object(
521 make_context(),
522 Some("test-key".into()),
523 Metadata::default(),
524 stream::single("hello world"),
525 )
526 .await
527 .unwrap();
528
529 let (_, _, stream) = service.get_object(id, None).await.unwrap().unwrap();
530 let body: BytesMut = stream.try_collect().await.unwrap();
531 assert_eq!(body.as_ref(), b"hello world");
532 }
533
534 #[tokio::test]
535 async fn basic_spawn_metadata_and_delete() {
536 let service = make_service();
537
538 let id = service
539 .insert_object(
540 make_context(),
541 Some("meta-key".into()),
542 Metadata::default(),
543 stream::single("data"),
544 )
545 .await
546 .unwrap();
547
548 let metadata = service.get_metadata(id.clone()).await.unwrap();
549 assert!(metadata.is_some());
550
551 service.delete_object(id.clone()).await.unwrap();
552
553 let after = service.get_object(id, None).await.unwrap();
554 assert!(after.is_none());
555 }
556
557 #[derive(Debug)]
558 struct PanicOnGet;
559
560 #[async_trait::async_trait]
561 impl Hooks for PanicOnGet {
562 async fn get_object(
563 &self,
564 _inner: &InMemoryBackend,
565 _id: &ObjectId,
566 _range: Option<ByteRange>,
567 ) -> Result<GetResponse> {
568 panic!("intentional panic in get_object");
569 }
570 }
571
572 #[tokio::test]
573 async fn panic_in_backend_returns_task_failed() {
574 let service = StorageService::new(Box::new(TestBackend::new(PanicOnGet)));
575
576 let id = ObjectId::new(make_context(), "panic-test".into());
577 let result = service.get_object(id, None).await;
578
579 let Err(Error::Panic(msg)) = result else {
580 panic!("expected Panic error");
581 };
582 assert!(msg.contains("intentional panic in get_object"), "{msg}");
583 }
584
585 #[derive(Clone, Debug, Default)]
589 struct GateOnPut {
590 pause: bool,
591 paused: Arc<tokio::sync::Notify>,
592 resume: Arc<tokio::sync::Notify>,
593 on_put: Arc<tokio::sync::Notify>,
594 }
595
596 impl GateOnPut {
597 fn with_pause() -> Self {
598 Self {
599 pause: true,
600 ..Default::default()
601 }
602 }
603 }
604
605 #[async_trait::async_trait]
606 impl Hooks for GateOnPut {
607 async fn put_object(
608 &self,
609 inner: &InMemoryBackend,
610 id: &ObjectId,
611 metadata: &Metadata,
612 stream: ClientStream,
613 ) -> Result<PutResponse> {
614 if self.pause {
615 self.paused.notify_one();
616 self.resume.notified().await;
617 }
618 inner.put_object(id, metadata, stream).await?;
619 self.on_put.notify_one();
620 Ok(())
621 }
622
623 async fn compare_and_write(
624 &self,
625 inner: &InMemoryBackend,
626 id: &ObjectId,
627 current: Option<&ObjectId>,
628 write: TieredWrite,
629 ) -> Result<bool> {
630 let notify = matches!(write, TieredWrite::Tombstone(_) | TieredWrite::Object(_, _));
631 let result = inner.compare_and_write(id, current, write).await?;
632 if notify {
633 self.on_put.notify_one();
634 }
635 Ok(result)
636 }
637 }
638
639 #[tokio::test]
640 async fn receiver_drop_does_not_prevent_completion() {
641 let hv = Box::new(TestBackend::new(GateOnPut::default()));
642 let lt = Box::new(TestBackend::new(GateOnPut::with_pause()));
643 let backend = TieredStorage::new(hv.clone(), lt.clone(), Box::new(NoopChangeLog));
644 let service = StorageService::new(Box::new(backend));
645
646 let payload = vec![0xABu8; 2 * 1024 * 1024]; let request = service.insert_object(
648 make_context(),
649 Some("completion-test".into()),
650 Metadata::default(),
651 stream::single(payload),
652 );
653
654 let paused = Arc::clone(<.hooks.paused);
657 tokio::select! {
658 _ = request => panic!("insert should not complete while backend is paused"),
659 _ = paused.notified() => {}
660 }
661
662 lt.hooks.resume.notify_one();
666
667 let on_put = Arc::clone(&hv.hooks.on_put);
670 tokio::time::timeout(Duration::from_secs(5), on_put.notified())
671 .await
672 .expect("timed out waiting for tombstone write");
673
674 let id = ObjectId::new(make_context(), "completion-test".into());
677 let tombstone = hv.inner.get(&id).expect_tombstone();
678 let lt_id = tombstone.target;
679 assert!(lt.inner.contains(<_id), "long-term object missing");
680 }
681
682 fn make_limited_service(limit: u32) -> (StorageService, TestBackend<GateOnPut>) {
685 let backend = TestBackend::new(GateOnPut::with_pause());
686 let service = StorageService::new(Box::new(backend.clone()))
687 .with_concurrency(ConcurrencyLimiter::new(limit));
688 (service, backend)
689 }
690
691 #[tokio::test]
692 async fn at_capacity_rejects() {
693 let (service, hv) = make_limited_service(1);
694
695 let svc = service.clone();
697 let first = tokio::spawn(async move {
698 svc.insert_object(
699 make_context(),
700 Some("first".into()),
701 Metadata::default(),
702 stream::single("data"),
703 )
704 .await
705 });
706
707 hv.hooks.paused.notified().await;
709
710 let result = service
712 .insert_object(
713 make_context(),
714 Some("second".into()),
715 Metadata::default(),
716 stream::single("data"),
717 )
718 .await;
719
720 assert!(
721 matches!(result, Err(Error::AtCapacity)),
722 "expected AtCapacity, got {result:?}"
723 );
724
725 hv.hooks.resume.notify_one();
727 first.await.unwrap().unwrap();
728
729 service
731 .get_metadata(ObjectId::new(make_context(), "first".into()))
732 .await
733 .unwrap();
734 }
735
736 #[tokio::test]
737 async fn tasks_limit_returns_configured_limit() {
738 let backend = Box::new(InMemoryBackend::new("cap"));
739 let service = StorageService::new(backend).with_concurrency(ConcurrencyLimiter::new(7));
740 assert_eq!(service.tasks_limit(), 7);
741 }
742
743 #[tokio::test]
744 async fn tasks_running_tracks_in_flight() {
745 let (service, hv) = make_limited_service(5);
746
747 assert_eq!(service.tasks_running(), 0);
748
749 let svc = service.clone();
751 let _blocked = tokio::spawn(async move {
752 svc.insert_object(
753 make_context(),
754 Some("in-use-test".into()),
755 Metadata::default(),
756 stream::single("data"),
757 )
758 .await
759 });
760
761 hv.hooks.paused.notified().await;
762 assert_eq!(service.tasks_running(), 1);
763
764 hv.hooks.resume.notify_one();
765 }
766
767 #[tokio::test]
768 async fn permits_released_after_panic() {
769 let service = StorageService::new(Box::new(TestBackend::new(PanicOnGet)))
770 .with_concurrency(ConcurrencyLimiter::new(1));
771
772 let id = ObjectId::new(make_context(), "panic-permit".into());
774 let result = service.get_object(id.clone(), None).await;
775 assert!(matches!(result, Err(Error::Panic(_))));
776
777 let result = service.get_object(id, None).await;
779 assert!(
780 !matches!(result, Err(Error::AtCapacity)),
781 "permit was not released after panic"
782 );
783 }
784}