Skip to main content

slatedb_mirror/
mirror.rs

1//! A policy-driven local mirror of whole objects from an [`ObjectStore`].
2//!
3//! [`ObjectStoreMirror`] implements the object-store API and keeps local copies
4//! with the same bytes as their remote objects. A [`MirrorPolicy`] selects the
5//! read and write routes. The mirror handles local files, downloads, and cleanup.
6//! The policy decides which objects to keep and when callers can read them locally.
7//!
8//! Use [`ObjectStoreMirror::builder`] to choose a local directory, a remote store,
9//! and a policy. [`ObjectStoreMirrorBuilder::build`] recovers local copies and
10//! starts the background tasks. Policies use a [`MirrorHandle`] to fetch,
11//! prefetch, and evict objects.
12//!
13//! Mirrored paths must never be overwritten. Local copies are disposable and
14//! require callers to detect corruption. See [`ObjectStoreMirror`] for routing,
15//! error handling, lifecycle details, and a complete example.
16
17use std::fmt;
18use std::path::PathBuf;
19use std::sync::Arc;
20use std::time::Duration;
21
22use async_trait::async_trait;
23use bytes::Bytes;
24use futures::future;
25use futures::stream::{self, BoxStream, StreamExt};
26use object_store::path::Path;
27use object_store::{
28    CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
29    ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions,
30};
31use slatedb_common::clock::{DefaultSystemClock, SystemClock};
32
33use crate::error::MIRROR_STORE_NAME;
34use crate::inner::{local_error, Inner};
35use crate::layout::{LocalName, LocalObject, LOCK_FILE};
36use crate::multipart::{self, committed};
37use crate::ordering::OpKind;
38use crate::retry::Retry;
39use crate::startup::recover;
40use crate::{scan, MirrorError, MirrorHandle, MirrorPolicy, MultipartRoute, ReadRoute, WriteRoute};
41use slatedb_common::vfs::{StdVfs, Vfs};
42
43/// The default for [`ObjectStoreMirrorBuilder::with_download_concurrency`].
44pub const DEFAULT_DOWNLOAD_CONCURRENCY: usize = 8;
45
46/// The default for [`ObjectStoreMirrorBuilder::with_remote_scan_interval`].
47pub const DEFAULT_REMOTE_SCAN_INTERVAL: Duration = Duration::from_secs(600);
48
49/// Wraps a route error from the policy in `MirrorError::Policy`, unless it's
50/// already a mirror error.
51fn policy_error(err: object_store::Error) -> object_store::Error {
52    if MirrorError::find(&err).is_some() {
53        err
54    } else {
55        MirrorError::Policy {
56            source: Box::new(err),
57        }
58        .into()
59    }
60}
61
62fn range_error(source: impl std::error::Error + Send + Sync + 'static) -> object_store::Error {
63    object_store::Error::Generic {
64        store: MIRROR_STORE_NAME,
65        source: Box::new(source),
66    }
67}
68
69fn bytes_payload(bytes: Bytes) -> GetResultPayload {
70    GetResultPayload::Stream(stream::once(async move { Ok(bytes) }).boxed())
71}
72
73/// An [`ObjectStore`] wrapper that keeps whole objects in a local directory.
74///
75/// The remote store holds the authoritative objects. Local copies contain the
76/// same bytes, without splitting objects into blocks or changing their format.
77/// A [`MirrorPolicy`] chooses which objects to mirror. The mirror does not impose
78/// a size limit or automatically evict objects to reclaim disk space.
79///
80/// # Read and write routes
81///
82/// A [`MirrorPolicy`] selects how each request uses the remote store and
83/// local copies:
84///
85/// - GET and HEAD use [`ReadRoute`].
86/// - PUT uses [`WriteRoute`].
87/// - Multipart uploads use [`MultipartRoute`].
88///
89/// The `Observe` routes pass the whole object to [`MirrorPolicy::observe`].
90/// Multipart uploads do not support observation.
91///
92/// Policies manage local copies through [`MirrorHandle`]:
93///
94/// - [`fetch`](MirrorHandle::fetch) waits for a local copy.
95/// - [`prefetch`](MirrorHandle::prefetch) schedules a background download.
96/// - [`evict`](MirrorHandle::evict) removes a local copy.
97///
98/// Fetching an existing local copy does not refresh it. Use
99/// [`ReadRoute::Refetch`] to replace it.
100///
101/// Downloads, mirrored writes, deletes, and evictions coordinate per path.
102/// A later download or mirrored write can cancel an eviction that has not started.
103///
104/// Listings always use the remote store. Deletes remove the remote object and
105/// its local copy. Copy and rename operations return [`MirrorError::Unsupported`].
106///
107/// # Object lifetime and local integrity
108///
109/// Policies must only mirror paths that are never overwritten. The mirror does
110/// not track remote updates to existing objects. Route mutable objects through
111/// `Remote` or `Observe`.
112///
113/// Local copies are disposable. The default [`StdVfs`] flushes writes without
114/// syncing them to disk. Local reads can return corrupted bytes without an
115/// error, including after a machine crash. Startup checks metadata, filenames,
116/// and file sizes, but does not check object contents.
117///
118/// Callers must detect corruption and signal their policy to route a retry
119/// through [`ReadRoute::Refetch`]. Repeating a [`ReadRoute::Local`] read does
120/// not repair the local copy. Callers must also check the refetched data for
121/// corruption and limit the number of retries.
122///
123/// # Startup and background tasks
124///
125/// Call [`ObjectStoreMirrorBuilder::build`] inside a Tokio runtime. It creates
126/// the local directory, acquires its exclusive lock, and recovers complete local
127/// copies. Use a dedicated directory for each mirror and remote store.
128///
129/// The builder starts [`MirrorPolicy::run`] and a periodic remote scan. The scan
130/// removes local copies of objects deleted remotely. It does not download new
131/// objects or refresh existing copies. The scan requires strongly consistent
132/// remote reads, writes, deletes, and listings. If the remote store lacks these
133/// guarantees, disable the scan with
134/// [`with_remote_scan_interval(None)`](ObjectStoreMirrorBuilder::with_remote_scan_interval).
135///
136/// The defaults allow [eight concurrent downloads](DEFAULT_DOWNLOAD_CONCURRENCY)
137/// and [a remote scan every ten minutes](DEFAULT_REMOTE_SCAN_INTERVAL). Downloads
138/// and remote scan listings retry transient errors without a limit by default.
139/// Use [`ObjectStoreMirrorBuilder::with_max_retries`] to bound these retries.
140///
141/// Dropping the last `Arc<ObjectStoreMirror>` cancels the background tasks.
142/// Cloned [`MirrorHandle`] values keep the local state and directory lock alive
143/// until those handles are also dropped.
144///
145/// # Errors
146///
147/// Mirror errors travel inside [`object_store::Error`]. Use [`MirrorError::find`]
148/// to recover the specific error from its source chain.
149/// [`MirrorError::WriteCommitted`] means that the remote write succeeded, but
150/// local installation or the observation hook failed. Account for that remote
151/// write before deciding whether to retry it.
152///
153/// # Example
154///
155/// This policy mirrors every write and serves every read locally. It assumes
156/// that callers write each path once. Existing remote objects need an explicit
157/// [`MirrorHandle::fetch`] before a local read.
158///
159/// ```
160/// use std::sync::Arc;
161///
162/// use async_trait::async_trait;
163/// use object_store::memory::InMemory;
164/// use object_store::path::Path;
165/// use object_store::{
166///     GetOptions, ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions,
167/// };
168/// use slatedb_mirror::{MirrorPolicy, MultipartRoute, ObjectStoreMirror, ReadRoute, WriteRoute};
169///
170/// struct ImmutableObjects;
171///
172/// #[async_trait]
173/// impl MirrorPolicy for ImmutableObjects {
174///     fn read_route(&self, _: &Path, _: &GetOptions) -> object_store::Result<ReadRoute> {
175///         Ok(ReadRoute::Local)
176///     }
177///
178///     fn put_route(&self, _: &Path, _: &PutOptions) -> object_store::Result<WriteRoute> {
179///         Ok(WriteRoute::Mirror)
180///     }
181///
182///     fn put_multipart_route(
183///         &self,
184///         _: &Path,
185///         _: &PutMultipartOptions,
186///     ) -> object_store::Result<MultipartRoute> {
187///         Ok(MultipartRoute::Mirror)
188///     }
189/// }
190///
191/// # #[tokio::main(flavor = "current_thread")]
192/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
193/// let local_dir = tempfile::tempdir()?;
194/// let remote: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
195/// let mirror = ObjectStoreMirror::builder(
196///     local_dir.path(),
197///     remote,
198///     Arc::new(ImmutableObjects),
199/// )
200/// .build()
201/// .await?;
202///
203/// let path = Path::from("objects/first.bin");
204/// mirror.put(&path, "hello".into()).await?;
205/// assert!(mirror.handle().contains(&path));
206/// let bytes = mirror.get(&path).await?.bytes().await?;
207/// assert_eq!(bytes.as_ref(), b"hello");
208/// # Ok(())
209/// # }
210/// ```
211pub struct ObjectStoreMirror {
212    pub(crate) inner: Arc<Inner>,
213}
214
215impl ObjectStoreMirror {
216    pub fn builder(
217        local_dir: impl Into<PathBuf>,
218        object_store: Arc<dyn ObjectStore>,
219        policy: Arc<dyn MirrorPolicy>,
220    ) -> ObjectStoreMirrorBuilder {
221        ObjectStoreMirrorBuilder {
222            local_dir: local_dir.into(),
223            remote: object_store,
224            policy,
225            vfs: Arc::new(StdVfs::new()),
226            clock: Arc::new(DefaultSystemClock::new()),
227            download_concurrency: DEFAULT_DOWNLOAD_CONCURRENCY,
228            remote_scan_interval: Some(DEFAULT_REMOTE_SCAN_INTERVAL),
229            max_retries: None,
230        }
231    }
232
233    /// A handle to this mirror, like the one the policy gets.
234    pub fn handle(&self) -> MirrorHandle {
235        MirrorHandle {
236            inner: Arc::clone(&self.inner),
237        }
238    }
239
240    async fn get_observed(
241        &self,
242        location: &Path,
243        mut options: GetOptions,
244    ) -> object_store::Result<GetResult> {
245        let inner = &self.inner;
246        if options.head {
247            return inner.remote.get_opts(location, options).await;
248        }
249        // `observe` needs the whole object. A `NotModified` or other error
250        // returns here, before `observe`.
251        let range = options.range.take();
252        let result = inner.remote.get_opts(location, options).await?;
253        let meta = result.meta.clone();
254        let attributes = result.attributes.clone();
255        let extensions = result.extensions.clone();
256        let bytes = result.bytes().await?;
257        let range = match range {
258            Some(range) => range.as_range(bytes.len() as u64).map_err(range_error)?,
259            None => 0..bytes.len() as u64,
260        };
261        inner
262            .policy
263            .observe(location, &bytes, &self.handle())
264            .await?;
265        Ok(GetResult {
266            payload: bytes_payload(bytes.slice(range.start as usize..range.end as usize)),
267            meta,
268            range,
269            attributes,
270            extensions,
271        })
272    }
273
274    async fn get_local(
275        &self,
276        location: &Path,
277        options: GetOptions,
278    ) -> object_store::Result<GetResult> {
279        let inner = &self.inner;
280        let Some(object) = inner.local(location) else {
281            return Err(MirrorError::NotLocal {
282                path: location.clone(),
283            }
284            .into());
285        };
286        inner.read_local(location, &object, &options).await
287    }
288
289    async fn get_refetched(
290        &self,
291        location: &Path,
292        options: GetOptions,
293    ) -> object_store::Result<GetResult> {
294        let inner = &self.inner;
295        let ticket = inner
296            .ordering
297            .register(location, OpKind::Download)
298            .expect("downloads are always queued");
299        ticket.ready().await;
300        let object = inner.download(location).await?;
301        // Keep the ticket until the read is done, so an eviction can't remove
302        // the file underneath it.
303        let result = inner.read_local(location, &object, &options).await;
304        drop(ticket);
305        result
306    }
307
308    async fn put_mirrored(
309        &self,
310        location: &Path,
311        payload: PutPayload,
312        opts: PutOptions,
313    ) -> object_store::Result<PutResult> {
314        let inner = &self.inner;
315        let name = LocalName::new(location);
316        inner.prefixes.lock().check_or_insert(&name)?;
317        let ticket = inner
318            .ordering
319            .register(location, OpKind::Write)
320            .expect("writes are always queued");
321        ticket.ready().await;
322
323        let size = payload.content_length() as u64;
324        let attributes = opts.attributes.clone();
325        let temp = inner.temp_file(&name);
326        let (local, remote) = future::join(
327            inner.write_payload(&temp.path, &payload),
328            inner.remote.put_opts(location, payload.clone(), opts),
329        )
330        .await;
331        let put = match remote {
332            Ok(put) => put,
333            Err(err) => {
334                temp.remove().await;
335                return Err(err);
336            }
337        };
338        if let Err(err) = local {
339            temp.remove().await;
340            return Err(committed(location, local_error(err)));
341        }
342
343        let object = LocalObject {
344            meta: ObjectMeta {
345                location: location.clone(),
346                last_modified: inner.clock.now(),
347                size,
348                e_tag: put.e_tag.clone(),
349                version: put.version.clone(),
350            },
351            attributes,
352        };
353        inner
354            .install(location, &name, temp, object)
355            .await
356            .map_err(|err| committed(location, err))?;
357        drop(ticket);
358        Ok(put)
359    }
360}
361
362impl Drop for ObjectStoreMirror {
363    fn drop(&mut self) {
364        self.inner.shutdown.cancel();
365    }
366}
367
368impl fmt::Debug for ObjectStoreMirror {
369    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
370        f.debug_struct("ObjectStoreMirror")
371            .field("root", &self.inner.root)
372            .field("remote", &self.inner.remote)
373            .finish_non_exhaustive()
374    }
375}
376
377impl fmt::Display for ObjectStoreMirror {
378    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
379        write!(f, "ObjectStoreMirror({})", self.inner.remote)
380    }
381}
382
383#[async_trait]
384impl ObjectStore for ObjectStoreMirror {
385    async fn put_opts(
386        &self,
387        location: &Path,
388        payload: PutPayload,
389        opts: PutOptions,
390    ) -> object_store::Result<PutResult> {
391        let inner = &self.inner;
392        let route = inner
393            .policy
394            .put_route(location, &opts)
395            .map_err(policy_error)?;
396        match route {
397            WriteRoute::Remote => inner.remote.put_opts(location, payload, opts).await,
398            WriteRoute::Observe => {
399                let bytes = Bytes::from(payload.clone());
400                let put = inner.remote.put_opts(location, payload, opts).await?;
401                inner
402                    .policy
403                    .observe(location, &bytes, &self.handle())
404                    .await
405                    .map_err(|err| committed(location, err))?;
406                Ok(put)
407            }
408            WriteRoute::Mirror => self.put_mirrored(location, payload, opts).await,
409        }
410    }
411
412    async fn put_multipart_opts(
413        &self,
414        location: &Path,
415        opts: PutMultipartOptions,
416    ) -> object_store::Result<Box<dyn MultipartUpload>> {
417        let inner = &self.inner;
418        let route = inner
419            .policy
420            .put_multipart_route(location, &opts)
421            .map_err(policy_error)?;
422        match route {
423            MultipartRoute::Remote => inner.remote.put_multipart_opts(location, opts).await,
424            MultipartRoute::Mirror => multipart::start(Arc::clone(inner), location, opts).await,
425        }
426    }
427
428    async fn get_opts(
429        &self,
430        location: &Path,
431        options: GetOptions,
432    ) -> object_store::Result<GetResult> {
433        let inner = &self.inner;
434        let route = inner
435            .policy
436            .read_route(location, &options)
437            .map_err(policy_error)?;
438        match route {
439            ReadRoute::Remote => inner.remote.get_opts(location, options).await,
440            ReadRoute::Observe => self.get_observed(location, options).await,
441            ReadRoute::Local => self.get_local(location, options).await,
442            ReadRoute::Refetch => self.get_refetched(location, options).await,
443        }
444    }
445
446    fn delete_stream(
447        &self,
448        locations: BoxStream<'static, object_store::Result<Path>>,
449    ) -> BoxStream<'static, object_store::Result<Path>> {
450        let inner = Arc::clone(&self.inner);
451        locations
452            .then(move |location| {
453                let inner = Arc::clone(&inner);
454                async move {
455                    let location = location?;
456                    inner.delete(&location).await?;
457                    Ok(location)
458                }
459            })
460            .boxed()
461    }
462
463    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
464        self.inner.remote.list(prefix)
465    }
466
467    fn list_with_offset(
468        &self,
469        prefix: Option<&Path>,
470        offset: &Path,
471    ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
472        self.inner.remote.list_with_offset(prefix, offset)
473    }
474
475    async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result<ListResult> {
476        self.inner.remote.list_with_delimiter(prefix).await
477    }
478
479    async fn copy_opts(
480        &self,
481        _from: &Path,
482        _to: &Path,
483        _options: CopyOptions,
484    ) -> object_store::Result<()> {
485        Err(MirrorError::Unsupported { operation: "copy" }.into())
486    }
487
488    async fn rename_opts(
489        &self,
490        _from: &Path,
491        _to: &Path,
492        _options: RenameOptions,
493    ) -> object_store::Result<()> {
494        Err(MirrorError::Unsupported {
495            operation: "rename",
496        }
497        .into())
498    }
499}
500
501/// Configures and builds an [`ObjectStoreMirror`].
502pub struct ObjectStoreMirrorBuilder {
503    local_dir: PathBuf,
504    remote: Arc<dyn ObjectStore>,
505    policy: Arc<dyn MirrorPolicy>,
506    vfs: Arc<dyn Vfs>,
507    clock: Arc<dyn SystemClock>,
508    download_concurrency: usize,
509    remote_scan_interval: Option<Duration>,
510    max_retries: Option<u32>,
511}
512
513impl fmt::Debug for ObjectStoreMirrorBuilder {
514    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
515        f.debug_struct("ObjectStoreMirrorBuilder")
516            .field("local_dir", &self.local_dir)
517            .field("remote", &self.remote)
518            .field("vfs", &self.vfs)
519            .field("download_concurrency", &self.download_concurrency)
520            .field("remote_scan_interval", &self.remote_scan_interval)
521            .field("max_retries", &self.max_retries)
522            .finish_non_exhaustive()
523    }
524}
525
526impl ObjectStoreMirrorBuilder {
527    /// Sets the virtual filesystem used for local I/O. The default is
528    /// [`StdVfs`].
529    pub fn with_vfs(mut self, vfs: Arc<dyn Vfs>) -> Self {
530        self.vfs = vfs;
531        self
532    }
533
534    /// Sets the maximum number of concurrent downloads used by `fetch`,
535    /// `prefetch`, and `Refetch`. The default is
536    /// [`DEFAULT_DOWNLOAD_CONCURRENCY`].
537    pub fn with_download_concurrency(mut self, concurrency: usize) -> Self {
538        self.download_concurrency = concurrency;
539        self
540    }
541
542    /// Sets the interval of the remote scan that removes local copies of
543    /// objects deleted remotely. The default is
544    /// [`DEFAULT_REMOTE_SCAN_INTERVAL`]. `None` disables the scan.
545    ///
546    /// The scan assumes strongly consistent reads, writes, deletes, and
547    /// listings. Disable it if the remote store doesn't provide them.
548    pub fn with_remote_scan_interval(mut self, interval: Option<Duration>) -> Self {
549        self.remote_scan_interval = interval;
550        self
551    }
552
553    /// Sets the clock used for retry backoff, the remote scan, and the
554    /// last-modified time of mirrored writes. The default uses Tokio's clock.
555    pub fn with_system_clock(mut self, clock: Arc<dyn SystemClock>) -> Self {
556        self.clock = clock;
557        self
558    }
559
560    /// Sets how many times the mirror retries a transient error from its own
561    /// remote calls (downloads and remote-scan LISTs). The default, `None`,
562    /// retries forever, like SlateDB's `object_store_max_retries`.
563    pub fn with_max_retries(mut self, max_retries: Option<u32>) -> Self {
564        self.max_retries = max_retries;
565        self
566    }
567
568    /// Validates the configuration, acquires the cache-directory lock, cleans
569    /// invalid local entries, and spawns the remote scan and
570    /// `MirrorPolicy::run`.
571    pub async fn build(self) -> object_store::Result<Arc<ObjectStoreMirror>> {
572        let invalid = |message: &str| -> object_store::Error {
573            MirrorError::InvalidConfig {
574                message: message.to_string(),
575            }
576            .into()
577        };
578        if self.download_concurrency == 0 {
579            return Err(invalid("download_concurrency must be greater than zero"));
580        }
581        if self.remote_scan_interval == Some(Duration::ZERO) {
582            return Err(invalid("remote_scan_interval must be greater than zero"));
583        }
584        let runtime = tokio::runtime::Handle::try_current()
585            .map_err(|_| invalid("the mirror must be built inside a Tokio runtime"))?;
586
587        self.vfs
588            .create_dir_all(&self.local_dir)
589            .await
590            .map_err(local_error)?;
591        let lock = self
592            .vfs
593            .lock(&self.local_dir.join(LOCK_FILE))
594            .await
595            .map_err(local_error)?;
596        let recovered = recover(self.vfs.as_ref(), &self.local_dir).await?;
597
598        let inner = Arc::new(Inner::new(
599            self.local_dir,
600            self.remote,
601            self.policy,
602            self.vfs,
603            Arc::clone(&self.clock),
604            Retry::new(self.clock, self.max_retries),
605            runtime,
606            recovered.entries,
607            recovered.prefixes,
608            self.download_concurrency,
609            lock,
610        ));
611        if let Some(interval) = self.remote_scan_interval {
612            inner.spawn(scan::run(Arc::clone(&inner), interval));
613        }
614        let handle = MirrorHandle {
615            inner: Arc::clone(&inner),
616        };
617        inner.spawn(Arc::clone(&inner.policy).run(handle));
618        Ok(Arc::new(ObjectStoreMirror { inner }))
619    }
620}
621
622/// Mirrors every object. For unit tests.
623#[cfg(test)]
624struct MirrorAll;
625
626#[cfg(test)]
627#[async_trait]
628impl MirrorPolicy for MirrorAll {
629    fn read_route(&self, _path: &Path, _options: &GetOptions) -> object_store::Result<ReadRoute> {
630        Ok(ReadRoute::Local)
631    }
632
633    fn put_route(&self, _path: &Path, _options: &PutOptions) -> object_store::Result<WriteRoute> {
634        Ok(WriteRoute::Mirror)
635    }
636
637    fn put_multipart_route(
638        &self,
639        _path: &Path,
640        _options: &PutMultipartOptions,
641    ) -> object_store::Result<MultipartRoute> {
642        Ok(MultipartRoute::Mirror)
643    }
644}
645
646#[cfg(test)]
647impl ObjectStoreMirror {
648    /// Builds a mirror of `remote` in `root` that mirrors every object and has
649    /// no remote scan. For unit tests.
650    pub(crate) async fn for_test(
651        root: &std::path::Path,
652        remote: Arc<dyn ObjectStore>,
653    ) -> Arc<Self> {
654        Self::builder(root, remote, Arc::new(MirrorAll))
655            .with_remote_scan_interval(None)
656            .build()
657            .await
658            .unwrap()
659    }
660}
661
662#[cfg(test)]
663mod tests {
664    //! End-to-end tests of `ObjectStoreMirror` over an in-memory remote store.
665
666    use std::collections::BTreeSet;
667    use std::fmt;
668    use std::path::{Path as StdPath, PathBuf};
669    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
670    use std::sync::Arc;
671    use std::time::Duration;
672
673    use async_trait::async_trait;
674    use bytes::Bytes;
675    use futures::stream::BoxStream;
676    use futures::TryStreamExt;
677    use md5::{Digest, Md5};
678    use object_store::memory::InMemory;
679    use object_store::path::Path;
680    use object_store::{
681        CopyOptions, GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta,
682        ObjectStore, ObjectStoreExt, PutMode, PutMultipartOptions, PutOptions, PutPayload,
683        PutResult,
684    };
685    use parking_lot::Mutex;
686    use tokio::sync::Semaphore;
687
688    use crate::{
689        MirrorError, MirrorHandle, MirrorPolicy, MultipartRoute, ReadRoute, StdVfs, Vfs, VfsEntry,
690        VfsLock, VfsWriter, WriteRoute,
691    };
692
693    use super::ObjectStoreMirror;
694
695    /// An in-memory store that counts GETs and can fail or hold them.
696    #[derive(Debug, Default)]
697    struct TestStore {
698        inner: InMemory,
699        /// Full-object GETs (not HEADs).
700        gets: AtomicUsize,
701        in_flight: AtomicUsize,
702        max_in_flight: AtomicUsize,
703        /// Fail this many upcoming GETs with a transient error.
704        fail_gets: AtomicUsize,
705        /// When set, each GET takes a permit first.
706        gate: Mutex<Option<Arc<Semaphore>>>,
707        /// PUTs started.
708        puts: AtomicUsize,
709        /// When set, each PUT takes a permit first.
710        put_gate: Mutex<Option<Arc<Semaphore>>>,
711    }
712
713    impl fmt::Display for TestStore {
714        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
715            write!(f, "TestStore")
716        }
717    }
718
719    #[async_trait]
720    impl ObjectStore for TestStore {
721        async fn put_opts(
722            &self,
723            location: &Path,
724            payload: PutPayload,
725            opts: PutOptions,
726        ) -> object_store::Result<PutResult> {
727            self.puts.fetch_add(1, Ordering::SeqCst);
728            let gate = self.put_gate.lock().clone();
729            if let Some(gate) = gate {
730                gate.acquire().await.unwrap().forget();
731            }
732            self.inner.put_opts(location, payload, opts).await
733        }
734
735        async fn put_multipart_opts(
736            &self,
737            location: &Path,
738            opts: PutMultipartOptions,
739        ) -> object_store::Result<Box<dyn MultipartUpload>> {
740            self.inner.put_multipart_opts(location, opts).await
741        }
742
743        async fn get_opts(
744            &self,
745            location: &Path,
746            options: GetOptions,
747        ) -> object_store::Result<GetResult> {
748            if options.head {
749                return self.inner.get_opts(location, options).await;
750            }
751            self.gets.fetch_add(1, Ordering::SeqCst);
752            let in_flight = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
753            self.max_in_flight.fetch_max(in_flight, Ordering::SeqCst);
754            let gate = self.gate.lock().clone();
755            if let Some(gate) = gate {
756                gate.acquire().await.unwrap().forget();
757            }
758            let result = if self
759                .fail_gets
760                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
761                .is_ok()
762            {
763                Err(object_store::Error::Generic {
764                    store: "TestStore",
765                    source: "transient".into(),
766                })
767            } else {
768                self.inner.get_opts(location, options).await
769            };
770            self.in_flight.fetch_sub(1, Ordering::SeqCst);
771            result
772        }
773
774        fn delete_stream(
775            &self,
776            locations: BoxStream<'static, object_store::Result<Path>>,
777        ) -> BoxStream<'static, object_store::Result<Path>> {
778            self.inner.delete_stream(locations)
779        }
780
781        fn list(
782            &self,
783            prefix: Option<&Path>,
784        ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
785            self.inner.list(prefix)
786        }
787
788        async fn list_with_delimiter(
789            &self,
790            prefix: Option<&Path>,
791        ) -> object_store::Result<ListResult> {
792            self.inner.list_with_delimiter(prefix).await
793        }
794
795        async fn copy_opts(
796            &self,
797            from: &Path,
798            to: &Path,
799            options: CopyOptions,
800        ) -> object_store::Result<()> {
801            self.inner.copy_opts(from, to, options).await
802        }
803    }
804
805    /// A [`StdVfs`] whose writes can be held.
806    #[derive(Debug, Default)]
807    struct GatedVfs {
808        inner: StdVfs,
809        /// When set, each write takes a permit first.
810        write_gate: Arc<Mutex<Option<Arc<Semaphore>>>>,
811        /// Writes that have waited on the gate.
812        held_writes: Arc<AtomicUsize>,
813    }
814
815    #[async_trait]
816    impl Vfs for GatedVfs {
817        async fn create_dir_all(&self, path: &StdPath) -> std::io::Result<()> {
818            self.inner.create_dir_all(path).await
819        }
820
821        async fn lock(&self, path: &StdPath) -> std::io::Result<Box<dyn VfsLock>> {
822            self.inner.lock(path).await
823        }
824
825        async fn list(&self, dir: &StdPath) -> std::io::Result<Vec<VfsEntry>> {
826            self.inner.list(dir).await
827        }
828
829        async fn read_range(
830            &self,
831            path: &StdPath,
832            range: std::ops::Range<u64>,
833        ) -> std::io::Result<Bytes> {
834            self.inner.read_range(path, range).await
835        }
836
837        async fn read(&self, path: &StdPath) -> std::io::Result<Bytes> {
838            self.inner.read(path).await
839        }
840
841        async fn create(&self, path: &StdPath) -> std::io::Result<Box<dyn VfsWriter>> {
842            Ok(Box::new(GatedWriter {
843                inner: self.inner.create(path).await?,
844                gate: Arc::clone(&self.write_gate),
845                held: Arc::clone(&self.held_writes),
846            }))
847        }
848
849        async fn rename(&self, from: &StdPath, to: &StdPath) -> std::io::Result<()> {
850            self.inner.rename(from, to).await
851        }
852
853        async fn remove(&self, path: &StdPath) -> std::io::Result<()> {
854            self.inner.remove(path).await
855        }
856    }
857
858    #[derive(Debug)]
859    struct GatedWriter {
860        inner: Box<dyn VfsWriter>,
861        gate: Arc<Mutex<Option<Arc<Semaphore>>>>,
862        held: Arc<AtomicUsize>,
863    }
864
865    #[async_trait]
866    impl VfsWriter for GatedWriter {
867        async fn write(&mut self, bytes: Bytes) -> std::io::Result<()> {
868            let gate = self.gate.lock().clone();
869            if let Some(gate) = gate {
870                self.held.fetch_add(1, Ordering::SeqCst);
871                gate.acquire().await.unwrap().forget();
872            }
873            self.inner.write(bytes).await
874        }
875
876        async fn finish(&mut self) -> std::io::Result<()> {
877            self.inner.finish().await
878        }
879    }
880
881    /// Routes `*.manifest` through `Observe` and `*.sst` to the mirror. A
882    /// manifest's body is a newline-separated list of paths to fetch.
883    struct TestPolicy {
884        sst_read: Mutex<ReadRoute>,
885        fail_observe: AtomicBool,
886        observed: Mutex<Vec<(Path, Bytes)>>,
887        run_started: AtomicBool,
888        run_stopped: Arc<AtomicBool>,
889    }
890
891    impl TestPolicy {
892        fn new() -> Arc<Self> {
893            Arc::new(Self {
894                sst_read: Mutex::new(ReadRoute::Local),
895                fail_observe: AtomicBool::new(false),
896                observed: Mutex::new(Vec::new()),
897                run_started: AtomicBool::new(false),
898                run_stopped: Arc::new(AtomicBool::new(false)),
899            })
900        }
901
902        fn observed(&self) -> Vec<(Path, Bytes)> {
903            self.observed.lock().clone()
904        }
905    }
906
907    struct SetOnDrop(Arc<AtomicBool>);
908
909    impl Drop for SetOnDrop {
910        fn drop(&mut self) {
911            self.0.store(true, Ordering::SeqCst);
912        }
913    }
914
915    #[async_trait]
916    impl MirrorPolicy for TestPolicy {
917        fn read_route(
918            &self,
919            path: &Path,
920            _options: &GetOptions,
921        ) -> object_store::Result<ReadRoute> {
922            Ok(match path.extension() {
923                Some("manifest") => ReadRoute::Observe,
924                Some("sst") => *self.sst_read.lock(),
925                Some("bad") => {
926                    return Err(object_store::Error::Generic {
927                        store: "TestPolicy",
928                        source: "bad route".into(),
929                    })
930                }
931                _ => ReadRoute::Remote,
932            })
933        }
934
935        fn put_route(
936            &self,
937            path: &Path,
938            _options: &PutOptions,
939        ) -> object_store::Result<WriteRoute> {
940            Ok(match path.extension() {
941                Some("manifest") => WriteRoute::Observe,
942                Some("sst") | Some("123") | Some("meta") => WriteRoute::Mirror,
943                _ => WriteRoute::Remote,
944            })
945        }
946
947        fn put_multipart_route(
948            &self,
949            path: &Path,
950            _options: &PutMultipartOptions,
951        ) -> object_store::Result<MultipartRoute> {
952            Ok(match path.extension() {
953                Some("sst") => MultipartRoute::Mirror,
954                _ => MultipartRoute::Remote,
955            })
956        }
957
958        async fn observe(
959            &self,
960            path: &Path,
961            bytes: &Bytes,
962            mirror: &MirrorHandle,
963        ) -> object_store::Result<()> {
964            self.observed.lock().push((path.clone(), bytes.clone()));
965            if self.fail_observe.load(Ordering::SeqCst) {
966                return Err(object_store::Error::Generic {
967                    store: "TestPolicy",
968                    source: "observe failed".into(),
969                });
970            }
971            let paths: Vec<Path> = std::str::from_utf8(bytes)
972                .unwrap_or_default()
973                .lines()
974                .filter(|line| !line.is_empty())
975                .map(Path::from)
976                .collect();
977            mirror.fetch(paths).await
978        }
979
980        async fn run(self: Arc<Self>, _mirror: MirrorHandle) {
981            let _stopped = SetOnDrop(Arc::clone(&self.run_stopped));
982            self.run_started.store(true, Ordering::SeqCst);
983            std::future::pending::<()>().await;
984        }
985    }
986
987    struct Fixture {
988        _dir: tempfile::TempDir,
989        root: PathBuf,
990        store: Arc<TestStore>,
991        policy: Arc<TestPolicy>,
992        mirror: Arc<ObjectStoreMirror>,
993    }
994
995    impl Fixture {
996        async fn new() -> Self {
997            Self::with_concurrency(8).await
998        }
999
1000        async fn with_concurrency(concurrency: usize) -> Self {
1001            let dir = tempfile::tempdir().unwrap();
1002            let root = dir.path().join("mirror");
1003            let store = Arc::new(TestStore::default());
1004            let policy = TestPolicy::new();
1005            let mirror = build(&root, &store, &policy, concurrency).await.unwrap();
1006            Self {
1007                _dir: dir,
1008                root,
1009                store,
1010                policy,
1011                mirror,
1012            }
1013        }
1014
1015        fn handle(&self) -> MirrorHandle {
1016            self.mirror.handle()
1017        }
1018
1019        fn files(&self) -> BTreeSet<String> {
1020            files(&self.root)
1021        }
1022
1023        fn local_bytes(&self, path: &str) -> Vec<u8> {
1024            std::fs::read(self.root.join(data_file(path))).unwrap()
1025        }
1026
1027        /// The files a set of complete local copies should leave behind.
1028        fn expected_files(&self, paths: &[&str]) -> BTreeSet<String> {
1029            let mut expected = BTreeSet::from(["LOCK".to_string()]);
1030            for path in paths {
1031                let data = data_file(path);
1032                expected.insert(format!("{data}.meta"));
1033                expected.insert(data);
1034            }
1035            expected
1036        }
1037
1038        async fn put_remote(&self, path: &str, bytes: &'static [u8]) {
1039            self.store
1040                .put(&Path::from(path), PutPayload::from_static(bytes))
1041                .await
1042                .unwrap();
1043        }
1044    }
1045
1046    async fn build(
1047        root: &StdPath,
1048        store: &Arc<TestStore>,
1049        policy: &Arc<TestPolicy>,
1050        concurrency: usize,
1051    ) -> object_store::Result<Arc<ObjectStoreMirror>> {
1052        ObjectStoreMirror::builder(
1053            root,
1054            Arc::clone(store) as Arc<dyn ObjectStore>,
1055            Arc::clone(policy) as Arc<dyn MirrorPolicy>,
1056        )
1057        .with_download_concurrency(concurrency)
1058        .with_remote_scan_interval(None)
1059        .build()
1060        .await
1061    }
1062
1063    /// The local data file name for `path`: the hex MD5 of its parent, a `.`, and
1064    /// its file name. See `src/layout.rs`.
1065    fn data_file(path: &str) -> String {
1066        let (parent, name) = path.rsplit_once('/').unwrap_or(("", path));
1067        let prefix = hex::encode(Md5::digest(parent.as_bytes()));
1068        format!("{prefix}.{name}.obj")
1069    }
1070
1071    fn files(root: &StdPath) -> BTreeSet<String> {
1072        std::fs::read_dir(root)
1073            .unwrap()
1074            .map(|entry| entry.unwrap().file_name().into_string().unwrap())
1075            .collect()
1076    }
1077
1078    fn mirror_error(err: &object_store::Error) -> &MirrorError {
1079        MirrorError::find(err).unwrap_or_else(|| panic!("expected a MirrorError, got {err:?}"))
1080    }
1081
1082    /// Polls `condition` until it holds, yielding to background tasks.
1083    async fn eventually(mut condition: impl FnMut() -> bool) {
1084        for _ in 0..1000 {
1085            if condition() {
1086                return;
1087            }
1088            tokio::time::sleep(Duration::from_millis(5)).await;
1089        }
1090        panic!("condition never held");
1091    }
1092
1093    const SST: &str = "db/compacted/01A.sst";
1094    const SST2: &str = "db/compacted/01B.sst";
1095    const MANIFEST: &str = "db/manifest/00001.manifest";
1096
1097    #[tokio::test]
1098    async fn should_reject_invalid_options() {
1099        let dir = tempfile::tempdir().unwrap();
1100        let store = Arc::new(TestStore::default());
1101        let policy = TestPolicy::new();
1102
1103        let err = build(dir.path(), &store, &policy, 0).await.unwrap_err();
1104        assert!(matches!(
1105            mirror_error(&err),
1106            MirrorError::InvalidConfig { .. }
1107        ));
1108
1109        let err = ObjectStoreMirror::builder(dir.path(), store, policy)
1110            .with_remote_scan_interval(Some(Duration::ZERO))
1111            .build()
1112            .await
1113            .unwrap_err();
1114        assert!(matches!(
1115            mirror_error(&err),
1116            MirrorError::InvalidConfig { .. }
1117        ));
1118    }
1119
1120    #[tokio::test]
1121    async fn should_lock_directory_and_stop_background_tasks_on_drop() {
1122        let fixture = Fixture::new().await;
1123        eventually(|| fixture.policy.run_started.load(Ordering::SeqCst)).await;
1124
1125        let err = build(&fixture.root, &fixture.store, &fixture.policy, 8)
1126            .await
1127            .unwrap_err();
1128        assert!(matches!(mirror_error(&err), MirrorError::Local { .. }));
1129
1130        let Fixture {
1131            _dir,
1132            root,
1133            store,
1134            policy,
1135            mirror,
1136        } = fixture;
1137        drop(mirror);
1138        eventually(|| policy.run_stopped.load(Ordering::SeqCst)).await;
1139
1140        // The lock is released once the background tasks let go of the state.
1141        let mut reopened = None;
1142        for _ in 0..100 {
1143            match build(&root, &store, &policy, 8).await {
1144                Ok(mirror) => {
1145                    reopened = Some(mirror);
1146                    break;
1147                }
1148                Err(_) => tokio::time::sleep(Duration::from_millis(5)).await,
1149            }
1150        }
1151        assert!(reopened.is_some());
1152    }
1153
1154    #[tokio::test]
1155    async fn should_pass_remote_routes_through() {
1156        let fixture = Fixture::new().await;
1157        let path = Path::from("db/wal/00001.wal");
1158        fixture
1159            .mirror
1160            .put(&path, PutPayload::from_static(b"wal"))
1161            .await
1162            .unwrap();
1163
1164        assert_eq!(
1165            fixture
1166                .mirror
1167                .get(&path)
1168                .await
1169                .unwrap()
1170                .bytes()
1171                .await
1172                .unwrap(),
1173            "wal"
1174        );
1175        assert!(!fixture.handle().contains(&path));
1176        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1177
1178        let listed: Vec<_> = fixture
1179            .mirror
1180            .list(Some(&Path::from("db")))
1181            .try_collect()
1182            .await
1183            .unwrap();
1184        assert_eq!(listed.len(), 1);
1185    }
1186
1187    #[tokio::test]
1188    async fn should_wrap_route_errors_in_policy_error() {
1189        let fixture = Fixture::new().await;
1190        let err = fixture
1191            .mirror
1192            .get(&Path::from("a/b.bad"))
1193            .await
1194            .unwrap_err();
1195        assert!(matches!(mirror_error(&err), MirrorError::Policy { .. }));
1196    }
1197
1198    #[tokio::test]
1199    async fn should_mirror_puts_and_serve_local_reads() {
1200        let fixture = Fixture::new().await;
1201        let path = Path::from(SST);
1202        let put = fixture
1203            .mirror
1204            .put(&path, PutPayload::from_static(b"0123456789"))
1205            .await
1206            .unwrap();
1207
1208        assert!(fixture.handle().contains(&path));
1209        assert_eq!(fixture.local_bytes(SST), b"0123456789");
1210        assert_eq!(fixture.files(), fixture.expected_files(&[SST]));
1211
1212        // Served locally, without touching the remote store.
1213        let result = fixture
1214            .mirror
1215            .get_opts(
1216                &path,
1217                GetOptions {
1218                    range: Some(GetRange::Bounded(2..5)),
1219                    ..Default::default()
1220                },
1221            )
1222            .await
1223            .unwrap();
1224        assert_eq!(result.range, 2..5);
1225        assert_eq!(result.bytes().await.unwrap(), "234");
1226        assert!(fixture.mirror.get_range(&path, 20..30).await.is_err());
1227        let head = fixture.mirror.head(&path).await.unwrap();
1228        assert_eq!(head.size, 10);
1229        assert_eq!(head.e_tag, put.e_tag);
1230        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 0);
1231
1232        // Preconditions are checked against the local metadata.
1233        let err = fixture
1234            .mirror
1235            .get_opts(
1236                &path,
1237                GetOptions {
1238                    if_none_match: put.e_tag.clone(),
1239                    ..Default::default()
1240                },
1241            )
1242            .await
1243            .unwrap_err();
1244        assert!(matches!(err, object_store::Error::NotModified { .. }));
1245    }
1246
1247    #[tokio::test]
1248    async fn should_return_not_local_on_local_miss() {
1249        let fixture = Fixture::new().await;
1250        fixture.put_remote(SST, b"remote only").await;
1251
1252        let err = fixture.mirror.get(&Path::from(SST)).await.unwrap_err();
1253        assert!(matches!(mirror_error(&err), MirrorError::NotLocal { .. }));
1254        let err = fixture.mirror.head(&Path::from(SST)).await.unwrap_err();
1255        assert!(matches!(mirror_error(&err), MirrorError::NotLocal { .. }));
1256        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 0);
1257    }
1258
1259    #[tokio::test]
1260    async fn should_not_write_remotely_when_put_fails_remotely() {
1261        let fixture = Fixture::new().await;
1262        fixture.put_remote(SST, b"first").await;
1263
1264        let err = fixture
1265            .mirror
1266            .put_opts(
1267                &Path::from(SST),
1268                PutPayload::from_static(b"second"),
1269                PutOptions {
1270                    mode: PutMode::Create,
1271                    ..Default::default()
1272                },
1273            )
1274            .await
1275            .unwrap_err();
1276        assert!(matches!(err, object_store::Error::AlreadyExists { .. }));
1277        assert!(!fixture.handle().contains(&Path::from(SST)));
1278        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1279    }
1280
1281    #[tokio::test]
1282    async fn should_mirror_names_that_look_like_layout_suffixes() {
1283        let fixture = Fixture::new().await;
1284        eventually(|| fixture.policy.run_started.load(Ordering::SeqCst)).await;
1285        // `01A.sst.meta` and `01A.sst.123` look like the metadata and a temp
1286        // file of `01A.sst`.
1287        let objects: [(&str, &'static [u8]); 3] = [
1288            (SST, b"sst"),
1289            ("db/compacted/01A.sst.meta", b"meta"),
1290            ("db/compacted/01A.sst.123", b"temp"),
1291        ];
1292        for (path, bytes) in objects {
1293            fixture
1294                .mirror
1295                .put(&Path::from(path), PutPayload::from_static(bytes))
1296                .await
1297                .unwrap();
1298        }
1299        let paths = objects.map(|(path, _)| path);
1300        assert_eq!(fixture.files(), fixture.expected_files(&paths));
1301        for (path, bytes) in objects {
1302            assert_eq!(fixture.local_bytes(path), bytes, "{path}");
1303        }
1304
1305        let Fixture {
1306            _dir,
1307            root,
1308            store,
1309            policy,
1310            mirror,
1311        } = fixture;
1312        drop(mirror);
1313        eventually(|| policy.run_stopped.load(Ordering::SeqCst)).await;
1314        let mut reopened = None;
1315        for _ in 0..100 {
1316            if let Ok(mirror) = build(&root, &store, &policy, 8).await {
1317                reopened = Some(mirror);
1318                break;
1319            }
1320            tokio::time::sleep(Duration::from_millis(5)).await;
1321        }
1322        let mirror = reopened.expect("mirror reopened");
1323
1324        for path in paths {
1325            assert!(mirror.handle().contains(&Path::from(path)), "{path}");
1326        }
1327        assert_eq!(
1328            mirror
1329                .get(&Path::from(SST))
1330                .await
1331                .unwrap()
1332                .bytes()
1333                .await
1334                .unwrap(),
1335            "sst"
1336        );
1337        assert_eq!(store.gets.load(Ordering::SeqCst), 0);
1338    }
1339
1340    #[tokio::test]
1341    async fn should_observe_whole_objects_on_ranged_reads() {
1342        let fixture = Fixture::new().await;
1343        fixture.put_remote(SST, b"sst").await;
1344        fixture
1345            .put_remote(MANIFEST, b"db/compacted/01A.sst\n")
1346            .await;
1347        let manifest = Path::from(MANIFEST);
1348
1349        let result = fixture
1350            .mirror
1351            .get_opts(
1352                &manifest,
1353                GetOptions {
1354                    range: Some(GetRange::Bounded(0..2)),
1355                    ..Default::default()
1356                },
1357            )
1358            .await
1359            .unwrap();
1360        assert_eq!(result.range, 0..2);
1361        assert_eq!(result.bytes().await.unwrap(), "db");
1362        assert_eq!(
1363            fixture.policy.observed(),
1364            vec![(manifest.clone(), Bytes::from("db/compacted/01A.sst\n"))]
1365        );
1366
1367        // `observe` fetched the SST it listed.
1368        assert!(fixture.handle().contains(&Path::from(SST)));
1369        assert_eq!(fixture.local_bytes(SST), b"sst");
1370
1371        // HEAD and bodiless conditional GETs skip `observe`.
1372        let meta = fixture.mirror.head(&manifest).await.unwrap();
1373        let err = fixture
1374            .mirror
1375            .get_opts(
1376                &manifest,
1377                GetOptions {
1378                    if_none_match: meta.e_tag,
1379                    ..Default::default()
1380                },
1381            )
1382            .await
1383            .unwrap_err();
1384        assert!(matches!(err, object_store::Error::NotModified { .. }));
1385        assert_eq!(fixture.policy.observed().len(), 1);
1386    }
1387
1388    #[tokio::test]
1389    async fn should_report_observe_failures() {
1390        let fixture = Fixture::new().await;
1391        fixture.policy.fail_observe.store(true, Ordering::SeqCst);
1392        let manifest = Path::from(MANIFEST);
1393
1394        // Writes are durable, so the caller has to know.
1395        let err = fixture
1396            .mirror
1397            .put(&manifest, PutPayload::from_static(b""))
1398            .await
1399            .unwrap_err();
1400        assert!(matches!(
1401            mirror_error(&err),
1402            MirrorError::WriteCommitted { .. }
1403        ));
1404        assert!(fixture.store.head(&manifest).await.is_ok());
1405
1406        // Reads return the policy's error unchanged.
1407        let err = fixture.mirror.get(&manifest).await.unwrap_err();
1408        assert!(MirrorError::find(&err).is_none());
1409        assert!(err.to_string().contains("observe failed"));
1410    }
1411
1412    #[tokio::test]
1413    async fn should_refetch_and_replace_local_copy() {
1414        let fixture = Fixture::new().await;
1415        let path = Path::from(SST);
1416        fixture
1417            .mirror
1418            .put(&path, PutPayload::from_static(b"good bytes"))
1419            .await
1420            .unwrap();
1421
1422        // Corrupt the local copy without changing its size.
1423        std::fs::write(fixture.root.join(data_file(SST)), b"bad bytes!").unwrap();
1424
1425        *fixture.policy.sst_read.lock() = ReadRoute::Refetch;
1426        assert_eq!(fixture.mirror.get_range(&path, 0..4).await.unwrap(), "good");
1427        assert_eq!(fixture.local_bytes(SST), b"good bytes");
1428        assert_eq!(fixture.mirror.head(&path).await.unwrap().size, 10);
1429        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 2);
1430        assert_eq!(fixture.files(), fixture.expected_files(&[SST]));
1431    }
1432
1433    #[tokio::test]
1434    async fn should_tee_multipart_uploads() {
1435        let fixture = Fixture::new().await;
1436        let path = Path::from(SST);
1437        let mut upload = fixture.mirror.put_multipart(&path).await.unwrap();
1438
1439        // Complete the parts out of order; the local file still gets them in
1440        // order.
1441        let first = upload.put_part(PutPayload::from_static(b"aaaaa"));
1442        let second = upload.put_part(PutPayload::from_static(b"bbbbb"));
1443        let third = upload.put_part(PutPayload::from_static(b"ccccc"));
1444        third.await.unwrap();
1445        second.await.unwrap();
1446        first.await.unwrap();
1447        assert!(!fixture.handle().contains(&path));
1448        upload.complete().await.unwrap();
1449
1450        assert!(fixture.handle().contains(&path));
1451        assert_eq!(fixture.local_bytes(SST), b"aaaaabbbbbccccc");
1452        assert_eq!(
1453            fixture
1454                .store
1455                .get(&path)
1456                .await
1457                .unwrap()
1458                .bytes()
1459                .await
1460                .unwrap(),
1461            "aaaaabbbbbccccc"
1462        );
1463        assert_eq!(fixture.mirror.head(&path).await.unwrap().size, 15);
1464        assert_eq!(fixture.files(), fixture.expected_files(&[SST]));
1465    }
1466
1467    #[tokio::test]
1468    async fn should_remove_temp_file_when_multipart_upload_is_aborted_or_dropped() {
1469        let fixture = Fixture::new().await;
1470        let path = Path::from(SST);
1471
1472        let mut upload = fixture.mirror.put_multipart(&path).await.unwrap();
1473        upload
1474            .put_part(PutPayload::from_static(b"part"))
1475            .await
1476            .unwrap();
1477        upload.abort().await.unwrap();
1478        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1479
1480        let mut upload = fixture.mirror.put_multipart(&path).await.unwrap();
1481        upload
1482            .put_part(PutPayload::from_static(b"part"))
1483            .await
1484            .unwrap();
1485        drop(upload);
1486        eventually(|| fixture.files() == fixture.expected_files(&[])).await;
1487
1488        // The dropped upload released the path.
1489        fixture
1490            .mirror
1491            .put(&path, PutPayload::from_static(b"x"))
1492            .await
1493            .unwrap();
1494        assert!(fixture.handle().contains(&path));
1495    }
1496
1497    #[tokio::test]
1498    async fn should_abort_multipart_upload_with_a_cancelled_part_mid_write() {
1499        let dir = tempfile::tempdir().unwrap();
1500        let root = dir.path().join("mirror");
1501        let store = Arc::new(TestStore::default());
1502        let policy = TestPolicy::new();
1503        let vfs = Arc::new(GatedVfs::default());
1504        let mirror = ObjectStoreMirror::builder(
1505            &root,
1506            Arc::clone(&store) as Arc<dyn ObjectStore>,
1507            Arc::clone(&policy) as Arc<dyn MirrorPolicy>,
1508        )
1509        .with_remote_scan_interval(None)
1510        .with_vfs(Arc::clone(&vfs) as Arc<dyn Vfs>)
1511        .build()
1512        .await
1513        .unwrap();
1514        let path = Path::from(SST);
1515        let mut upload = mirror.put_multipart(&path).await.unwrap();
1516
1517        // Cancel a part while its local write is held, the way
1518        // `WriteMultipart` drops its part tasks before aborting.
1519        let gate = Arc::new(Semaphore::new(0));
1520        *vfs.write_gate.lock() = Some(Arc::clone(&gate));
1521        let part = tokio::spawn(upload.put_part(PutPayload::from_static(b"part")));
1522        eventually(|| vfs.held_writes.load(Ordering::SeqCst) == 1).await;
1523        part.abort();
1524        let _ = part.await;
1525
1526        tokio::time::timeout(Duration::from_secs(5), upload.abort())
1527            .await
1528            .expect("abort hung")
1529            .unwrap();
1530        eventually(|| files(&root) == BTreeSet::from(["LOCK".to_string()])).await;
1531
1532        // The aborted upload released the path.
1533        *vfs.write_gate.lock() = None;
1534        mirror
1535            .put(&path, PutPayload::from_static(b"x"))
1536            .await
1537            .unwrap();
1538        assert!(mirror.handle().contains(&path));
1539    }
1540
1541    #[tokio::test]
1542    async fn should_delete_local_and_remote_copies() {
1543        let fixture = Fixture::new().await;
1544        let path = Path::from(SST);
1545        fixture
1546            .mirror
1547            .put(&path, PutPayload::from_static(b"x"))
1548            .await
1549            .unwrap();
1550
1551        fixture.mirror.delete(&path).await.unwrap();
1552        assert!(!fixture.handle().contains(&path));
1553        assert!(fixture.store.head(&path).await.is_err());
1554        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1555    }
1556
1557    #[tokio::test]
1558    async fn should_not_delete_remotely_before_earlier_mirrored_put() {
1559        let fixture = Fixture::new().await;
1560        let path = Path::from(SST);
1561        let gate = Arc::new(Semaphore::new(0));
1562        *fixture.store.put_gate.lock() = Some(Arc::clone(&gate));
1563
1564        let mirror = Arc::clone(&fixture.mirror);
1565        let put_path = path.clone();
1566        let put =
1567            tokio::spawn(async move { mirror.put(&put_path, PutPayload::from_static(b"x")).await });
1568        eventually(|| fixture.store.puts.load(Ordering::SeqCst) == 1).await;
1569
1570        let mirror = Arc::clone(&fixture.mirror);
1571        let delete_path = path.clone();
1572        let delete = tokio::spawn(async move { mirror.delete(&delete_path).await });
1573        tokio::time::sleep(Duration::from_millis(20)).await;
1574        assert!(!delete.is_finished());
1575
1576        gate.add_permits(1);
1577        put.await.unwrap().unwrap();
1578        delete.await.unwrap().unwrap();
1579        assert!(fixture.store.head(&path).await.is_err());
1580        assert!(!fixture.handle().contains(&path));
1581        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1582    }
1583
1584    #[tokio::test]
1585    async fn should_reject_copy_and_rename() {
1586        let fixture = Fixture::new().await;
1587        fixture.put_remote("a/b", b"x").await;
1588        let (from, to) = (Path::from("a/b"), Path::from("a/c"));
1589
1590        let err = fixture.mirror.copy(&from, &to).await.unwrap_err();
1591        assert!(matches!(
1592            mirror_error(&err),
1593            MirrorError::Unsupported { .. }
1594        ));
1595        let err = fixture.mirror.rename(&from, &to).await.unwrap_err();
1596        assert!(matches!(
1597            mirror_error(&err),
1598            MirrorError::Unsupported { .. }
1599        ));
1600        assert!(fixture.store.head(&to).await.is_err());
1601    }
1602
1603    #[tokio::test]
1604    async fn should_restore_local_copies_after_restart() {
1605        let fixture = Fixture::new().await;
1606        eventually(|| fixture.policy.run_started.load(Ordering::SeqCst)).await;
1607        let path = Path::from(SST);
1608        let put = fixture
1609            .mirror
1610            .put(&path, PutPayload::from_static(b"persisted"))
1611            .await
1612            .unwrap();
1613
1614        let Fixture {
1615            _dir,
1616            root,
1617            store,
1618            policy,
1619            mirror,
1620        } = fixture;
1621        drop(mirror);
1622        eventually(|| policy.run_stopped.load(Ordering::SeqCst)).await;
1623        let mut reopened = None;
1624        for _ in 0..100 {
1625            if let Ok(mirror) = build(&root, &store, &policy, 8).await {
1626                reopened = Some(mirror);
1627                break;
1628            }
1629            tokio::time::sleep(Duration::from_millis(5)).await;
1630        }
1631        let mirror = reopened.expect("mirror reopened");
1632
1633        assert!(mirror.handle().contains(&path));
1634        assert_eq!(mirror.head(&path).await.unwrap().e_tag, put.e_tag);
1635        assert_eq!(
1636            mirror.get(&path).await.unwrap().bytes().await.unwrap(),
1637            "persisted"
1638        );
1639        assert_eq!(store.gets.load(Ordering::SeqCst), 0);
1640    }
1641
1642    #[tokio::test]
1643    async fn should_refetch_corrupt_local_copy_after_restart() {
1644        let fixture = Fixture::new().await;
1645        // The policy task must create its shutdown guard before the mirror is dropped.
1646        eventually(|| fixture.policy.run_started.load(Ordering::SeqCst)).await;
1647        let path = Path::from(SST);
1648        fixture
1649            .mirror
1650            .put(&path, PutPayload::from_static(b"good bytes"))
1651            .await
1652            .unwrap();
1653
1654        let Fixture {
1655            _dir,
1656            root,
1657            store,
1658            policy,
1659            mirror,
1660        } = fixture;
1661        drop(mirror);
1662        eventually(|| policy.run_stopped.load(Ordering::SeqCst)).await;
1663
1664        // Model corruption that startup's file-size check cannot detect.
1665        std::fs::write(root.join(data_file(SST)), b"bad bytes!").unwrap();
1666        let mut reopened = None;
1667        for _ in 0..100 {
1668            if let Ok(mirror) = build(&root, &store, &policy, 8).await {
1669                reopened = Some(mirror);
1670                break;
1671            }
1672            tokio::time::sleep(Duration::from_millis(5)).await;
1673        }
1674        let mirror = reopened.expect("mirror reopened");
1675
1676        // Local reads succeed with corrupt bytes, even when repeated.
1677        assert!(mirror.handle().contains(&path));
1678        for _ in 0..2 {
1679            assert_eq!(
1680                mirror.get(&path).await.unwrap().bytes().await.unwrap(),
1681                "bad bytes!"
1682            );
1683        }
1684        assert_eq!(store.gets.load(Ordering::SeqCst), 0);
1685
1686        // The caller signals its policy to refetch, then checks the result.
1687        *policy.sst_read.lock() = ReadRoute::Refetch;
1688        assert_eq!(
1689            mirror.get(&path).await.unwrap().bytes().await.unwrap(),
1690            "good bytes"
1691        );
1692        *policy.sst_read.lock() = ReadRoute::Local;
1693        assert_eq!(
1694            mirror.get(&path).await.unwrap().bytes().await.unwrap(),
1695            "good bytes"
1696        );
1697        assert_eq!(store.gets.load(Ordering::SeqCst), 1);
1698    }
1699
1700    #[tokio::test]
1701    async fn should_dedupe_concurrent_fetches() {
1702        let fixture = Fixture::new().await;
1703        fixture.put_remote(SST, b"x").await;
1704        let gate = Arc::new(Semaphore::new(0));
1705        *fixture.store.gate.lock() = Some(Arc::clone(&gate));
1706
1707        let first = tokio::spawn(fixture.handle().fetch([Path::from(SST)]));
1708        let second = tokio::spawn(fixture.handle().fetch([Path::from(SST)]));
1709        eventually(|| fixture.store.gets.load(Ordering::SeqCst) == 1).await;
1710        gate.add_permits(10);
1711        first.await.unwrap().unwrap();
1712        second.await.unwrap().unwrap();
1713
1714        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 1);
1715        assert!(fixture.handle().contains(&Path::from(SST)));
1716
1717        // Already local, so no download.
1718        fixture.handle().fetch([Path::from(SST)]).await.unwrap();
1719        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 1);
1720    }
1721
1722    #[tokio::test]
1723    async fn should_limit_download_concurrency() {
1724        let fixture = Fixture::with_concurrency(1).await;
1725        fixture.put_remote(SST, b"a").await;
1726        fixture.put_remote(SST2, b"b").await;
1727        let gate = Arc::new(Semaphore::new(0));
1728        *fixture.store.gate.lock() = Some(Arc::clone(&gate));
1729
1730        let fetch = tokio::spawn(fixture.handle().fetch([Path::from(SST), Path::from(SST2)]));
1731        eventually(|| fixture.store.gets.load(Ordering::SeqCst) == 1).await;
1732        tokio::time::sleep(Duration::from_millis(20)).await;
1733        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 1);
1734
1735        gate.add_permits(10);
1736        fetch.await.unwrap().unwrap();
1737        assert_eq!(fixture.store.max_in_flight.load(Ordering::SeqCst), 1);
1738        assert!(fixture.handle().contains(&Path::from(SST2)));
1739    }
1740
1741    #[tokio::test]
1742    async fn should_retry_transient_download_errors() {
1743        let fixture = Fixture::new().await;
1744        fixture.put_remote(SST, b"x").await;
1745        fixture.store.fail_gets.store(2, Ordering::SeqCst);
1746
1747        fixture.handle().fetch([Path::from(SST)]).await.unwrap();
1748        assert_eq!(fixture.store.gets.load(Ordering::SeqCst), 3);
1749        assert_eq!(fixture.files(), fixture.expected_files(&[SST]));
1750    }
1751
1752    #[tokio::test]
1753    async fn should_fail_fetch_of_missing_object() {
1754        let fixture = Fixture::new().await;
1755        let err = fixture.handle().fetch([Path::from(SST)]).await.unwrap_err();
1756        assert!(matches!(err, object_store::Error::NotFound { .. }));
1757        assert_eq!(fixture.files(), fixture.expected_files(&[]));
1758    }
1759
1760    #[tokio::test]
1761    async fn should_evict_local_copies() {
1762        let fixture = Fixture::new().await;
1763        let path = Path::from(SST);
1764        fixture
1765            .mirror
1766            .put(&path, PutPayload::from_static(b"x"))
1767            .await
1768            .unwrap();
1769
1770        fixture.handle().evict([path.clone()]);
1771        eventually(|| !fixture.handle().contains(&path)).await;
1772        eventually(|| fixture.files() == fixture.expected_files(&[])).await;
1773        // The remote copy is untouched.
1774        assert!(fixture.store.head(&path).await.is_ok());
1775    }
1776
1777    #[tokio::test]
1778    async fn should_run_remote_scan_in_background() {
1779        let dir = tempfile::tempdir().unwrap();
1780        let store = Arc::new(TestStore::default());
1781        let policy = TestPolicy::new();
1782        let mirror = ObjectStoreMirror::builder(
1783            dir.path(),
1784            Arc::clone(&store) as Arc<dyn ObjectStore>,
1785            policy,
1786        )
1787        .with_remote_scan_interval(Some(Duration::from_millis(10)))
1788        .build()
1789        .await
1790        .unwrap();
1791        let path = Path::from(SST);
1792        mirror
1793            .put(&path, PutPayload::from_static(b"x"))
1794            .await
1795            .unwrap();
1796
1797        store.delete(&path).await.unwrap();
1798        eventually(|| !mirror.handle().contains(&path)).await;
1799    }
1800}