objectstore_service/

service.rs

//! Core storage service and configuration.
//!
//! [`StorageService`] is the main entry point for storing and retrieving
//! objects. Each operation runs in a separate tokio task for panic isolation.
//!
//! Callers supply an access timestamp, normally the HTTP request start time.
//! It stays fixed across admission, backend calls, and TTI calculations. Initial
//! creation and expiry are already resolved in the supplied metadata. Background
//! renewals retain the read-derived deadline but check liveness at worker start.
//!
//! See the [crate-level documentation](crate) for full architecture details.

use std::future::Future;
use std::num::NonZeroU64;
use std::sync::Arc;

use objectstore_types::metadata::Metadata;
use objectstore_types::range::{ByteRange, ContentRange};
use objectstore_types::resumable::{SessionToken as EncryptedSessionToken, UploadProgress};
use objectstore_types::time::Timestamp;

use crate::backend::common::Backend;
use crate::backend::counting::CountingBackend;
use crate::background::RenewalScheduler;
use crate::concurrency::ConcurrencyLimiter;
use crate::encryption::Cipher;
use crate::error::{ErrorKind, Result, ResultExt as _};
use crate::id::{ObjectContext, ObjectId};
use crate::multipart::{
    AbortMultipartResponse, CompleteMultipartResponse, CompletedPart, InitiateMultipartResponse,
    ListPartsResponse, PartNumber, UploadId, UploadPartResponse,
};
use crate::resumable::{BackendToken, SessionToken};
use crate::stream::{ClientStream, PayloadStream};
use crate::streaming::StreamExecutor;

/// Service response for [`StorageService::get_object`].
pub type GetResponse = Option<(Metadata, Option<ContentRange>, PayloadStream)>;
/// Service response for [`StorageService::get_metadata`].
pub type MetadataResponse = Option<Metadata>;
/// Service response for [`StorageService::insert_object`].
pub type InsertResponse = ObjectId;
/// Service response for [`StorageService::delete_object`].
pub type DeleteResponse = ();

/// Default concurrency limit for [`StorageService`].
///
/// This value is used when no explicit limiter is set via
/// [`StorageService::with_concurrency`].
pub const DEFAULT_CONCURRENCY_LIMIT: u32 = 500;

/// Default number of TTI renewals that may wait for background processing.
pub const DEFAULT_BACKGROUND_QUEUE_LIMIT: usize = 1_000;

/// Asynchronous storage service wrapping a single [`Backend`].
///
/// `StorageService` is the main entry point for storing and retrieving objects.
/// It delegates all storage operations to the backend supplied at construction,
/// adding task spawning, panic isolation, and a concurrency limit on top.
///
/// The typical backend is [`TieredStorage`](crate::backend::tiered::TieredStorage),
/// which provides size-based routing to high-volume and long-term backends along
/// with redirect tombstone management. Any type implementing [`Backend`] can be used.
///
/// # Lifecycle
///
/// After construction, call [`start`](StorageService::start) to start the
/// service's background processes.
///
/// # Run-to-Completion and Panic Isolation
///
/// Each operation runs to completion even if the caller is cancelled (e.g., on
/// client disconnect). This ensures that multi-step operations in the backend
/// are never left partially applied. Post-commit cleanup (e.g. deleting
/// unreferenced long-term blobs) runs in background tasks so callers are not
/// blocked. Call [`join`](StorageService::join) during shutdown to wait for
/// outstanding cleanup. Operations are also isolated from panics in backend
/// code — a failure in one operation does not bring down other in-flight work.
///
/// # Concurrency Limit
///
/// A [`ConcurrencyLimiter`] caps the number of in-flight backend operations.
/// Pass a custom limiter via
/// [`with_concurrency`](StorageService::with_concurrency); without one the
/// default is [`DEFAULT_CONCURRENCY_LIMIT`] permits with no queue.
#[derive(Clone, Debug)]
pub struct StorageService {
    inner: Arc<dyn Backend>,
    concurrency: ConcurrencyLimiter,
    renewals: RenewalScheduler,
    cipher: Arc<Cipher>,
}

impl StorageService {
    /// Creates a new `StorageService` wrapping the given backend.
    ///
    /// The backend is wrapped in a [`CountingBackend`] which increments a COGS usage counter for
    /// each operation run. Single-object operations served directly by `StorageService` are covered
    /// as we batched operations served by [`StreamExecutor`]. See
    /// [`backend::counting`](crate::backend::counting) for details.
    pub fn new(backend: Box<dyn Backend>, cipher: Cipher) -> Self {
        let inner: Arc<dyn Backend> = Arc::new(CountingBackend::new(backend));
        let concurrency = ConcurrencyLimiter::new(DEFAULT_CONCURRENCY_LIMIT);
        Self {
            inner: Arc::clone(&inner),
            concurrency: concurrency.clone(),
            renewals: RenewalScheduler::new(inner, concurrency, DEFAULT_BACKGROUND_QUEUE_LIMIT),
            cipher: Arc::new(cipher),
        }
    }

    /// Replaces the default concurrency limiter.
    ///
    /// Must be called before [`start`](Self::start). Without this, the
    /// service uses a limiter with [`DEFAULT_CONCURRENCY_LIMIT`] permits
    /// and no queue.
    pub fn with_concurrency(mut self, limiter: ConcurrencyLimiter) -> Self {
        self.renewals.set_concurrency(limiter.clone());
        self.concurrency = limiter;
        self
    }

    /// Replaces the default background expiry-renewal queue capacity.
    pub fn with_background_queue_limit(mut self, limit: usize) -> Self {
        self.renewals.set_capacity(limit);
        self
    }

    /// Returns a reference to the concurrency limiter.
    pub fn concurrency_limiter(&self) -> &ConcurrencyLimiter {
        &self.concurrency
    }

    /// Returns the number of backend tasks currently running.
    pub fn tasks_running(&self) -> u32 {
        self.concurrency.used_permits()
    }

    /// Returns the configured limit for concurrent backend tasks.
    pub fn tasks_limit(&self) -> u32 {
        self.concurrency.total_permits()
    }

    /// Prepares to stream multiple operations concurrently against this service.
    ///
    /// Each operation acquires a bulk permit individually via
    /// [`ConcurrencyLimiter::acquire_bulk`], which caps bulk traffic at a
    /// configurable percentage of execution slots while allowing operations
    /// to queue for permits instead of requiring upfront reservation.
    pub fn stream(&self) -> StreamExecutor {
        StreamExecutor::new(
            Arc::clone(&self.inner),
            self.concurrency.clone(),
            self.renewals.clone(),
        )
    }

    /// Starts background processes for the storage service.
    ///
    /// At startup, this tracks the following gauges:
    ///
    ///  - `service.concurrency.limit`: concurrent task execution slots
    ///  - `service.concurrency.queue_limit`: queue size for waiting tasks
    ///  - `service.concurrency.bulk_limit`: concurrent task execution slots for bulk operations
    ///
    /// Also spawns tasks that emit runtime gauges once per second:
    ///  - `service.concurrency.in_use`: currently running tasks
    ///  - `service.concurrency.queued`: currently queued tasks
    ///  - `service.concurrency.bulk_in_use`: currently running bulk tasks
    ///  - `service.expiry_renewal.queued`: expiry renewals waiting for the background worker
    pub fn start(&mut self) {
        self.renewals.start();

        let concurrency = self.concurrency.clone();
        objectstore_metrics::gauge!("service.concurrency.limit" = concurrency.total_permits());
        objectstore_metrics::gauge!("service.concurrency.queue_limit" = concurrency.total_queue());
        objectstore_metrics::gauge!("service.concurrency.bulk_limit" = concurrency.total_bulk());

        tokio::spawn(async move {
            concurrency
                .run_emitter(|stats| async move {
                    objectstore_metrics::gauge!("service.concurrency.in_use" = stats.in_use);
                    objectstore_metrics::gauge!("service.concurrency.queued" = stats.queued);
                    objectstore_metrics::gauge!(
                        "service.concurrency.bulk_in_use" = stats.bulk_in_use
                    );
                })
                .await;
        });

        let renewals = self.renewals.clone();
        tokio::spawn(async move {
            renewals
                .run_emitter(|queued| async move {
                    objectstore_metrics::gauge!("service.expiry_renewal.queued" = queued);
                })
                .await;
        });
    }

    /// Spawns a future in a separate task and awaits its result.
    ///
    /// # Observability
    ///
    /// This tracks two metrics:
    ///
    /// - `service.task.start` (counter) after acquiring a permit
    /// - `service.task.duration` (distribution) when the task completes
    ///
    /// Both are tagged with the given `operation` name and an `outcome`
    /// of `"success"` or `"error"`.
    ///
    /// # Errors
    ///
    /// - `AtCapacity` if the concurrency limit is reached
    /// - `Panic` if the spawned task panics (the panic message is captured for diagnostics)
    /// - `Dropped` if the task is dropped before sending its result.
    async fn spawn<T, F>(&self, operation: &'static str, f: F) -> Result<T>
    where
        T: Send + 'static,
        F: Future<Output = Result<T>> + Send + 'static,
    {
        let timer = objectstore_metrics::timer!("service.concurrency.wait");
        let permit = self.concurrency.acquire().await.inspect_err(|_| {
            objectstore_metrics::count!("service.concurrency.rejected", class = "normal");
            objectstore_log::warn!("Request rejected: service at capacity");
        })?;

        timer.record();
        crate::concurrency::run_metered(operation, permit, f).await
    }

    /// Creates or overwrites an object.
    ///
    /// The object is identified by the components of an [`ObjectId`]. The
    /// `context` is required, while the `key` can be assigned automatically if
    /// set to `None`.
    ///
    /// # Run-to-completion
    ///
    /// Once called, the operation runs to completion even if the returned future
    /// is dropped (e.g., on client disconnect). This guarantees that partially
    /// written objects in the backend are never left in an inconsistent state.
    pub async fn insert_object(
        &self,
        context: ObjectContext,
        key: Option<String>,
        metadata: Metadata,
        stream: ClientStream,
        access_time: Timestamp,
    ) -> Result<InsertResponse> {
        metadata.validate().kind(ErrorKind::InvalidMetadata)?;
        let id = ObjectId::optional(context, key);
        let inner = Arc::clone(&self.inner);
        self.spawn("insert", async move {
            inner
                .put_object(&id, &metadata, stream, access_time)
                .await?;
            Ok(id)
        })
        .await
    }

    /// Retrieves only the metadata for an object, without the payload.
    pub async fn get_metadata(
        &self,
        id: ObjectId,
        access_time: Timestamp,
    ) -> Result<MetadataResponse> {
        let inner = Arc::clone(&self.inner);
        let renewals = self.renewals.clone();
        self.spawn("get_metadata", async move {
            let response = inner.get_metadata(&id, access_time).await?;
            if let Some(ref metadata) = response
                && let Some(expire_at) = metadata.check_tti_bump(access_time)
            {
                renewals.schedule(id, expire_at);
            }
            Ok(response)
        })
        .await
    }

    /// Streams (part of) the contents of an object.
    pub async fn get_object(
        &self,
        id: ObjectId,
        access_time: Timestamp,
        range: Option<ByteRange>,
    ) -> Result<GetResponse> {
        let inner = Arc::clone(&self.inner);
        let renewals = self.renewals.clone();
        self.spawn("get", async move {
            let response = inner.get_object(&id, access_time, range).await?;
            if let Some((ref metadata, _, _)) = response
                && let Some(expire_at) = metadata.check_tti_bump(access_time)
            {
                renewals.schedule(id, expire_at);
            }
            Ok(response)
        })
        .await
    }

    /// Extends an existing TTL or TTI object's deadline.
    pub async fn set_expiry(
        &self,
        id: ObjectId,
        expire_at: Timestamp,
        access_time: Timestamp,
    ) -> Result<bool> {
        let inner = Arc::clone(&self.inner);
        self.spawn("set_expiry", async move {
            inner.set_expiry(&id, expire_at, access_time).await
        })
        .await
    }

    /// Deletes an object, if it exists.
    ///
    /// # Run-to-completion
    ///
    /// Once called, the operation runs to completion even if the returned future
    /// is dropped. This guarantees that multi-step delete sequences in the backend
    /// are never left partially applied.
    pub async fn delete_object(
        &self,
        id: ObjectId,
        access_time: Timestamp,
    ) -> Result<DeleteResponse> {
        let inner = Arc::clone(&self.inner);
        self.spawn("delete", async move {
            inner.delete_object(&id, access_time).await
        })
        .await
    }

    /// Waits for all outstanding background operations to complete.
    ///
    /// Blocks until any pending background cleanup tasks finish, up to the
    /// backend's configured timeout. Should be called during graceful shutdown
    /// after the HTTP server has stopped accepting new requests.
    pub async fn join(&self) {
        self.renewals.join().await;
        self.inner.join().await;
    }

    // --- Multipart upload operations ---

    /// Initiates a new multipart upload.
    pub async fn initiate_multipart(
        &self,
        id: ObjectId,
        metadata: Metadata,
    ) -> Result<InitiateMultipartResponse> {
        metadata.validate().kind(ErrorKind::InvalidMetadata)?;
        self.inner.as_multipart_upload_backend()?; // Fail before clone/spawn if unsupported
        let inner = self.inner.clone();
        self.spawn("initiate_multipart", async move {
            inner
                .as_multipart_upload_backend()?
                .initiate_multipart(&id, &metadata)
                .await
        })
        .await
    }

    /// Uploads a single part.
    ///
    /// Note that this requires a `content_length`.
    /// This grants us the broadest and most seamless compatibility when it comes to backends.
    /// For example, MinIO rejects `UploadPart` requests without a `Content-Length` on plain PUT
    /// requests.
    /// This can be worked around by using AWS SigV4 chunked streaming requests, which we could use
    /// if one day we'll have a usecase where the client doesn't know the part length upfront.
    pub async fn upload_part(
        &self,
        id: ObjectId,
        upload_id: UploadId,
        part_number: PartNumber,
        content_length: u64,
        content_md5: Option<String>,
        body: ClientStream,
    ) -> Result<UploadPartResponse> {
        self.inner.as_multipart_upload_backend()?; // Fail before clone/spawn if unsupported
        let inner = self.inner.clone();
        self.spawn("upload_part", async move {
            inner
                .as_multipart_upload_backend()?
                .upload_part(
                    &id,
                    &upload_id,
                    part_number,
                    content_length,
                    content_md5.as_deref(),
                    body,
                )
                .await
        })
        .await
    }

    /// Lists the parts uploaded so far.
    pub async fn list_parts(
        &self,
        id: ObjectId,
        upload_id: UploadId,
        max_parts: Option<u32>,
        part_number_marker: Option<PartNumber>,
    ) -> Result<ListPartsResponse> {
        self.inner.as_multipart_upload_backend()?; // Fail before clone/spawn if unsupported
        let inner = self.inner.clone();
        self.spawn("list_parts", async move {
            inner
                .as_multipart_upload_backend()?
                .list_parts(&id, &upload_id, max_parts, part_number_marker)
                .await
        })
        .await
    }

    /// Aborts a multipart upload.
    pub async fn abort_multipart(
        &self,
        id: ObjectId,
        upload_id: UploadId,
    ) -> Result<AbortMultipartResponse> {
        self.inner.as_multipart_upload_backend()?; // Fail before clone/spawn if unsupported
        let inner = self.inner.clone();
        self.spawn("abort_multipart", async move {
            inner
                .as_multipart_upload_backend()?
                .abort_multipart(&id, &upload_id)
                .await
        })
        .await
    }

    /// Finalizes a multipart upload.
    pub async fn complete_multipart(
        &self,
        id: ObjectId,
        upload_id: UploadId,
        parts: Vec<CompletedPart>,
        access_time: Timestamp,
    ) -> Result<CompleteMultipartResponse> {
        self.inner.as_multipart_upload_backend()?; // Fail before clone/spawn if unsupported
        let inner = self.inner.clone();
        self.spawn("complete_multipart", async move {
            inner
                .as_multipart_upload_backend()?
                .complete_multipart(&id, &upload_id, parts, access_time)
                .await
        })
        .await
    }

    // --- Resumable upload operations ---

    /// Opens a resumable upload session for an object of `total_length` bytes.
    ///
    /// Returns `Ok(None)` for zero-length objects or when the backend declines resumable uploads
    /// for this object, in which case the caller should fall back to [`Self::insert_object`].
    pub async fn create_upload_session(
        &self,
        id: ObjectId,
        metadata: Metadata,
        total_length: u64,
    ) -> Result<Option<EncryptedSessionToken>> {
        let Some(total_length) = NonZeroU64::new(total_length) else {
            return Ok(None);
        };
        metadata.validate().kind(ErrorKind::InvalidMetadata)?;
        let inner = Arc::clone(&self.inner);
        let cipher = Arc::clone(&self.cipher);
        self.spawn("create_upload_session", async move {
            let session = inner
                .create_upload_session(&id, &metadata, total_length)
                .await?;
            session
                .map(|backend_token| {
                    cipher
                        .encrypt(&SessionToken {
                            object_id: id,
                            backend_token,
                        })
                        .map(EncryptedSessionToken::new)
                })
                .transpose()
        })
        .await
    }

    fn backend_token_for(
        &self,
        expected_id: &ObjectId,
        token: EncryptedSessionToken,
    ) -> Result<BackendToken> {
        let session: SessionToken = self
            .cipher
            .decrypt(token.as_bytes())
            .map_err(|_| ErrorKind::UnknownUploadSession)?;
        if session.object_id != *expected_id {
            return Err(ErrorKind::UnknownUploadSession.into());
        }
        Ok(session.backend_token)
    }

    /// Writes a chunk of `content_length` bytes at `offset` into an open session.
    ///
    /// Completes the upload once the chunk carrying the last byte is persisted.
    ///
    /// # Run-to-completion
    ///
    /// Once called, the operation runs to completion even if the returned future is dropped.
    /// This matters most for the final chunk, which completes the upload.
    pub async fn put_chunk(
        &self,
        id: ObjectId,
        token: EncryptedSessionToken,
        offset: u64,
        content_length: u64,
        body: ClientStream,
    ) -> Result<UploadProgress> {
        let session = self.backend_token_for(&id, token)?;
        let inner = Arc::clone(&self.inner);
        self.spawn("put_chunk", async move {
            inner
                .put_chunk(&id, &session, offset, content_length, body)
                .await
        })
        .await
    }

    /// Reports how far a session has progressed.
    ///
    /// This can observe completion after the final chunk's response was lost. A composed backend
    /// may also finish pending publication work, so this requires write permission at the API
    /// layer.
    pub async fn upload_offset(
        &self,
        id: ObjectId,
        token: EncryptedSessionToken,
    ) -> Result<UploadProgress> {
        let session = self.backend_token_for(&id, token)?;
        let inner = Arc::clone(&self.inner);
        self.spawn("upload_offset", async move {
            inner.upload_offset(&id, &session).await
        })
        .await
    }

    /// Cancels an upload session, discarding whatever was uploaded.
    pub async fn cancel_upload(&self, id: ObjectId, token: EncryptedSessionToken) -> Result<()> {
        let session = self.backend_token_for(&id, token)?;
        let inner = Arc::clone(&self.inner);
        self.spawn("cancel_upload", async move {
            inner.cancel_upload(&id, &session).await
        })
        .await
    }
}

#[cfg(test)]
mod tests {
    use std::error::Error as _;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::{Arc, Mutex};
    use std::time::Duration;

    use bytes::BytesMut;
    use futures_util::TryStreamExt;
    use objectstore_types::metadata::{ExpirationPolicy, Metadata};
    use objectstore_types::range::ByteRange;
    use objectstore_types::scope::{Scope, Scopes};

    use super::*;
    use crate::backend::bigtable::{BigTableBackend, BigTableConfig};
    use crate::backend::changelog::NoopChangeLog;
    use crate::backend::common::{HighVolumeBackend, PutResponse, TieredWrite};
    use crate::backend::gcs::{GcsBackend, GcsConfig};
    use crate::backend::in_memory::InMemoryBackend;
    use crate::backend::testing::{Hooks, TestBackend};
    use crate::backend::tiered::TieredStorage;
    use crate::change_stream::ChangeStreamFactory;
    use crate::stream::{self, ClientStream};

    #[derive(Clone, Debug, Default)]
    struct ResumableTokenHooks {
        seen_tokens: Arc<Mutex<Vec<String>>>,
    }

    #[async_trait::async_trait]
    impl Hooks for ResumableTokenHooks {
        async fn create_upload_session(
            &self,
            _inner: &InMemoryBackend,
            _id: &ObjectId,
            _metadata: &Metadata,
            _total_length: NonZeroU64,
        ) -> Result<Option<BackendToken>> {
            Ok(Some("backend token".to_owned()))
        }

        async fn upload_offset(
            &self,
            _inner: &InMemoryBackend,
            _id: &ObjectId,
            token: &BackendToken,
        ) -> Result<UploadProgress> {
            self.seen_tokens.lock().unwrap().push(token.to_owned());
            Ok(UploadProgress::Incomplete { offset: 0 })
        }
    }

    fn make_context() -> ObjectContext {
        ObjectContext {
            usecase: "testing".into(),
            scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
        }
    }

    fn make_service() -> StorageService {
        StorageService::new(
            Box::new(InMemoryBackend::new("in-memory")),
            Cipher::ephemeral().unwrap(),
        )
    }

    #[tokio::test]
    async fn insert_without_key_generates_unique_id() {
        let service = make_service();

        let id = service
            .insert_object(
                make_context(),
                None,
                Metadata::default(),
                stream::single("auto-keyed"),
                Timestamp::now(),
            )
            .await
            .unwrap();

        assert!(uuid::Uuid::parse_str(id.key()).is_ok());
    }

    #[tokio::test]
    async fn stores_files() {
        let service = make_service();

        let key = service
            .insert_object(
                make_context(),
                Some("testing".into()),
                Metadata::default(),
                stream::single("oh hai!"),
                Timestamp::now(),
            )
            .await
            .unwrap();

        let (_metadata, _, stream) = service
            .get_object(key, Timestamp::now(), None)
            .await
            .unwrap()
            .unwrap();
        let file_contents: BytesMut = stream.try_collect().await.unwrap();

        assert_eq!(file_contents.as_ref(), b"oh hai!");
    }

    #[tokio::test]
    async fn works_with_gcs() {
        let config = GcsConfig {
            endpoint: Some("http://localhost:8087".into()),
            bucket: "test-bucket".into(), // aligned with the env var in devservices and CI
            cogs: None,
        };

        let backend = GcsBackend::new(config, &ChangeStreamFactory::default())
            .await
            .unwrap();
        let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());

        let key = service
            .insert_object(
                make_context(),
                Some("testing".into()),
                Metadata::default(),
                stream::single("oh hai!"),
                Timestamp::now(),
            )
            .await
            .unwrap();

        let (_metadata, _, stream) = service
            .get_object(key, Timestamp::now(), None)
            .await
            .unwrap()
            .unwrap();
        let file_contents: BytesMut = stream.try_collect().await.unwrap();

        assert_eq!(file_contents.as_ref(), b"oh hai!");
    }

    #[tokio::test]
    async fn tombstone_redirect_and_delete() {
        let bigtable_config = BigTableConfig {
            endpoint: Some("localhost:8086".into()),
            project_id: "testing".into(),
            instance_name: "objectstore".into(),
            table_name: "objectstore".into(),
            connections: None,
            rpc_timeout: Duration::from_secs(2),
            cogs: None,
        };
        let gcs_config = GcsConfig {
            endpoint: Some("http://localhost:8087".into()),
            bucket: "test-bucket".into(),
            cogs: None,
        };

        let high_volume = Box::new(
            BigTableBackend::new(bigtable_config, &ChangeStreamFactory::default())
                .await
                .unwrap(),
        );
        let long_term = Box::new(
            GcsBackend::new(gcs_config.clone(), &ChangeStreamFactory::default())
                .await
                .unwrap(),
        );
        let backend = TieredStorage::new(high_volume, long_term, Box::new(NoopChangeLog));
        let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());

        // A separate GCS backend to directly inspect the long-term storage.
        let gcs_backend = GcsBackend::new(gcs_config.clone(), &ChangeStreamFactory::default())
            .await
            .unwrap();

        // Insert a >1 MiB object with a key.  This forces the long-term path:
        // the real payload goes to GCS, and a redirect tombstone is written to BigTable.
        let payload_len = 2 * 1024 * 1024;
        let payload = vec![0xAB; payload_len]; // 2 MiB
        let id = service
            .insert_object(
                make_context(),
                Some("delete-cleanup-test".into()),
                Metadata::default(),
                stream::single(payload),
                Timestamp::now(),
            )
            .await
            .unwrap();

        // Sanity: the object is readable through the service (follows the tombstone).
        let (_, _, stream) = service
            .get_object(id.clone(), Timestamp::now(), None)
            .await
            .unwrap()
            .unwrap();
        let body: BytesMut = stream.try_collect().await.unwrap();
        assert_eq!(body.len(), payload_len);

        // Delete through the service layer.
        service
            .delete_object(id.clone(), Timestamp::now())
            .await
            .unwrap();

        // The tombstone in BigTable should be gone, so the service returns None.
        let after_delete = service
            .get_object(id.clone(), Timestamp::now(), None)
            .await
            .unwrap();
        assert!(after_delete.is_none(), "tombstone not deleted");

        // The real object in GCS must also be gone — no orphan.
        let orphan = gcs_backend
            .get_object(&id, Timestamp::now(), None)
            .await
            .unwrap();
        assert!(orphan.is_none(), "object leaked");
    }

    // --- Task spawning tests (public API) ---

    #[tokio::test]
    async fn basic_spawn_insert_and_get() {
        let service = make_service();

        let id = service
            .insert_object(
                make_context(),
                Some("test-key".into()),
                Metadata::default(),
                stream::single("hello world"),
                Timestamp::now(),
            )
            .await
            .unwrap();

        let (_, _, stream) = service
            .get_object(id, Timestamp::now(), None)
            .await
            .unwrap()
            .unwrap();
        let body: BytesMut = stream.try_collect().await.unwrap();
        assert_eq!(body.as_ref(), b"hello world");
    }

    #[tokio::test]
    async fn basic_spawn_metadata_and_delete() {
        let service = make_service();

        let id = service
            .insert_object(
                make_context(),
                Some("meta-key".into()),
                Metadata::default(),
                stream::single("data"),
                Timestamp::now(),
            )
            .await
            .unwrap();

        let metadata = service
            .get_metadata(id.clone(), Timestamp::now())
            .await
            .unwrap();
        assert!(metadata.is_some());

        service
            .delete_object(id.clone(), Timestamp::now())
            .await
            .unwrap();

        let after = service
            .get_object(id, Timestamp::now(), None)
            .await
            .unwrap();
        assert!(after.is_none());
    }

    #[tokio::test]
    async fn set_expiry() {
        let service = make_service();
        let old_expiry = Timestamp::now() + Duration::from_hours(1);
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
            time_expires: Some(old_expiry),
            ..Default::default()
        };
        let id = service
            .insert_object(
                make_context(),
                Some("explicit-expiry".into()),
                metadata,
                stream::single("payload"),
                Timestamp::now(),
            )
            .await
            .unwrap();
        let requested = old_expiry + Duration::from_hours(1);

        assert!(
            service
                .set_expiry(id.clone(), requested, Timestamp::now())
                .await
                .unwrap()
        );
        assert_eq!(
            service
                .get_metadata(id, Timestamp::now())
                .await
                .unwrap()
                .unwrap()
                .time_expires,
            Some(requested)
        );
    }

    #[derive(Debug)]
    struct PanicOnGet;

    #[async_trait::async_trait]
    impl Hooks for PanicOnGet {
        async fn get_object(
            &self,
            _inner: &InMemoryBackend,
            _id: &ObjectId,
            _access_time: Timestamp,
            _range: Option<ByteRange>,
        ) -> Result<GetResponse> {
            panic!("intentional panic in get_object");
        }
    }

    #[tokio::test]
    async fn panic_in_backend_returns_task_failed() {
        let service = StorageService::new(
            Box::new(TestBackend::new(PanicOnGet)),
            Cipher::ephemeral().unwrap(),
        );

        let id = ObjectId::new(make_context(), "panic-test".into());
        let result = service.get_object(id, Timestamp::now(), None).await;

        let Err(error) = result else {
            panic!("expected Panic error");
        };
        assert_eq!(error.kind(), ErrorKind::Panic);
        assert_eq!(error.to_string(), "service task panicked");
        assert_eq!(
            std::error::Error::source(&error).unwrap().to_string(),
            "intentional panic in get_object"
        );
    }

    #[derive(Clone, Debug, Default)]
    struct GateOnExpiry {
        calls: Arc<AtomicUsize>,
        access_time: Arc<Mutex<Option<Timestamp>>>,
        started: Arc<tokio::sync::Notify>,
        resume: Arc<tokio::sync::Notify>,
    }

    #[async_trait::async_trait]
    impl Hooks for GateOnExpiry {
        async fn set_expiry(
            &self,
            inner: &InMemoryBackend,
            id: &ObjectId,
            expire_at: Timestamp,
            access_time: Timestamp,
        ) -> Result<bool> {
            *self.access_time.lock().unwrap() = Some(access_time);
            self.calls.fetch_add(1, Ordering::SeqCst);
            self.started.notify_one();
            self.resume.notified().await;
            inner.set_expiry(id, expire_at, access_time).await
        }
    }

    fn stale_tti_metadata() -> Metadata {
        Metadata {
            expiration_policy: ExpirationPolicy::TimeToIdle(Duration::from_hours(1)),
            time_expires: Some(Timestamp::now() + Duration::from_mins(1)),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn background_renewal() {
        let now = Timestamp::now();
        let access_time = now - Duration::from_secs(10);
        let backend = TestBackend::new(GateOnExpiry::default());
        let id = ObjectId::new(make_context(), "background-renewal".into());
        let metadata = stale_tti_metadata();
        backend
            .inner
            .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
            .await
            .unwrap();
        let mut service =
            StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
        service.start();

        let response = tokio::time::timeout(
            Duration::from_secs(1),
            service.get_object(id.clone(), access_time, Some(ByteRange::Bounded(0, 2))),
        )
        .await
        .expect("GET waited for its background renewal")
        .unwrap()
        .unwrap();
        assert_eq!(response.0.time_expires, metadata.time_expires);
        backend.hooks.started.notified().await;
        assert!(backend.hooks.access_time.lock().unwrap().unwrap() >= now);

        let join = tokio::spawn({
            let service = service.clone();
            async move { service.join().await }
        });
        tokio::pin!(join);
        assert!(
            tokio::time::timeout(Duration::from_millis(25), &mut join)
                .await
                .is_err()
        );

        backend.hooks.resume.notify_waiters();
        tokio::time::timeout(Duration::from_secs(1), &mut join)
            .await
            .expect("shutdown did not drain renewal")
            .unwrap();
        assert_eq!(
            backend.inner.get(&id).expect_object().0.time_expires,
            Some(access_time + Duration::from_hours(1))
        );
    }

    #[tokio::test]
    async fn renewal_deduplication() {
        let backend = TestBackend::new(GateOnExpiry::default());
        let id = ObjectId::new(make_context(), "deduplicated-renewal".into());
        backend
            .inner
            .put_object(
                &id,
                &stale_tti_metadata(),
                stream::single("payload"),
                Timestamp::now(),
            )
            .await
            .unwrap();
        let mut service =
            StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
        service.start();

        service
            .get_metadata(id.clone(), Timestamp::now())
            .await
            .unwrap();
        backend.hooks.started.notified().await;
        service.get_metadata(id, Timestamp::now()).await.unwrap();
        tokio::task::yield_now().await;
        assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 1);

        backend.hooks.resume.notify_waiters();
        service.join().await;
    }

    #[tokio::test]
    async fn ttl_read() {
        let backend = TestBackend::new(GateOnExpiry::default());
        let id = ObjectId::new(make_context(), "ttl-no-renewal".into());
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
            time_expires: Some(Timestamp::now() + Duration::from_mins(1)),
            ..Default::default()
        };
        backend
            .inner
            .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
            .await
            .unwrap();
        let mut service =
            StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap());
        service.start();

        service.get_metadata(id, Timestamp::now()).await.unwrap();
        tokio::task::yield_now().await;
        assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 0);
        service.join().await;
    }

    #[tokio::test]
    async fn renewal_queueing() {
        let backend = TestBackend::new(GateOnExpiry::default());
        let first = ObjectId::new(make_context(), "first-renewal".into());
        let second = ObjectId::new(make_context(), "queued-renewal".into());
        let metadata = stale_tti_metadata();
        for id in [&first, &second] {
            backend
                .inner
                .put_object(id, &metadata, stream::single("payload"), Timestamp::now())
                .await
                .unwrap();
        }

        let concurrency = ConcurrencyLimiter::new(1);
        let mut scheduler = RenewalScheduler::new(Arc::new(backend.clone()), concurrency, 1);
        let expire_at = Timestamp::now() + Duration::from_hours(1);
        scheduler.schedule(first, expire_at);
        assert_eq!(scheduler.queued(), 1);
        scheduler.start();
        backend.hooks.started.notified().await;
        assert_eq!(scheduler.queued(), 0);

        scheduler.schedule(second.clone(), expire_at);
        tokio::task::yield_now().await;
        assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 1);
        assert_eq!(
            backend.inner.get(&second).expect_object().0.time_expires,
            metadata.time_expires
        );

        backend.hooks.resume.notify_waiters();
        tokio::time::timeout(Duration::from_secs(1), async {
            while backend.hooks.calls.load(Ordering::SeqCst) < 2 {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("queued renewal did not start");
        backend.hooks.resume.notify_waiters();
        scheduler.join().await;
        assert_eq!(scheduler.queued(), 0);
        assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 2);
        assert!(backend.inner.get(&second).expect_object().0.time_expires > metadata.time_expires);
    }

    #[derive(Clone, Debug, Default)]
    struct FailFirstExpiry {
        calls: Arc<AtomicUsize>,
        panic: bool,
    }

    #[async_trait::async_trait]
    impl Hooks for FailFirstExpiry {
        async fn set_expiry(
            &self,
            inner: &InMemoryBackend,
            id: &ObjectId,
            expire_at: Timestamp,
            access_time: Timestamp,
        ) -> Result<bool> {
            if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
                assert!(!self.panic, "intentional renewal panic");
                return Err(ErrorKind::BackendFailure.into());
            }
            inner.set_expiry(id, expire_at, access_time).await
        }
    }

    #[tokio::test]
    async fn renewal_failure_cleanup() {
        for panic in [false, true] {
            let backend = TestBackend::new(FailFirstExpiry {
                panic,
                ..Default::default()
            });
            let id = ObjectId::new(make_context(), "failed-renewal".into());
            let metadata = stale_tti_metadata();
            backend
                .inner
                .put_object(&id, &metadata, stream::single("payload"), Timestamp::now())
                .await
                .unwrap();
            let concurrency = ConcurrencyLimiter::new(1);
            let mut scheduler = RenewalScheduler::new(Arc::new(backend.clone()), concurrency, 1);
            scheduler.start();
            let expire_at = metadata.check_tti_bump(Timestamp::now()).unwrap();

            scheduler.schedule(id.clone(), expire_at);
            tokio::time::timeout(Duration::from_secs(1), async {
                while scheduler.pending() != 0 {
                    tokio::task::yield_now().await;
                }
            })
            .await
            .expect("failure did not release renewal guards");

            scheduler.schedule(id, expire_at);
            scheduler.join().await;
            assert_eq!(backend.hooks.calls.load(Ordering::SeqCst), 2);
        }
    }

    /// In-memory backend with optional synchronization for `put_object`.
    ///
    /// When `pause` is enabled, each `put_object` call notifies `paused` and
    #[derive(Clone, Debug, Default)]
    struct GateOnPut {
        pause: bool,
        paused: Arc<tokio::sync::Notify>,
        resume: Arc<tokio::sync::Notify>,
        on_put: Arc<tokio::sync::Notify>,
    }

    impl GateOnPut {
        fn with_pause() -> Self {
            Self {
                pause: true,
                ..Default::default()
            }
        }
    }

    #[async_trait::async_trait]
    impl Hooks for GateOnPut {
        async fn put_object(
            &self,
            inner: &InMemoryBackend,
            id: &ObjectId,
            metadata: &Metadata,
            stream: ClientStream,
            access_time: Timestamp,
        ) -> Result<PutResponse> {
            if self.pause {
                self.paused.notify_one();
                self.resume.notified().await;
            }
            inner.put_object(id, metadata, stream, access_time).await?;
            self.on_put.notify_one();
            Ok(())
        }

        async fn compare_and_write(
            &self,
            inner: &InMemoryBackend,
            id: &ObjectId,
            current: Option<&ObjectId>,
            write: TieredWrite,
            access_time: Timestamp,
        ) -> Result<bool> {
            let notify = matches!(write, TieredWrite::Tombstone(_) | TieredWrite::Object(_, _));
            let result = inner
                .compare_and_write(id, current, write, access_time)
                .await?;
            if notify {
                self.on_put.notify_one();
            }
            Ok(result)
        }
    }

    #[tokio::test]
    async fn receiver_drop_does_not_prevent_completion() {
        let hv = Box::new(TestBackend::new(GateOnPut::default()));
        let lt = Box::new(TestBackend::new(GateOnPut::with_pause()));
        let backend = TieredStorage::new(hv.clone(), lt.clone(), Box::new(NoopChangeLog));
        let service = StorageService::new(Box::new(backend), Cipher::ephemeral().unwrap());

        let payload = vec![0xABu8; 2 * 1024 * 1024]; // 2 MiB → long-term path
        let request = service.insert_object(
            make_context(),
            Some("completion-test".into()),
            Metadata::default(),
            stream::single(payload),
            Timestamp::now(),
        );

        // Start insert through the public API. select! drops the future once the
        // backend signals it has paused, simulating a client disconnect mid-write.
        let paused = Arc::clone(&lt.hooks.paused);
        tokio::select! {
            _ = request => panic!("insert should not complete while backend is paused"),
            _ = paused.notified() => {}
        }

        // The spawned task is now blocked inside put_object, and the caller
        // request (including the oneshot receiver) has been dropped. Unpause so
        // the task can finish writing.
        lt.hooks.resume.notify_one();

        // Wait for the tombstone write to the high-volume backend, which is the
        // last step of the long-term insert path.
        let on_put = Arc::clone(&hv.hooks.on_put);
        tokio::time::timeout(Duration::from_secs(5), on_put.notified())
            .await
            .expect("timed out waiting for tombstone write");

        // Verify the object was fully written despite the caller being dropped.
        // The tombstone in HV points to the revision key in LT.
        let id = ObjectId::new(make_context(), "completion-test".into());
        let tombstone = hv.inner.get(&id).expect_tombstone();
        let lt_id = tombstone.target;
        assert!(lt.inner.contains(&lt_id), "long-term object missing");
    }

    // --- Concurrency limit tests ---

    fn make_limited_service(limit: u32) -> (StorageService, TestBackend<GateOnPut>) {
        let backend = TestBackend::new(GateOnPut::with_pause());
        let service = StorageService::new(Box::new(backend.clone()), Cipher::ephemeral().unwrap())
            .with_concurrency(ConcurrencyLimiter::new(limit));
        (service, backend)
    }

    #[tokio::test]
    async fn at_capacity_rejects() {
        let (service, hv) = make_limited_service(1);

        // First insert blocks on the gated backend, holding the single permit.
        let svc = service.clone();
        let first = tokio::spawn(async move {
            svc.insert_object(
                make_context(),
                Some("first".into()),
                Metadata::default(),
                stream::single("data"),
                Timestamp::now(),
            )
            .await
        });

        // Wait for the backend to signal it has paused (permit is held).
        hv.hooks.paused.notified().await;

        // Second insert should be rejected immediately.
        let result = service
            .insert_object(
                make_context(),
                Some("second".into()),
                Metadata::default(),
                stream::single("data"),
                Timestamp::now(),
            )
            .await;

        assert!(
            result
                .as_ref()
                .is_err_and(|error| error.kind() == ErrorKind::AtCapacity),
            "expected AtCapacity, got {result:?}"
        );

        // Unblock the first operation.
        hv.hooks.resume.notify_one();
        first.await.unwrap().unwrap();

        // Now that the permit is released, a new operation should succeed.
        service
            .get_metadata(
                ObjectId::new(make_context(), "first".into()),
                Timestamp::now(),
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn tasks_limit_returns_configured_limit() {
        let backend = Box::new(InMemoryBackend::new("cap"));
        let service = StorageService::new(backend, Cipher::ephemeral().unwrap())
            .with_concurrency(ConcurrencyLimiter::new(7));
        assert_eq!(service.tasks_limit(), 7);
    }

    #[tokio::test]
    async fn tasks_running_tracks_in_flight() {
        let (service, hv) = make_limited_service(5);

        assert_eq!(service.tasks_running(), 0);

        // Kick off a request that blocks in the backend, holding a permit.
        let svc = service.clone();
        let _blocked = tokio::spawn(async move {
            svc.insert_object(
                make_context(),
                Some("in-use-test".into()),
                Metadata::default(),
                stream::single("data"),
                Timestamp::now(),
            )
            .await
        });

        hv.hooks.paused.notified().await;
        assert_eq!(service.tasks_running(), 1);

        hv.hooks.resume.notify_one();
    }

    #[tokio::test]
    async fn permits_released_after_panic() {
        let service = StorageService::new(
            Box::new(TestBackend::new(PanicOnGet)),
            Cipher::ephemeral().unwrap(),
        )
        .with_concurrency(ConcurrencyLimiter::new(1));

        // First operation panics — the permit must still be released.
        let id = ObjectId::new(make_context(), "panic-permit".into());
        let result = service.get_object(id.clone(), Timestamp::now(), None).await;
        assert!(result.is_err_and(|error| error.kind() == ErrorKind::Panic));

        // Second operation should succeed in acquiring the permit (not AtCapacity).
        let result = service.get_object(id, Timestamp::now(), None).await;
        assert!(
            !result.is_err_and(|error| error.kind() == ErrorKind::AtCapacity),
            "permit was not released after panic"
        );
    }

    // --- Resumable uploads ---

    #[tokio::test]
    async fn resumable_create_preserves_backend_refusal_as_none() {
        let service = make_service();
        let id = ObjectId::new(make_context(), "resumable".into());

        let result = service
            .create_upload_session(id, Metadata::default(), 1024)
            .await;

        assert!(matches!(result, Ok(None)), "{result:?}");
    }

    #[tokio::test]
    async fn resumable_create_declines_zero_length() {
        let service = StorageService::new(
            Box::new(TestBackend::new(ResumableTokenHooks::default())),
            Cipher::ephemeral().unwrap(),
        );
        let id = ObjectId::new(make_context(), "resumable".into());

        let result = service
            .create_upload_session(id, Metadata::default(), 0)
            .await;

        assert!(matches!(result, Ok(None)), "{result:?}");
    }

    #[tokio::test]
    async fn resumable_create_validates_metadata() {
        let service = make_service();
        let id = ObjectId::new(make_context(), "resumable".into());

        // A timeout policy with no resolved `time_expires` is rejected before the backend
        // is consulted, exactly as it is for a regular insert.
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_secs(60)),
            ..Default::default()
        };

        let result = service.create_upload_session(id, metadata, 1024).await;
        assert!(result.is_err_and(|error| error.kind() == ErrorKind::InvalidMetadata));
    }

    #[tokio::test]
    async fn resumable_tokens_are_encrypted_by_default() -> Result<()> {
        let hooks = ResumableTokenHooks::default();
        let service = StorageService::new(
            Box::new(TestBackend::new(hooks.clone())),
            Cipher::ephemeral().unwrap(),
        );
        let id = ObjectId::new(make_context(), "resumable".into());

        let token = service
            .create_upload_session(id.clone(), Metadata::default(), 4)
            .await?
            .expect("test backend supports resumable uploads");
        assert_ne!(token.as_bytes(), b"backend token");
        assert!(matches!(
            service
                .upload_offset(id.clone(), EncryptedSessionToken::new(b"backend token"))
                .await,
            Err(error) if error.kind() == ErrorKind::UnknownUploadSession
        ));
        let other_id = ObjectId::new(make_context(), "other".into());
        assert!(matches!(
            service.upload_offset(other_id, token.clone()).await,
            Err(error) if error.kind() == ErrorKind::UnknownUploadSession
        ));
        assert!(hooks.seen_tokens.lock().unwrap().is_empty());
        service.upload_offset(id, token).await?;
        assert_eq!(
            hooks.seen_tokens.lock().unwrap().as_slice(),
            &["backend token"]
        );
        Ok(())
    }

    #[tokio::test]
    async fn configured_encryption_only_crosses_the_service_boundary() -> Result<()> {
        let hooks = ResumableTokenHooks::default();
        let encryption = Cipher::new(
            "v1",
            std::collections::BTreeMap::from([("v1".into(), vec![7; 32])]),
        )
        .unwrap();
        let service = StorageService::new(Box::new(TestBackend::new(hooks.clone())), encryption);
        let id = ObjectId::new(make_context(), "resumable".into());

        let encrypted = service
            .create_upload_session(id.clone(), Metadata::default(), 4)
            .await?
            .expect("test backend supports resumable uploads");
        assert_ne!(encrypted.as_bytes(), b"backend token");
        service.upload_offset(id, encrypted).await?;
        assert_eq!(
            hooks.seen_tokens.lock().unwrap().as_slice(),
            &["backend token"]
        );
        Ok(())
    }

    #[tokio::test]
    async fn configured_encryption_rejects_plaintext_tokens() {
        let hooks = ResumableTokenHooks::default();
        let encryption = Cipher::new(
            "v1",
            std::collections::BTreeMap::from([("v1".into(), vec![7; 32])]),
        )
        .unwrap();
        let service = StorageService::new(Box::new(TestBackend::new(hooks.clone())), encryption);
        let id = ObjectId::new(make_context(), "resumable".into());

        let result = service
            .upload_offset(id, EncryptedSessionToken::new(b"backend token"))
            .await;
        let error = result.unwrap_err();
        assert_eq!(error.kind(), ErrorKind::UnknownUploadSession);
        assert_eq!(error.to_string(), "unknown upload session");
        assert!(error.source().is_none());
        assert!(hooks.seen_tokens.lock().unwrap().is_empty());
    }
}