1use std::future::Future;
14use std::num::NonZeroU64;
15use std::sync::Arc;
16
17use objectstore_types::metadata::Metadata;
18use objectstore_types::range::{ByteRange, ContentRange};
19use objectstore_types::resumable::{SessionToken as EncryptedSessionToken, UploadProgress};
20use objectstore_types::time::Timestamp;
21
22use crate::backend::common::{Backend, ExpiryUpdate, SetExpiryResponse};
23use crate::backend::counting::CountingBackend;
24use crate::background::RenewalScheduler;
25use crate::concurrency::ConcurrencyLimiter;
26use crate::encryption::Cipher;
27use crate::error::{ErrorKind, Result, ResultExt as _};
28use crate::id::{ObjectContext, ObjectId};
29use crate::resumable::Session;
30use crate::stream::{ClientStream, PayloadStream};
31use crate::streaming::StreamExecutor;
32
33pub type GetResponse = Option<(Metadata, Option<ContentRange>, PayloadStream)>;
35pub type MetadataResponse = Option<Metadata>;
37pub type InsertResponse = ObjectId;
39pub type DeleteResponse = ();
41
42#[derive(Debug)]
44pub struct CreateUploadSessionResponse {
45 pub session: EncryptedSessionToken,
47 pub granularity: u64,
49}
50
51pub const DEFAULT_CONCURRENCY_LIMIT: u32 = 500;
56
57pub const DEFAULT_BACKGROUND_QUEUE_LIMIT: usize = 1_000;
59
60#[derive(Clone, Debug)]
92pub struct StorageService {
93 inner: Arc<dyn Backend>,
94 concurrency: ConcurrencyLimiter,
95 renewals: RenewalScheduler,
96 cipher: Arc<Cipher>,
97}
98
99impl StorageService {
100 pub fn new(backend: Box<dyn Backend>, cipher: Cipher) -> Self {
107 let inner: Arc<dyn Backend> = Arc::new(CountingBackend::new(backend));
108 let concurrency = ConcurrencyLimiter::new(DEFAULT_CONCURRENCY_LIMIT);
109 Self {
110 inner: Arc::clone(&inner),
111 concurrency: concurrency.clone(),
112 renewals: RenewalScheduler::new(inner, concurrency, DEFAULT_BACKGROUND_QUEUE_LIMIT),
113 cipher: Arc::new(cipher),
114 }
115 }
116
117 pub fn with_concurrency(mut self, limiter: ConcurrencyLimiter) -> Self {
123 self.renewals.set_concurrency(limiter.clone());
124 self.concurrency = limiter;
125 self
126 }
127
128 pub fn with_background_queue_limit(mut self, limit: usize) -> Self {
130 self.renewals.set_capacity(limit);
131 self
132 }
133
134 pub fn concurrency_limiter(&self) -> &ConcurrencyLimiter {
136 &self.concurrency
137 }
138
139 pub fn tasks_running(&self) -> u32 {
141 self.concurrency.used_permits()
142 }
143
144 pub fn tasks_limit(&self) -> u32 {
146 self.concurrency.total_permits()
147 }
148
149 pub fn stream(&self) -> StreamExecutor {
156 StreamExecutor::new(
157 Arc::clone(&self.inner),
158 self.concurrency.clone(),
159 self.renewals.clone(),
160 )
161 }
162
163 pub fn start(&mut self) {
177 self.renewals.start();
178
179 let concurrency = self.concurrency.clone();
180 objectstore_metrics::gauge!("service.concurrency.limit" = concurrency.total_permits());
181 objectstore_metrics::gauge!("service.concurrency.queue_limit" = concurrency.total_queue());
182 objectstore_metrics::gauge!("service.concurrency.bulk_limit" = concurrency.total_bulk());
183
184 tokio::spawn(async move {
185 concurrency
186 .run_emitter(|stats| async move {
187 objectstore_metrics::gauge!("service.concurrency.in_use" = stats.in_use);
188 objectstore_metrics::gauge!("service.concurrency.queued" = stats.queued);
189 objectstore_metrics::gauge!(
190 "service.concurrency.bulk_in_use" = stats.bulk_in_use
191 );
192 })
193 .await;
194 });
195
196 let renewals = self.renewals.clone();
197 tokio::spawn(async move {
198 renewals
199 .run_emitter(|queued| async move {
200 objectstore_metrics::gauge!("service.expiry_renewal.queued" = queued);
201 })
202 .await;
203 });
204 }
205
206 async fn spawn<T, F>(&self, operation: &'static str, f: F) -> Result<T>
224 where
225 T: Send + 'static,
226 F: Future<Output = Result<T>> + Send + 'static,
227 {
228 let timer = objectstore_metrics::timer!("service.concurrency.wait");
229 let permit = self.concurrency.acquire().await.inspect_err(|_| {
230 objectstore_metrics::count!("service.concurrency.rejected", class = "normal");
231 objectstore_log::warn!("Request rejected: service at capacity");
232 })?;
233
234 timer.record();
235 crate::concurrency::run_metered(operation, permit, f).await
236 }
237
238 pub async fn insert_object(
250 &self,
251 context: ObjectContext,
252 key: Option<String>,
253 metadata: Metadata,
254 stream: ClientStream,
255 access_time: Timestamp,
256 ) -> Result<InsertResponse> {
257 metadata.validate().kind(ErrorKind::InvalidMetadata)?;
258 let id = ObjectId::optional(context, key);
259 let inner = Arc::clone(&self.inner);
260 self.spawn("insert", async move {
261 inner
262 .put_object(&id, &metadata, stream, access_time)
263 .await?;
264 Ok(id)
265 })
266 .await
267 }
268
269 pub async fn get_metadata(
271 &self,
272 id: ObjectId,
273 access_time: Timestamp,
274 ) -> Result<MetadataResponse> {
275 let inner = Arc::clone(&self.inner);
276 let renewals = self.renewals.clone();
277 self.spawn("get_metadata", async move {
278 let response = inner.get_metadata(&id, access_time).await?;
279 if let Some(ref metadata) = response
280 && let Some(expire_at) = metadata.check_tti_bump(access_time)
281 {
282 renewals.schedule(id, expire_at);
283 }
284 Ok(response)
285 })
286 .await
287 }
288
289 pub async fn get_object(
291 &self,
292 id: ObjectId,
293 access_time: Timestamp,
294 range: Option<ByteRange>,
295 ) -> Result<GetResponse> {
296 let inner = Arc::clone(&self.inner);
297 let renewals = self.renewals.clone();
298 self.spawn("get", async move {
299 let response = inner.get_object(&id, access_time, range).await?;
300 if let Some((ref metadata, _, _)) = response
301 && let Some(expire_at) = metadata.check_tti_bump(access_time)
302 {
303 renewals.schedule(id, expire_at);
304 }
305 Ok(response)
306 })
307 .await
308 }
309
310 pub async fn set_expiry(
320 &self,
321 id: ObjectId,
322 target: ExpiryUpdate,
323 access_time: Timestamp,
324 ) -> Result<SetExpiryResponse> {
325 let inner = Arc::clone(&self.inner);
326 self.spawn("set_expiry", async move {
327 inner.set_expiry(&id, target, access_time).await
328 })
329 .await
330 }
331
332 pub async fn delete_object(
340 &self,
341 id: ObjectId,
342 access_time: Timestamp,
343 ) -> Result<DeleteResponse> {
344 let inner = Arc::clone(&self.inner);
345 self.spawn("delete", async move {
346 inner.delete_object(&id, access_time).await
347 })
348 .await
349 }
350
351 pub async fn join(&self) {
357 self.renewals.join().await;
358 self.inner.join().await;
359 }
360
361 pub async fn create_upload_session(
368 &self,
369 id: ObjectId,
370 metadata: Metadata,
371 upload_length: u64,
372 ) -> Result<Option<CreateUploadSessionResponse>> {
373 let Some(upload_length) = NonZeroU64::new(upload_length) else {
374 return Ok(None);
375 };
376 metadata.validate().kind(ErrorKind::InvalidMetadata)?;
377 let inner = Arc::clone(&self.inner);
378 let cipher = Arc::clone(&self.cipher);
379 self.spawn("create_upload_session", async move {
380 let session = inner
381 .create_upload_session(&id, &metadata, upload_length)
382 .await?;
383 session
384 .map(|backend_token| {
385 let session = cipher
386 .encrypt(&Session {
387 object_id: id,
388 upload_length,
389 backend_token,
390 })
391 .map(EncryptedSessionToken::new)?;
392 Ok(CreateUploadSessionResponse {
393 session,
394 granularity: inner.upload_granularity(),
395 })
396 })
397 .transpose()
398 })
399 .await
400 }
401
402 fn session_for(&self, expected_id: &ObjectId, token: EncryptedSessionToken) -> Result<Session> {
403 let session: Session = self
404 .cipher
405 .decrypt(token.as_bytes())
406 .map_err(|_| ErrorKind::InvalidUploadSessionToken)?;
407 if session.object_id != *expected_id {
408 return Err(ErrorKind::UnknownUploadSession.into());
409 }
410 Ok(session)
411 }
412
413 pub async fn put_chunk(
422 &self,
423 id: ObjectId,
424 token: EncryptedSessionToken,
425 offset: u64,
426 content_length: u64,
427 body: ClientStream,
428 ) -> Result<UploadProgress> {
429 let session = self.session_for(&id, token)?;
430 let inner = Arc::clone(&self.inner);
431 self.spawn("put_chunk", async move {
432 inner
433 .put_chunk(&session, offset, content_length, body)
434 .await
435 })
436 .await
437 }
438
439 pub async fn upload_offset(
445 &self,
446 id: ObjectId,
447 token: EncryptedSessionToken,
448 ) -> Result<UploadProgress> {
449 let session = self.session_for(&id, token)?;
450 let inner = Arc::clone(&self.inner);
451 self.spawn("upload_offset", async move {
452 inner.upload_offset(&session).await
453 })
454 .await
455 }
456
457 pub async fn cancel_upload(&self, id: ObjectId, token: EncryptedSessionToken) -> Result<()> {
459 let session = self.session_for(&id, token)?;
460 let inner = Arc::clone(&self.inner);
461 self.spawn("cancel_upload", async move {
462 inner.cancel_upload(&session).await
463 })
464 .await
465 }
466}
467
468#[cfg(test)]
469mod tests {
470 use std::error::Error as _;
471 use std::sync::atomic::{AtomicUsize, Ordering};
472 use std::sync::{Arc, Mutex};
473 use std::time::Duration;
474
475 use bytes::BytesMut;
476 use futures_util::TryStreamExt;
477 use objectstore_types::metadata::{ExpirationPolicy, Metadata};
478 use objectstore_types::range::ByteRange;
479 use objectstore_types::scope::{Scope, Scopes};
480
481 use super::*;
482 use crate::backend::bigtable::{BigTableBackend, BigTableConfig};
483 use crate::backend::changelog::NoopChangeLog;
484 use crate::backend::common::{ExpiryTarget, HighVolumeBackend, PutResponse, TieredWrite};
485 use crate::backend::gcs::{GcsBackend, GcsConfig};
486 use crate::backend::in_memory::InMemoryBackend;
487 use crate::backend::testing::{Hooks, TestBackend};
488 use crate::backend::tiered::TieredStorage;
489 use crate::change_stream::ChangeStreamFactory;
490 use crate::resumable::BackendToken;
491 use crate::stream::{self, ClientStream};
492
493 #[derive(Clone, Debug, Default)]
494 struct ResumableTokenHooks {
495 seen_tokens: Arc<Mutex<Vec<String>>>,
496 }
497
498 #[async_trait::async_trait]
499 impl Hooks for ResumableTokenHooks {
500 fn upload_granularity(&self, _inner: &InMemoryBackend) -> u64 {
501 256 * 1024
502 }
503
504 async fn create_upload_session(
505 &self,
506 _inner: &InMemoryBackend,
507 _id: &ObjectId,
508 _metadata: &Metadata,
509 _upload_length: NonZeroU64,
510 ) -> Result<Option<BackendToken>> {
511 Ok(Some("backend token".to_owned()))
512 }
513
514 async fn upload_offset(
515 &self,
516 _inner: &InMemoryBackend,
517 session: &Session,
518 ) -> Result<UploadProgress> {
519 assert_eq!(session.upload_length.get(), 4);
520 self.seen_tokens
521 .lock()
522 .unwrap()
523 .push(session.backend_token.clone());
524 Ok(UploadProgress::Incomplete { offset: 0 })
525 }
526 }
527
528 fn make_context() -> ObjectContext {
529 ObjectContext {
530 usecase: "testing".parse().unwrap(),
531 scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
532 }
533 }
534
535 fn make_service() -> StorageService {
536 StorageService::new(
537 Box::new(InMemoryBackend::new("in-memory")),
538 Cipher::ephemeral().unwrap(),
539 )
540 }
541
542 #[tokio::test]
543 async fn insert_without_key_generates_unique_id() {
544 let service = make_service();
545
546 let id = service
547 .insert_object(
548 make_context(),
549 None,
550 Metadata::default(),
551 stream::single("auto-keyed"),
552 Timestamp::now(),
553 )
554 .await
555 .unwrap();
556
557 assert!(uuid::Uuid::parse_str(id.key()).is_ok());
558 }
559
560 #[tokio::test]
561 async fn stores_files() {
562 let service = make_service();
563
564 let key = service
565 .insert_object(
566 make_context(),
567 Some("testing".into()),
568 Metadata::default(),
569 stream::single("oh hai!"),
570 Timestamp::now(),
571 )
572 .await
573 .unwrap();
574
575 let (_metadata, _, stream) = service
576 .get_object(key, Timestamp::now(), None)
577 .await
578 .unwrap()
579 .unwrap();
580 let file_contents: BytesMut = stream.try_collect().await.unwrap();
581
582 assert_eq!(file_contents.as_ref(), b"oh hai!");
583 }
584
585 #[tokio::test]
586 async fn works_with_gcs() {
587 let config = GcsConfig {
588 endpoint: Some("http://localhost:8087".into()),
589 bucket: "test-bucket".into(), tcp_user_timeout: Duration::from_secs(30),
591 cogs: None,
592 };
593
594 let backend = GcsBackend::new(config, &ChangeStreamFactory::default())
595 .await
596 .unwrap();
597 let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());
598
599 let key = service
600 .insert_object(
601 make_context(),
602 Some("testing".into()),
603 Metadata::default(),
604 stream::single("oh hai!"),
605 Timestamp::now(),
606 )
607 .await
608 .unwrap();
609
610 let (_metadata, _, stream) = service
611 .get_object(key, Timestamp::now(), None)
612 .await
613 .unwrap()
614 .unwrap();
615 let file_contents: BytesMut = stream.try_collect().await.unwrap();
616
617 assert_eq!(file_contents.as_ref(), b"oh hai!");
618 }
619
620 #[tokio::test]
621 async fn tombstone_redirect_and_delete() {
622 let bigtable_config = BigTableConfig {
623 endpoint: Some("localhost:8086".into()),
624 project_id: "testing".into(),
625 instance_name: "objectstore".into(),
626 table_name: "objectstore".into(),
627 connections: None,
628 rpc_timeout: Duration::from_secs(2),
629 cogs: None,
630 };
631 let gcs_config = GcsConfig {
632 endpoint: Some("http://localhost:8087".into()),
633 bucket: "test-bucket".into(),
634 tcp_user_timeout: Duration::from_secs(30),
635 cogs: None,
636 };
637
638 let high_volume = Box::new(
639 BigTableBackend::new(bigtable_config, &ChangeStreamFactory::default())
640 .await
641 .unwrap(),
642 );
643 let long_term = Box::new(
644 GcsBackend::new(gcs_config.clone(), &ChangeStreamFactory::default())
645 .await
646 .unwrap(),
647 );
648 let backend = TieredStorage::new(high_volume, long_term, Box::new(NoopChangeLog));
649 let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());
650
651 let gcs_backend = GcsBackend::new(gcs_config.clone(), &ChangeStreamFactory::default())
653 .await
654 .unwrap();
655
656 let payload_len = 2 * 1024 * 1024;
659 let payload = vec![0xAB; payload_len]; let id = service
661 .insert_object(
662 make_context(),
663 Some("delete-cleanup-test".into()),
664 Metadata::default(),
665 stream::single(payload),
666 Timestamp::now(),
667 )
668 .await
669 .unwrap();
670
671 let (_, _, stream) = service
673 .get_object(id.clone(), Timestamp::now(), None)
674 .await
675 .unwrap()
676 .unwrap();
677 let body: BytesMut = stream.try_collect().await.unwrap();
678 assert_eq!(body.len(), payload_len);
679
680 service
682 .delete_object(id.clone(), Timestamp::now())
683 .await
684 .unwrap();
685
686 let after_delete = service
688 .get_object(id.clone(), Timestamp::now(), None)
689 .await
690 .unwrap();
691 assert!(after_delete.is_none(), "tombstone not deleted");
692
693 let orphan = gcs_backend
695 .get_object(&id, Timestamp::now(), None)
696 .await
697 .unwrap();
698 assert!(orphan.is_none(), "object leaked");
699 }
700
701 #[tokio::test]
704 async fn basic_spawn_insert_and_get() {
705 let service = make_service();
706
707 let id = service
708 .insert_object(
709 make_context(),
710 Some("test-key".into()),
711 Metadata::default(),
712 stream::single("hello world"),
713 Timestamp::now(),
714 )
715 .await
716 .unwrap();
717
718 let (_, _, stream) = service
719 .get_object(id, Timestamp::now(), None)
720 .await
721 .unwrap()
722 .unwrap();
723 let body: BytesMut = stream.try_collect().await.unwrap();
724 assert_eq!(body.as_ref(), b"hello world");
725 }
726
727 #[tokio::test]
728 async fn basic_spawn_metadata_and_delete() {
729 let service = make_service();
730
731 let id = service
732 .insert_object(
733 make_context(),
734 Some("meta-key".into()),
735 Metadata::default(),
736 stream::single("data"),
737 Timestamp::now(),
738 )
739 .await
740 .unwrap();
741
742 let metadata = service
743 .get_metadata(id.clone(), Timestamp::now())
744 .await
745 .unwrap();
746 assert!(metadata.is_some());
747
748 service
749 .delete_object(id.clone(), Timestamp::now())
750 .await
751 .unwrap();
752
753 let after = service
754 .get_object(id, Timestamp::now(), None)
755 .await
756 .unwrap();
757 assert!(after.is_none());
758 }
759
760 #[tokio::test]
761 async fn set_expiry() {
762 let service = make_service();
763 let old_expiry = Timestamp::now() + Duration::from_hours(1);
764 let metadata = Metadata {
765 expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
766 time_expires: Some(old_expiry),
767 ..Default::default()
768 };
769 let id = service
770 .insert_object(
771 make_context(),
772 Some("explicit-expiry".into()),
773 metadata,
774 stream::single("payload"),
775 Timestamp::now(),
776 )
777 .await
778 .unwrap();
779 let requested = old_expiry + Duration::from_hours(1);
780
781 assert_eq!(
782 service
783 .set_expiry(
784 id.clone(),
785 ExpiryTarget::At(requested).into(),
786 Timestamp::now()
787 )
788 .await
789 .unwrap(),
790 SetExpiryResponse::Satisfied(requested)
791 );
792 assert_eq!(
793 service
794 .get_metadata(id, Timestamp::now())
795 .await
796 .unwrap()
797 .unwrap()
798 .time_expires,
799 Some(requested)
800 );
801 }
802
803 #[derive(Debug)]
804 struct PanicOnGet;
805
806 #[async_trait::async_trait]
807 impl Hooks for PanicOnGet {
808 async fn get_object(
809 &self,
810 _inner: &InMemoryBackend,
811 _id: &ObjectId,
812 _access_time: Timestamp,
813 _range: Option<ByteRange>,
814 ) -> Result<GetResponse> {
815 panic!("intentional panic in get_object");
816 }
817 }
818
819 #[tokio::test]
820 async fn panic_in_backend_returns_task_failed() {
821 let service = StorageService::new(
822 Box::new(TestBackend::new(PanicOnGet)),
823 Cipher::ephemeral().unwrap(),
824 );
825
826 let id = ObjectId::new(make_context(), "panic-test".into());
827 let result = service.get_object(id, Timestamp::now(), None).await;
828
829 let Err(error) = result else {
830 panic!("expected Panic error");
831 };
832 assert_eq!(error.kind(), ErrorKind::Panic);
833 assert_eq!(error.to_string(), "service task panicked");
834 assert_eq!(
835 std::error::Error::source(&error).unwrap().to_string(),
836 "intentional panic in get_object"
837 );
838 }
839
840 #[derive(Clone, Debug, Default)]
841 struct GateOnExpiry {
842 calls: Arc<AtomicUsize>,
843 access_time: Arc<Mutex<Option<Timestamp>>>,
844 started: Arc<tokio::sync::Notify>,
845 resume: Arc<tokio::sync::Notify>,
846 }
847
848 #[async_trait::async_trait]
849 impl Hooks for GateOnExpiry {
850 async fn set_expiry(
851 &self,
852 inner: &InMemoryBackend,
853 id: &ObjectId,
854 target: ExpiryUpdate,
855 access_time: Timestamp,
856 ) -> Result<SetExpiryResponse> {
857 *self.access_time.lock().unwrap() = Some(access_time);
858 self.calls.fetch_add(1, Ordering::SeqCst);
859 self.started.notify_one();
860 self.resume.notified().await;
861 inner.set_expiry(id, target, access_time).await
862 }
863 }
864
865 fn stale_tti_metadata() -> Metadata {
866 Metadata {
867 expiration_policy: ExpirationPolicy::TimeToIdle(Duration::from_hours(1)),
868 time_expires: Some(Timestamp::now() + Duration::from_mins(1)),
869 ..Default::default()
870 }
871 }
872
873 #[tokio::test]
874 async fn background_renewal() {
875 let now = Timestamp::now();
876 let access_time = now - Duration::from_secs(10);
877 let backend = TestBackend::new(GateOnExpiry::default());
878 let id = ObjectId::new(make_context(), "background-renewal".into());
879 let metadata = stale_tti_metadata();
880 backend
881 .inner
882 .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
883 .await
884 .unwrap();
885 let mut service =
886 StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
887 service.start();
888
889 let response = tokio::time::timeout(
890 Duration::from_secs(1),
891 service.get_object(id.clone(), access_time, Some(ByteRange::Bounded(0, 2))),
892 )
893 .await
894 .expect("GET waited for its background renewal")
895 .unwrap()
896 .unwrap();
897 assert_eq!(response.0.time_expires, metadata.time_expires);
898 backend.hooks.started.notified().await;
899 assert!(backend.hooks.access_time.lock().unwrap().unwrap() >= now);
900
901 let join = tokio::spawn({
902 let service = service.clone();
903 async move { service.join().await }
904 });
905 tokio::pin!(join);
906 assert!(
907 tokio::time::timeout(Duration::from_millis(25), &mut join)
908 .await
909 .is_err()
910 );
911
912 backend.hooks.resume.notify_waiters();
913 tokio::time::timeout(Duration::from_secs(1), &mut join)
914 .await
915 .expect("shutdown did not drain renewal")
916 .unwrap();
917 assert_eq!(
918 backend.inner.get(&id).expect_object().0.time_expires,
919 Some(access_time + Duration::from_hours(1))
920 );
921 }
922
923 #[tokio::test]
924 async fn renewal_deduplication() {
925 let backend = TestBackend::new(GateOnExpiry::default());
926 let id = ObjectId::new(make_context(), "deduplicated-renewal".into());
927 backend
928 .inner
929 .put_object(
930 &id,
931 &stale_tti_metadata(),
932 stream::single("payload"),
933 Timestamp::now(),
934 )
935 .await
936 .unwrap();
937 let mut service =
938 StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
939 service.start();
940
941 service
942 .get_metadata(id.clone(), Timestamp::now())
943 .await
944 .unwrap();
945 backend.hooks.started.notified().await;
946 service.get_metadata(id, Timestamp::now()).await.unwrap();
947 tokio::task::yield_now().await;
948 assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 1);
949
950 backend.hooks.resume.notify_waiters();
951 service.join().await;
952 }
953
954 #[tokio::test]
955 async fn ttl_read() {
956 let backend = TestBackend::new(GateOnExpiry::default());
957 let id = ObjectId::new(make_context(), "ttl-no-renewal".into());
958 let metadata = Metadata {
959 expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
960 time_expires: Some(Timestamp::now() + Duration::from_mins(1)),
961 ..Default::default()
962 };
963 backend
964 .inner
965 .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
966 .await
967 .unwrap();
968 let mut service =
969 StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
970 service.start();
971
972 service.get_metadata(id, Timestamp::now()).await.unwrap();
973 tokio::task::yield_now().await;
974 assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 0);
975 service.join().await;
976 }
977
978 #[tokio::test]
979 async fn renewal_queueing() {
980 let backend = TestBackend::new(GateOnExpiry::default());
981 let first = ObjectId::new(make_context(), "first-renewal".into());
982 let second = ObjectId::new(make_context(), "queued-renewal".into());
983 let metadata = stale_tti_metadata();
984 for id in [&first, &second] {
985 backend
986 .inner
987 .put_object(id, &metadata, stream::single("payload"), Timestamp::now())
988 .await
989 .unwrap();
990 }
991
992 let concurrency = ConcurrencyLimiter::new(1);
993 let mut scheduler = RenewalScheduler::new(Arc::new(backend.clone()), concurrency, 1);
994 let expire_at = Timestamp::now() + Duration::from_hours(1);
995 scheduler.schedule(first, expire_at);
996 assert_eq!(scheduler.queued(), 1);
997 scheduler.start();
998 backend.hooks.started.notified().await;
999 assert_eq!(scheduler.queued(), 0);
1000
1001 scheduler.schedule(second.clone(), expire_at);
1002 tokio::task::yield_now().await;
1003 assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 1);
1004 assert_eq!(
1005 backend.inner.get(&second).expect_object().0.time_expires,
1006 metadata.time_expires
1007 );
1008
1009 backend.hooks.resume.notify_waiters();
1010 tokio::time::timeout(Duration::from_secs(1), async {
1011 while backend.hooks.calls.load(Ordering::SeqCst) < 2 {
1012 tokio::task::yield_now().await;
1013 }
1014 })
1015 .await
1016 .expect("queued renewal did not start");
1017 backend.hooks.resume.notify_waiters();
1018 scheduler.join().await;
1019 assert_eq!(scheduler.queued(), 0);
1020 assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 2);
1021 assert!(backend.inner.get(&second).expect_object().0.time_expires > metadata.time_expires);
1022 }
1023
1024 #[derive(Clone, Debug, Default)]
1025 struct FailFirstExpiry {
1026 calls: Arc<AtomicUsize>,
1027 panic: bool,
1028 }
1029
1030 #[async_trait::async_trait]
1031 impl Hooks for FailFirstExpiry {
1032 async fn set_expiry(
1033 &self,
1034 inner: &InMemoryBackend,
1035 id: &ObjectId,
1036 target: ExpiryUpdate,
1037 access_time: Timestamp,
1038 ) -> Result<SetExpiryResponse> {
1039 if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
1040 assert!(!self.panic, "intentional renewal panic");
1041 return Err(ErrorKind::BackendFailure.into());
1042 }
1043 inner.set_expiry(id, target, access_time).await
1044 }
1045 }
1046
1047 #[tokio::test]
1048 async fn renewal_failure_cleanup() {
1049 for panic in [false, true] {
1050 let backend = TestBackend::new(FailFirstExpiry {
1051 panic,
1052 ..Default::default()
1053 });
1054 let id = ObjectId::new(make_context(), "failed-renewal".into());
1055 let metadata = stale_tti_metadata();
1056 backend
1057 .inner
1058 .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
1059 .await
1060 .unwrap();
1061 let concurrency = ConcurrencyLimiter::new(1);
1062 let mut scheduler = RenewalScheduler::new(Arc::new(backend.clone()), concurrency, 1);
1063 scheduler.start();
1064 let expire_at = metadata.check_tti_bump(Timestamp::now()).unwrap();
1065
1066 scheduler.schedule(id.clone(), expire_at);
1067 tokio::time::timeout(Duration::from_secs(1), async {
1068 while scheduler.pending() != 0 {
1069 tokio::task::yield_now().await;
1070 }
1071 })
1072 .await
1073 .expect("failure did not release renewal guards");
1074
1075 scheduler.schedule(id, expire_at);
1076 scheduler.join().await;
1077 assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 2);
1078 }
1079 }
1080
1081 #[derive(Clone, Debug, Default)]
1085 struct GateOnPut {
1086 pause: bool,
1087 paused: Arc<tokio::sync::Notify>,
1088 resume: Arc<tokio::sync::Notify>,
1089 on_put: Arc<tokio::sync::Notify>,
1090 }
1091
1092 impl GateOnPut {
1093 fn with_pause() -> Self {
1094 Self {
1095 pause: true,
1096 ..Default::default()
1097 }
1098 }
1099 }
1100
1101 #[async_trait::async_trait]
1102 impl Hooks for GateOnPut {
1103 async fn put_object(
1104 &self,
1105 inner: &InMemoryBackend,
1106 id: &ObjectId,
1107 metadata: &Metadata,
1108 stream: ClientStream,
1109 access_time: Timestamp,
1110 ) -> Result<PutResponse> {
1111 if self.pause {
1112 self.paused.notify_one();
1113 self.resume.notified().await;
1114 }
1115 inner.put_object(id, metadata, stream, access_time).await?;
1116 self.on_put.notify_one();
1117 Ok(())
1118 }
1119
1120 async fn compare_and_write(
1121 &self,
1122 inner: &InMemoryBackend,
1123 id: &ObjectId,
1124 current: Option<&ObjectId>,
1125 write: TieredWrite,
1126 access_time: Timestamp,
1127 ) -> Result<bool> {
1128 let notify = matches!(write, TieredWrite::Tombstone(_) | TieredWrite::Object(_, _));
1129 let result = inner
1130 .compare_and_write(id, current, write, access_time)
1131 .await?;
1132 if notify {
1133 self.on_put.notify_one();
1134 }
1135 Ok(result)
1136 }
1137 }
1138
1139 #[tokio::test]
1140 async fn receiver_drop_does_not_prevent_completion() {
1141 let hv = Box::new(TestBackend::new(GateOnPut::default()));
1142 let lt = Box::new(TestBackend::new(GateOnPut::with_pause()));
1143 let backend = TieredStorage::new(hv.clone(), lt.clone(), Box::new(NoopChangeLog));
1144 let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());
1145
1146 let payload = vec![0xABu8; 2 * 1024 * 1024]; let request = service.insert_object(
1148 make_context(),
1149 Some("completion-test".into()),
1150 Metadata::default(),
1151 stream::single(payload),
1152 Timestamp::now(),
1153 );
1154
1155 let paused = Arc::clone(<.hooks.paused);
1158 tokio::select! {
1159 _ = request => panic!("insert should not complete while backend is paused"),
1160 _ = paused.notified() => {}
1161 }
1162
1163 lt.hooks.resume.notify_one();
1167
1168 let on_put = Arc::clone(&hv.hooks.on_put);
1171 tokio::time::timeout(Duration::from_secs(5), on_put.notified())
1172 .await
1173 .expect("timed out waiting for tombstone write");
1174
1175 let id = ObjectId::new(make_context(), "completion-test".into());
1178 let tombstone = hv.inner.get(&id).expect_tombstone();
1179 let lt_id = tombstone.target;
1180 assert!(lt.inner.contains(<_id), "long-term object missing");
1181 }
1182
1183 fn make_limited_service(limit: u32) -> (StorageService, TestBackend<GateOnPut>) {
1186 let backend = TestBackend::new(GateOnPut::with_pause());
1187 let service = StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap())
1188 .with_concurrency(ConcurrencyLimiter::new(limit));
1189 (service, backend)
1190 }
1191
1192 #[tokio::test]
1193 async fn at_capacity_rejects() {
1194 let (service, hv) = make_limited_service(1);
1195
1196 let svc = service.clone();
1198 let first = tokio::spawn(async move {
1199 svc.insert_object(
1200 make_context(),
1201 Some("first".into()),
1202 Metadata::default(),
1203 stream::single("data"),
1204 Timestamp::now(),
1205 )
1206 .await
1207 });
1208
1209 hv.hooks.paused.notified().await;
1211
1212 let result = service
1214 .insert_object(
1215 make_context(),
1216 Some("second".into()),
1217 Metadata::default(),
1218 stream::single("data"),
1219 Timestamp::now(),
1220 )
1221 .await;
1222
1223 assert!(
1224 result
1225 .as_ref()
1226 .is_err_and(|error| error.kind() == ErrorKind::AtCapacity),
1227 "expected AtCapacity, got {result:?}"
1228 );
1229
1230 hv.hooks.resume.notify_one();
1232 first.await.unwrap().unwrap();
1233
1234 service
1236 .get_metadata(
1237 ObjectId::new(make_context(), "first".into()),
1238 Timestamp::now(),
1239 )
1240 .await
1241 .unwrap();
1242 }
1243
1244 #[tokio::test]
1245 async fn tasks_limit_returns_configured_limit() {
1246 let backend = Box::new(InMemoryBackend::new("cap"));
1247 let service = StorageService::new(backend, Cipher::ephemeral().unwrap())
1248 .with_concurrency(ConcurrencyLimiter::new(7));
1249 assert_eq!(service.tasks_limit(), 7);
1250 }
1251
1252 #[tokio::test]
1253 async fn tasks_running_tracks_in_flight() {
1254 let (service, hv) = make_limited_service(5);
1255
1256 assert_eq!(service.tasks_running(), 0);
1257
1258 let svc = service.clone();
1260 let _blocked = tokio::spawn(async move {
1261 svc.insert_object(
1262 make_context(),
1263 Some("in-use-test".into()),
1264 Metadata::default(),
1265 stream::single("data"),
1266 Timestamp::now(),
1267 )
1268 .await
1269 });
1270
1271 hv.hooks.paused.notified().await;
1272 assert_eq!(service.tasks_running(), 1);
1273
1274 hv.hooks.resume.notify_one();
1275 }
1276
1277 #[tokio::test]
1278 async fn permits_released_after_panic() {
1279 let service = StorageService::new(
1280 Box::new(TestBackend::new(PanicOnGet)),
1281 Cipher::ephemeral().unwrap(),
1282 )
1283 .with_concurrency(ConcurrencyLimiter::new(1));
1284
1285 let id = ObjectId::new(make_context(), "panic-permit".into());
1287 let result = service.get_object(id.clone(), Timestamp::now(), None).await;
1288 assert!(result.is_err_and(|error| error.kind() == ErrorKind::Panic));
1289
1290 let result = service.get_object(id, Timestamp::now(), None).await;
1292 assert!(
1293 !result.is_err_and(|error| error.kind() == ErrorKind::AtCapacity),
1294 "permit was not released after panic"
1295 );
1296 }
1297
1298 #[tokio::test]
1301 async fn resumable_round_trip() {
1302 let service = make_service();
1303 let id = ObjectId::new(make_context(), "resumable".into());
1304 let created = service
1305 .create_upload_session(id.clone(), Metadata::default(), 3)
1306 .await
1307 .unwrap()
1308 .unwrap();
1309 assert_eq!(created.granularity, 0);
1310 let token = created.session;
1311 assert_eq!(
1312 service
1313 .put_chunk(id.clone(), token.clone(), 0, 1, stream::single("a"))
1314 .await
1315 .unwrap(),
1316 UploadProgress::Incomplete { offset: 1 }
1317 );
1318 assert_eq!(
1319 service
1320 .upload_offset(id.clone(), token.clone())
1321 .await
1322 .unwrap(),
1323 UploadProgress::Incomplete { offset: 1 }
1324 );
1325 assert_eq!(
1326 service
1327 .put_chunk(id.clone(), token, 1, 2, stream::single("bc"))
1328 .await
1329 .unwrap(),
1330 UploadProgress::Complete
1331 );
1332 let (metadata, _, body) = service
1333 .get_object(id, Timestamp::now(), None)
1334 .await
1335 .unwrap()
1336 .unwrap();
1337 assert_eq!(metadata.size, Some(3));
1338 assert_eq!(
1339 body.try_collect::<BytesMut>().await.unwrap().as_ref(),
1340 b"abc"
1341 );
1342 }
1343
1344 #[tokio::test]
1345 async fn resumable_create_declines_zero_length() {
1346 let service = StorageService::new(
1347 Box::new(TestBackend::new(ResumableTokenHooks::default())),
1348 Cipher::ephemeral().unwrap(),
1349 );
1350 let id = ObjectId::new(make_context(), "resumable".into());
1351
1352 let result = service
1353 .create_upload_session(id, Metadata::default(), 0)
1354 .await;
1355
1356 assert!(matches!(result, Ok(None)), "{result:?}");
1357 }
1358
1359 #[tokio::test]
1360 async fn resumable_create_validates_metadata() {
1361 let service = make_service();
1362 let id = ObjectId::new(make_context(), "resumable".into());
1363
1364 let metadata = Metadata {
1367 expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_secs(60)),
1368 ..Default::default()
1369 };
1370
1371 let result = service.create_upload_session(id, metadata, 1024).await;
1372 assert!(result.is_err_and(|error| error.kind() == ErrorKind::InvalidMetadata));
1373 }
1374
1375 #[tokio::test]
1376 async fn resumable_tokens_are_encrypted_by_default() -> Result<()> {
1377 let hooks = ResumableTokenHooks::default();
1378 let service = StorageService::new(
1379 Box::new(TestBackend::new(hooks.clone())),
1380 Cipher::ephemeral().unwrap(),
1381 );
1382 let id = ObjectId::new(make_context(), "resumable".into());
1383
1384 let created = service
1385 .create_upload_session(id.clone(), Metadata::default(), 4)
1386 .await?
1387 .expect("test backend supports resumable uploads");
1388 assert_eq!(created.granularity, 256 * 1024);
1389 let token = created.session;
1390 assert_ne!(token.as_bytes(), b"backend token");
1391 assert!(matches!(
1392 service
1393 .upload_offset(id.clone(), EncryptedSessionToken::new(b"backend token"))
1394 .await,
1395 Err(error) if error.kind() == ErrorKind::InvalidUploadSessionToken
1396 ));
1397 let other_service = StorageService::new(
1398 Box::new(TestBackend::new(hooks.clone())),
1399 Cipher::ephemeral().unwrap(),
1400 );
1401 let error = other_service
1402 .upload_offset(id.clone(), token.clone())
1403 .await
1404 .unwrap_err();
1405 assert_eq!(error.kind(), ErrorKind::InvalidUploadSessionToken);
1406 assert!(error.source().is_none());
1407 let other_id = ObjectId::new(make_context(), "other".into());
1408 assert!(matches!(
1409 service.upload_offset(other_id, token.clone()).await,
1410 Err(error) if error.kind() == ErrorKind::UnknownUploadSession
1411 ));
1412 assert!(hooks.seen_tokens.lock().unwrap().is_empty());
1413 service.upload_offset(id, token).await?;
1414 assert_eq!(
1415 hooks.seen_tokens.lock().unwrap().as_slice(),
1416 &["backend token"]
1417 );
1418 Ok(())
1419 }
1420
1421 #[tokio::test]
1422 async fn configured_encryption_only_crosses_the_service_boundary() -> Result<()> {
1423 let hooks = ResumableTokenHooks::default();
1424 let encryption = Cipher::new(
1425 "v1",
1426 std::collections::BTreeMap::from([("v1".into(), vec![7; 32])]),
1427 )
1428 .unwrap();
1429 let service = StorageService::new(Box::new(TestBackend::new(hooks.clone())), encryption);
1430 let id = ObjectId::new(make_context(), "resumable".into());
1431
1432 let encrypted = service
1433 .create_upload_session(id.clone(), Metadata::default(), 4)
1434 .await?
1435 .expect("test backend supports resumable uploads")
1436 .session;
1437 assert_ne!(encrypted.as_bytes(), b"backend token");
1438 service.upload_offset(id, encrypted).await?;
1439 assert_eq!(
1440 hooks.seen_tokens.lock().unwrap().as_slice(),
1441 &["backend token"]
1442 );
1443 Ok(())
1444 }
1445
1446 #[tokio::test]
1447 async fn configured_encryption_rejects_plaintext_tokens() {
1448 let hooks = ResumableTokenHooks::default();
1449 let encryption = Cipher::new(
1450 "v1",
1451 std::collections::BTreeMap::from([("v1".into(), vec![7; 32])]),
1452 )
1453 .unwrap();
1454 let service = StorageService::new(Box::new(TestBackend::new(hooks.clone())), encryption);
1455 let id = ObjectId::new(make_context(), "resumable".into());
1456
1457 let result = service
1458 .upload_offset(id, EncryptedSessionToken::new(b"backend token"))
1459 .await;
1460 let error = result.unwrap_err();
1461 assert_eq!(error.kind(), ErrorKind::InvalidUploadSessionToken);
1462 assert_eq!(
1463 error.to_string(),
1464 "invalid upload session token: could not decrypt or authenticate it"
1465 );
1466 assert!(error.source().is_none());
1467 assert!(hooks.seen_tokens.lock().unwrap().is_empty());
1468 }
1469}