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::{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)]
75pub struct StorageService {
76 inner: Arc<dyn Backend>,
77 concurrency: ConcurrencyLimiter,
78}
79
80impl StorageService {
81 pub fn new(backend: Box<dyn Backend>) -> Self {
88 Self {
89 inner: Arc::new(CountingBackend::new(backend)),
90 concurrency: ConcurrencyLimiter::new(DEFAULT_CONCURRENCY_LIMIT),
91 }
92 }
93
94 pub fn with_concurrency_limit(mut self, max: u32) -> Self {
99 self.concurrency = ConcurrencyLimiter::new(max);
100 self
101 }
102
103 pub fn tasks_available(&self) -> u32 {
105 self.concurrency.available_permits()
106 }
107
108 pub fn tasks_running(&self) -> u32 {
110 self.concurrency.used_permits()
111 }
112
113 pub fn tasks_limit(&self) -> u32 {
115 self.concurrency.total_permits()
116 }
117
118 pub fn stream(&self) -> Result<StreamExecutor> {
125 let available = self.tasks_available();
126 let window = available.div_ceil(10);
127
128 let acquire_result = match window {
129 0 => Err(Error::AtCapacity),
130 _ => self.concurrency.try_acquire_many(window),
131 };
132 let reservation = acquire_result.inspect_err(|_| {
133 objectstore_metrics::count!("service.concurrency.rejected");
134 objectstore_log::warn!("Request rejected: service at capacity");
135 })?;
136
137 Ok(StreamExecutor {
138 backend: Arc::clone(&self.inner),
139 window,
140 reservation,
141 })
142 }
143
144 pub fn start(&self) {
149 let concurrency = self.concurrency.clone();
150 let limit = concurrency.total_permits();
151 tokio::spawn(async move {
152 concurrency
153 .run_emitter(|permits| async move {
154 objectstore_metrics::gauge!("service.concurrency.in_use" = permits);
155 objectstore_metrics::gauge!("service.concurrency.limit" = limit);
156 })
157 .await;
158 });
159 }
160
161 async fn spawn<T, F>(&self, operation: &'static str, f: F) -> Result<T>
173 where
174 T: Send + 'static,
175 F: Future<Output = Result<T>> + Send + 'static,
176 {
177 let permit = self.concurrency.try_acquire().inspect_err(|_| {
178 objectstore_metrics::count!("service.concurrency.rejected");
179 objectstore_log::warn!("Request rejected: service at capacity");
180 })?;
181
182 crate::concurrency::spawn_metered(operation, permit, f).await
183 }
184
185 pub async fn insert_object(
197 &self,
198 context: ObjectContext,
199 key: Option<String>,
200 metadata: Metadata,
201 stream: ClientStream,
202 ) -> Result<InsertResponse> {
203 metadata.validate()?;
204 let id = ObjectId::optional(context, key);
205 let inner = Arc::clone(&self.inner);
206 self.spawn("insert", async move {
207 inner.put_object(&id, &metadata, stream).await?;
208 Ok(id)
209 })
210 .await
211 }
212
213 pub async fn get_metadata(&self, id: ObjectId) -> Result<MetadataResponse> {
215 let inner = Arc::clone(&self.inner);
216 self.spawn("get_metadata", async move { inner.get_metadata(&id).await })
217 .await
218 }
219
220 pub async fn get_object(&self, id: ObjectId, range: Option<ByteRange>) -> Result<GetResponse> {
222 let inner = Arc::clone(&self.inner);
223 self.spawn("get", async move { inner.get_object(&id, range).await })
224 .await
225 }
226
227 pub async fn delete_object(&self, id: ObjectId) -> Result<DeleteResponse> {
235 let inner = Arc::clone(&self.inner);
236 self.spawn("delete", async move { inner.delete_object(&id).await })
237 .await
238 }
239
240 pub async fn join(&self) {
246 self.inner.join().await;
247 }
248
249 pub async fn initiate_multipart(
253 &self,
254 id: ObjectId,
255 metadata: Metadata,
256 ) -> Result<InitiateMultipartResponse> {
257 metadata.validate()?;
258 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
260 self.spawn("initiate_multipart", async move {
261 inner
262 .as_multipart_upload_backend()?
263 .initiate_multipart(&id, &metadata)
264 .await
265 })
266 .await
267 }
268
269 pub async fn upload_part(
278 &self,
279 id: ObjectId,
280 upload_id: UploadId,
281 part_number: PartNumber,
282 content_length: u64,
283 content_md5: Option<String>,
284 body: ClientStream,
285 ) -> Result<UploadPartResponse> {
286 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
288 self.spawn("upload_part", async move {
289 inner
290 .as_multipart_upload_backend()?
291 .upload_part(
292 &id,
293 &upload_id,
294 part_number,
295 content_length,
296 content_md5.as_deref(),
297 body,
298 )
299 .await
300 })
301 .await
302 }
303
304 pub async fn list_parts(
306 &self,
307 id: ObjectId,
308 upload_id: UploadId,
309 max_parts: Option<u32>,
310 part_number_marker: Option<PartNumber>,
311 ) -> Result<ListPartsResponse> {
312 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
314 self.spawn("list_parts", async move {
315 inner
316 .as_multipart_upload_backend()?
317 .list_parts(&id, &upload_id, max_parts, part_number_marker)
318 .await
319 })
320 .await
321 }
322
323 pub async fn abort_multipart(
325 &self,
326 id: ObjectId,
327 upload_id: UploadId,
328 ) -> Result<AbortMultipartResponse> {
329 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
331 self.spawn("abort_multipart", async move {
332 inner
333 .as_multipart_upload_backend()?
334 .abort_multipart(&id, &upload_id)
335 .await
336 })
337 .await
338 }
339
340 pub async fn complete_multipart(
342 &self,
343 id: ObjectId,
344 upload_id: UploadId,
345 parts: Vec<CompletedPart>,
346 ) -> Result<CompleteMultipartResponse> {
347 self.inner.as_multipart_upload_backend()?; let inner = self.inner.clone();
349 self.spawn("complete_multipart", async move {
350 inner
351 .as_multipart_upload_backend()?
352 .complete_multipart(&id, &upload_id, parts)
353 .await
354 })
355 .await
356 }
357}
358
359#[cfg(test)]
360mod tests {
361 use std::sync::Arc;
362 use std::time::Duration;
363
364 use bytes::BytesMut;
365 use futures_util::TryStreamExt;
366 use objectstore_types::metadata::Metadata;
367 use objectstore_types::range::ByteRange;
368 use objectstore_types::scope::{Scope, Scopes};
369
370 use super::*;
371 use crate::backend::bigtable::{BigTableBackend, BigTableConfig};
372 use crate::backend::changelog::NoopChangeLog;
373 use crate::backend::common::{HighVolumeBackend, PutResponse, TieredWrite};
374 use crate::backend::gcs::{GcsBackend, GcsConfig};
375 use crate::backend::in_memory::InMemoryBackend;
376 use crate::backend::testing::{Hooks, TestBackend};
377 use crate::backend::tiered::TieredStorage;
378 use crate::error::Error;
379 use crate::stream::{self, ClientStream};
380
381 fn make_context() -> ObjectContext {
382 ObjectContext {
383 usecase: "testing".into(),
384 scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
385 }
386 }
387
388 fn make_service() -> StorageService {
389 StorageService::new(Box::new(InMemoryBackend::new("in-memory")))
390 }
391
392 #[tokio::test]
393 async fn insert_without_key_generates_unique_id() {
394 let service = make_service();
395
396 let id = service
397 .insert_object(
398 make_context(),
399 None,
400 Metadata::default(),
401 stream::single("auto-keyed"),
402 )
403 .await
404 .unwrap();
405
406 assert!(uuid::Uuid::parse_str(id.key()).is_ok());
407 }
408
409 #[tokio::test]
410 async fn stores_files() {
411 let service = make_service();
412
413 let key = service
414 .insert_object(
415 make_context(),
416 Some("testing".into()),
417 Metadata::default(),
418 stream::single("oh hai!"),
419 )
420 .await
421 .unwrap();
422
423 let (_metadata, _, stream) = service.get_object(key, None).await.unwrap().unwrap();
424 let file_contents: BytesMut = stream.try_collect().await.unwrap();
425
426 assert_eq!(file_contents.as_ref(), b"oh hai!");
427 }
428
429 #[tokio::test]
430 async fn works_with_gcs() {
431 let config = GcsConfig {
432 endpoint: Some("http://localhost:8087".into()),
433 bucket: "test-bucket".into(), };
435
436 let backend = GcsBackend::new(config).await.unwrap();
437 let service = StorageService::new(Box::new(backend));
438
439 let key = service
440 .insert_object(
441 make_context(),
442 Some("testing".into()),
443 Metadata::default(),
444 stream::single("oh hai!"),
445 )
446 .await
447 .unwrap();
448
449 let (_metadata, _, stream) = service.get_object(key, None).await.unwrap().unwrap();
450 let file_contents: BytesMut = stream.try_collect().await.unwrap();
451
452 assert_eq!(file_contents.as_ref(), b"oh hai!");
453 }
454
455 #[tokio::test]
456 async fn tombstone_redirect_and_delete() {
457 let bigtable_config = BigTableConfig {
458 endpoint: Some("localhost:8086".into()),
459 project_id: "testing".into(),
460 instance_name: "objectstore".into(),
461 table_name: "objectstore".into(),
462 connections: None,
463 };
464 let gcs_config = GcsConfig {
465 endpoint: Some("http://localhost:8087".into()),
466 bucket: "test-bucket".into(),
467 };
468
469 let high_volume = Box::new(BigTableBackend::new(bigtable_config).await.unwrap());
470 let long_term = Box::new(GcsBackend::new(gcs_config.clone()).await.unwrap());
471 let backend = TieredStorage::new(high_volume, long_term, Box::new(NoopChangeLog));
472 let service = StorageService::new(Box::new(backend));
473
474 let gcs_backend = GcsBackend::new(gcs_config.clone()).await.unwrap();
476
477 let payload_len = 2 * 1024 * 1024;
480 let payload = vec![0xAB; payload_len]; let id = service
482 .insert_object(
483 make_context(),
484 Some("delete-cleanup-test".into()),
485 Metadata::default(),
486 stream::single(payload),
487 )
488 .await
489 .unwrap();
490
491 let (_, _, stream) = service.get_object(id.clone(), None).await.unwrap().unwrap();
493 let body: BytesMut = stream.try_collect().await.unwrap();
494 assert_eq!(body.len(), payload_len);
495
496 service.delete_object(id.clone()).await.unwrap();
498
499 let after_delete = service.get_object(id.clone(), None).await.unwrap();
501 assert!(after_delete.is_none(), "tombstone not deleted");
502
503 let orphan = gcs_backend.get_object(&id, None).await.unwrap();
505 assert!(orphan.is_none(), "object leaked");
506 }
507
508 #[tokio::test]
511 async fn basic_spawn_insert_and_get() {
512 let service = make_service();
513
514 let id = service
515 .insert_object(
516 make_context(),
517 Some("test-key".into()),
518 Metadata::default(),
519 stream::single("hello world"),
520 )
521 .await
522 .unwrap();
523
524 let (_, _, stream) = service.get_object(id, None).await.unwrap().unwrap();
525 let body: BytesMut = stream.try_collect().await.unwrap();
526 assert_eq!(body.as_ref(), b"hello world");
527 }
528
529 #[tokio::test]
530 async fn basic_spawn_metadata_and_delete() {
531 let service = make_service();
532
533 let id = service
534 .insert_object(
535 make_context(),
536 Some("meta-key".into()),
537 Metadata::default(),
538 stream::single("data"),
539 )
540 .await
541 .unwrap();
542
543 let metadata = service.get_metadata(id.clone()).await.unwrap();
544 assert!(metadata.is_some());
545
546 service.delete_object(id.clone()).await.unwrap();
547
548 let after = service.get_object(id, None).await.unwrap();
549 assert!(after.is_none());
550 }
551
552 #[derive(Debug)]
553 struct PanicOnGet;
554
555 #[async_trait::async_trait]
556 impl Hooks for PanicOnGet {
557 async fn get_object(
558 &self,
559 _inner: &InMemoryBackend,
560 _id: &ObjectId,
561 _range: Option<ByteRange>,
562 ) -> Result<GetResponse> {
563 panic!("intentional panic in get_object");
564 }
565 }
566
567 #[tokio::test]
568 async fn panic_in_backend_returns_task_failed() {
569 let service = StorageService::new(Box::new(TestBackend::new(PanicOnGet)));
570
571 let id = ObjectId::new(make_context(), "panic-test".into());
572 let result = service.get_object(id, None).await;
573
574 let Err(Error::Panic(msg)) = result else {
575 panic!("expected Panic error");
576 };
577 assert!(msg.contains("intentional panic in get_object"), "{msg}");
578 }
579
580 #[derive(Clone, Debug, Default)]
584 struct GateOnPut {
585 pause: bool,
586 paused: Arc<tokio::sync::Notify>,
587 resume: Arc<tokio::sync::Notify>,
588 on_put: Arc<tokio::sync::Notify>,
589 }
590
591 impl GateOnPut {
592 fn with_pause() -> Self {
593 Self {
594 pause: true,
595 ..Default::default()
596 }
597 }
598 }
599
600 #[async_trait::async_trait]
601 impl Hooks for GateOnPut {
602 async fn put_object(
603 &self,
604 inner: &InMemoryBackend,
605 id: &ObjectId,
606 metadata: &Metadata,
607 stream: ClientStream,
608 ) -> Result<PutResponse> {
609 if self.pause {
610 self.paused.notify_one();
611 self.resume.notified().await;
612 }
613 inner.put_object(id, metadata, stream).await?;
614 self.on_put.notify_one();
615 Ok(())
616 }
617
618 async fn compare_and_write(
619 &self,
620 inner: &InMemoryBackend,
621 id: &ObjectId,
622 current: Option<&ObjectId>,
623 write: TieredWrite,
624 ) -> Result<bool> {
625 let notify = matches!(write, TieredWrite::Tombstone(_) | TieredWrite::Object(_, _));
626 let result = inner.compare_and_write(id, current, write).await?;
627 if notify {
628 self.on_put.notify_one();
629 }
630 Ok(result)
631 }
632 }
633
634 #[tokio::test]
635 async fn receiver_drop_does_not_prevent_completion() {
636 let hv = Box::new(TestBackend::new(GateOnPut::default()));
637 let lt = Box::new(TestBackend::new(GateOnPut::with_pause()));
638 let backend = TieredStorage::new(hv.clone(), lt.clone(), Box::new(NoopChangeLog));
639 let service = StorageService::new(Box::new(backend));
640
641 let payload = vec![0xABu8; 2 * 1024 * 1024]; let request = service.insert_object(
643 make_context(),
644 Some("completion-test".into()),
645 Metadata::default(),
646 stream::single(payload),
647 );
648
649 let paused = Arc::clone(<.hooks.paused);
652 tokio::select! {
653 _ = request => panic!("insert should not complete while backend is paused"),
654 _ = paused.notified() => {}
655 }
656
657 lt.hooks.resume.notify_one();
661
662 let on_put = Arc::clone(&hv.hooks.on_put);
665 tokio::time::timeout(Duration::from_secs(5), on_put.notified())
666 .await
667 .expect("timed out waiting for tombstone write");
668
669 let id = ObjectId::new(make_context(), "completion-test".into());
672 let tombstone = hv.inner.get(&id).expect_tombstone();
673 let lt_id = tombstone.target;
674 assert!(lt.inner.contains(<_id), "long-term object missing");
675 }
676
677 fn make_limited_service(limit: u32) -> (StorageService, TestBackend<GateOnPut>) {
680 let backend = TestBackend::new(GateOnPut::with_pause());
681 let service = StorageService::new(Box::new(backend.clone())).with_concurrency_limit(limit);
682 (service, backend)
683 }
684
685 #[tokio::test]
686 async fn at_capacity_rejects() {
687 let (service, hv) = make_limited_service(1);
688
689 let svc = service.clone();
691 let first = tokio::spawn(async move {
692 svc.insert_object(
693 make_context(),
694 Some("first".into()),
695 Metadata::default(),
696 stream::single("data"),
697 )
698 .await
699 });
700
701 hv.hooks.paused.notified().await;
703
704 let result = service
706 .insert_object(
707 make_context(),
708 Some("second".into()),
709 Metadata::default(),
710 stream::single("data"),
711 )
712 .await;
713
714 assert!(
715 matches!(result, Err(Error::AtCapacity)),
716 "expected AtCapacity, got {result:?}"
717 );
718
719 hv.hooks.resume.notify_one();
721 first.await.unwrap().unwrap();
722
723 service
725 .get_metadata(ObjectId::new(make_context(), "first".into()))
726 .await
727 .unwrap();
728 }
729
730 #[tokio::test]
731 async fn tasks_limit_returns_configured_limit() {
732 let backend = Box::new(InMemoryBackend::new("cap"));
733 let service = StorageService::new(backend).with_concurrency_limit(7);
734 assert_eq!(service.tasks_limit(), 7);
735 }
736
737 #[tokio::test]
738 async fn tasks_running_tracks_in_flight() {
739 let (service, hv) = make_limited_service(5);
740
741 assert_eq!(service.tasks_running(), 0);
742
743 let svc = service.clone();
745 let _blocked = tokio::spawn(async move {
746 svc.insert_object(
747 make_context(),
748 Some("in-use-test".into()),
749 Metadata::default(),
750 stream::single("data"),
751 )
752 .await
753 });
754
755 hv.hooks.paused.notified().await;
756 assert_eq!(service.tasks_running(), 1);
757
758 hv.hooks.resume.notify_one();
759 }
760
761 #[tokio::test]
762 async fn permits_released_after_panic() {
763 let service =
764 StorageService::new(Box::new(TestBackend::new(PanicOnGet))).with_concurrency_limit(1);
765
766 let id = ObjectId::new(make_context(), "panic-permit".into());
768 let result = service.get_object(id.clone(), None).await;
769 assert!(matches!(result, Err(Error::Panic(_))));
770
771 let result = service.get_object(id, None).await;
773 assert!(
774 !matches!(result, Err(Error::AtCapacity)),
775 "permit was not released after panic"
776 );
777 }
778}