log_store/object_store_wal/
io.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Object store access for WAL objects: a deterministic key layout and
16//! conditional creates that make a retry of the same object a no-op.
17
18use async_stream::try_stream;
19use bytes::Bytes;
20use common_telemetry::warn;
21use futures::TryStreamExt;
22use futures::stream::BoxStream;
23use object_store::{Entry, ErrorKind, ObjectStore};
24use snafu::{OptionExt, ResultExt, ensure};
25use tokio::sync::mpsc;
26use tokio_util::task::AbortOnDropHandle;
27use uuid::Uuid;
28
29use crate::error::{
30    CorruptedWalObjectSnafu, InvalidWalObjectStoreSnafu, Result, WalObjectConflictSnafu,
31    WalObjectStoreSnafu,
32};
33
34/// Width of the zero-padded object sequence in an object key, wide enough for
35/// [`u64::MAX`] so that lexicographic and numeric order agree.
36const OBJECT_SEQ_WIDTH: usize = 20;
37const OBJECT_SUFFIX: &str = ".wal";
38const PROBE_CONTENT: &[u8] = b"conditional create probe";
39/// Objects a listing may hold ahead of the reads that take them, four LIST
40/// pages of S3, so the listing stays pages ahead of the reads while what it
41/// holds is bounded.
42const LISTING_AHEAD_OBJECTS: usize = 4096;
43
44/// Outcome of a conditional create.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub(crate) enum PutResult {
47    /// This call created the object.
48    Created,
49    /// The object already held the same content, so the call was a retry.
50    AlreadyPresent,
51}
52
53/// An object discovered by a LIST.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub(crate) struct ListedObject {
56    pub(crate) object_seq: u64,
57    pub(crate) path: String,
58    /// Length of the object in bytes, as reported by the listing.
59    pub(crate) size: u64,
60}
61
62/// Reads and writes the WAL objects under one prefix.
63#[derive(Clone)]
64pub(crate) struct ObjectStoreIo {
65    store: ObjectStore,
66    object_prefix: String,
67    /// Holds the keys of [`Self::probe_conditional_create`], outside
68    /// `object_prefix` so that no listing of the objects returns one.
69    probe_prefix: String,
70}
71
72impl ObjectStoreIo {
73    /// Binds `store` to the objects under `prefix`. The store must support
74    /// conditional creates, which is what keeps a retried write from
75    /// overwriting a durable object.
76    pub(crate) fn new(store: ObjectStore, prefix: impl AsRef<str>) -> Result<Self> {
77        let prefix = normalize_prefix(prefix.as_ref())?;
78        ensure!(
79            store.info().capability().write_with_if_not_exists,
80            InvalidWalObjectStoreSnafu {
81                reason: "object store does not support conditional create",
82            }
83        );
84        Ok(Self {
85            store,
86            object_prefix: format!("{prefix}/objects/"),
87            probe_prefix: format!("{prefix}/probes/"),
88        })
89    }
90
91    /// Checks that the store refuses a conditional create of an object that
92    /// exists. The capability checked by [`Self::new`] is declared per
93    /// service, and a store behind the S3 service may accept the condition and
94    /// overwrite the object anyway, which would let a retried write replace a
95    /// durable object. The probe creates a key of its own, unique to this
96    /// call, twice, and then deletes it.
97    pub(crate) async fn probe_conditional_create(&self) -> Result<()> {
98        let path = format!("{}{}", self.probe_prefix, Uuid::new_v4());
99        self.put_path_if_absent(path.clone(), Bytes::from_static(PROBE_CONTENT))
100            .await?;
101        // Sent directly: only the store's own refusal shows that it honours
102        // the condition, and a read-back would find the first create's object.
103        let recreated = self
104            .store
105            .write_with(&path, PROBE_CONTENT)
106            .if_not_exists(true)
107            .await;
108        if let Err(error) = self.store.delete(&path).await {
109            warn!(error; "Failed to delete the conditional create probe {path} of the object store WAL");
110        }
111        match recreated {
112            // The first create stored this key, so `AlreadyExists` can only
113            // refer to it.
114            Err(error)
115                if matches!(
116                    error.kind(),
117                    ErrorKind::ConditionNotMatch | ErrorKind::AlreadyExists
118                ) =>
119            {
120                Ok(())
121            }
122            Err(error) => Err(error).context(WalObjectStoreSnafu {
123                operation: "write",
124                path,
125            }),
126            Ok(_) => InvalidWalObjectStoreSnafu {
127                reason: format!(
128                    "the object store ({}) overwrote the existing object {path} on a conditional create, but the object store WAL requires conditional create support; a versioned bucket or the S3-compatible endpoint of GCS can cause this",
129                    self.storage()
130                ),
131            }
132            .fail(),
133        }
134    }
135
136    /// Creates the object `object_seq` unless it already exists, see
137    /// [`Self::put_path_if_absent`].
138    pub(crate) async fn put_if_absent(&self, object_seq: u64, content: Bytes) -> Result<PutResult> {
139        self.put_path_if_absent(self.object_path(object_seq), content)
140            .await
141    }
142
143    /// Creates `path` unless it already exists. A retry with the same content
144    /// is accepted, different content is rejected: a create the store stored
145    /// but whose response was lost can be retried, by the caller or by the
146    /// retry layer, and then finds its own object.
147    ///
148    /// When the store reports that the object exists but the read that
149    /// compares its content fails, the read failure is returned, so that its
150    /// retry hint says whether comparing again may settle the outcome. Any
151    /// other create failure is returned as the write failure it is.
152    async fn put_path_if_absent(&self, path: String, content: Bytes) -> Result<PutResult> {
153        let write_result = self
154            .store
155            .write_with(&path, content.clone())
156            .if_not_exists(true)
157            .await;
158
159        match write_result {
160            Ok(_) => Ok(PutResult::Created),
161            Err(error) => match self.store.read(&path).await {
162                Ok(existing) if existing.to_bytes() == content => Ok(PutResult::AlreadyPresent),
163                Ok(_) => WalObjectConflictSnafu { path }.fail(),
164                Err(read_error) if reports_existing_object(&error) => {
165                    Err(read_error).context(WalObjectStoreSnafu {
166                        operation: "read",
167                        path,
168                    })
169                }
170                Err(_) => Err(error).context(WalObjectStoreSnafu {
171                    operation: "write",
172                    path,
173                }),
174            },
175        }
176    }
177
178    /// Reads the object `object_seq`.
179    pub(crate) async fn get(&self, object_seq: u64) -> Result<Bytes> {
180        let path = self.object_path(object_seq);
181        self.store
182            .read(&path)
183            .await
184            .map(|content| content.to_bytes())
185            .context(WalObjectStoreSnafu {
186                operation: "read",
187                path,
188            })
189    }
190
191    /// Reads `len` bytes of the object `object_seq` starting at `offset`. The
192    /// range must lie inside the object; one that reaches past its end fails.
193    pub(crate) async fn get_range(&self, object_seq: u64, offset: u64, len: u64) -> Result<Bytes> {
194        let path = self.object_path(object_seq);
195        let end = offset
196            .checked_add(len)
197            .with_context(|| CorruptedWalObjectSnafu {
198                reason: format!("byte range {offset}..{len} overflows the object"),
199            })?;
200        self.store
201            .read_with(&path)
202            .range(offset..end)
203            .await
204            .map(|content| content.to_bytes())
205            .context(WalObjectStoreSnafu {
206                operation: "read",
207                path,
208            })
209    }
210
211    /// Deletes the object `object_seq`. Deleting an object that does not
212    /// exist succeeds.
213    pub(crate) async fn delete(&self, object_seq: u64) -> Result<()> {
214        let path = self.object_path(object_seq);
215        self.store.delete(&path).await.context(WalObjectStoreSnafu {
216            operation: "delete",
217            path,
218        })
219    }
220
221    /// Lists the objects under the prefix, ordered by object sequence. Keys
222    /// that do not follow the object layout are ignored.
223    pub(crate) async fn list(&self) -> Result<Vec<ListedObject>> {
224        let entries = self
225            .store
226            .list(&self.object_prefix)
227            .await
228            .with_context(|_| WalObjectStoreSnafu {
229                operation: "list",
230                path: self.object_prefix.clone(),
231            })?;
232        let mut objects = entries
233            .iter()
234            .filter_map(|entry| self.listed_object(entry))
235            .collect::<Vec<_>>();
236        objects.sort_unstable_by_key(|object| object.object_seq);
237        Ok(objects)
238    }
239
240    /// Lists the objects under the prefix like [`list`](Self::list), but from
241    /// a task of its own and in the order the object store returns them, so
242    /// the listing goes on while the caller consumes the objects listed so
243    /// far. A listing that fails, however it fails, ends the stream with an
244    /// error, before any object it listed that the caller has not taken yet;
245    /// dropping the stream stops the listing.
246    pub(crate) fn list_ahead(&self) -> BoxStream<'static, Result<ListedObject>> {
247        let (listing, rx) = self.clone().listing();
248        let listing = AbortOnDropHandle::new(common_runtime::spawn_global(listing));
249        settle_listing(rx, listing, self.object_prefix.clone())
250    }
251
252    /// Returns the listing of the objects under the prefix and the receiver it
253    /// sends them to; it waits while the receiver holds
254    /// [`LISTING_AHEAD_OBJECTS`] of them.
255    fn listing(
256        self,
257    ) -> (
258        impl Future<Output = object_store::Result<()>> + Send + 'static,
259        mpsc::Receiver<ListedObject>,
260    ) {
261        let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS);
262        let listing = async move {
263            let mut lister = self.store.lister(&self.object_prefix).await?;
264            while let Some(entry) = lister.try_next().await? {
265                // The receiver is gone only once the stream was dropped.
266                if let Some(object) = self.listed_object(&entry)
267                    && tx.send(object).await.is_err()
268                {
269                    break;
270                }
271            }
272            Ok::<_, object_store::Error>(())
273        };
274        (listing, rx)
275    }
276
277    pub(crate) fn object_path(&self, object_seq: u64) -> String {
278        format!(
279            "{}{object_seq:0OBJECT_SEQ_WIDTH$}{OBJECT_SUFFIX}",
280            self.object_prefix
281        )
282    }
283
284    /// Describes the object store the objects live in by its scheme, its
285    /// bucket when it has one and its root, none of which is a credential.
286    pub(crate) fn storage(&self) -> String {
287        let info = self.store.info();
288        let name = info.name();
289        if name.is_empty() {
290            format!("{}, root {:?}", info.scheme(), info.root())
291        } else {
292            format!("{}, bucket {name:?}, root {:?}", info.scheme(), info.root())
293        }
294    }
295
296    fn listed_object(&self, entry: &Entry) -> Option<ListedObject> {
297        self.parse_object_seq(entry.path())
298            .map(|object_seq| ListedObject {
299                object_seq,
300                path: entry.path().to_string(),
301                size: entry.metadata().content_length(),
302            })
303    }
304
305    fn parse_object_seq(&self, path: &str) -> Option<u64> {
306        let value = path
307            .strip_prefix(&self.object_prefix)?
308            .strip_suffix(OBJECT_SUFFIX)?;
309        if value.len() != OBJECT_SEQ_WIDTH || !value.bytes().all(|byte| byte.is_ascii_digit()) {
310            return None;
311        }
312        value.parse().ok()
313    }
314}
315
316/// Yields the objects `listing` sends over `rx`. Once the listing finished it
317/// is settled before any further object is yielded, so its failure goes
318/// before the objects it listed that the caller has not taken yet.
319fn settle_listing(
320    mut rx: mpsc::Receiver<ListedObject>,
321    listing: AbortOnDropHandle<object_store::Result<()>>,
322    path: String,
323) -> BoxStream<'static, Result<ListedObject>> {
324    let mut listing = Some(listing);
325    Box::pin(try_stream! {
326        loop {
327            let object = rx.recv().await;
328            if let Some(listing) =
329                listing.take_if(|listing| object.is_none() || listing.is_finished())
330            {
331                listing
332                    .await
333                    .unwrap_or_else(|error| {
334                        let stopped = "the listing task did not complete";
335                        Err(object_store::Error::new(ErrorKind::Unexpected, stopped)
336                            .set_source(error))
337                    })
338                    .context(WalObjectStoreSnafu { operation: "list", path: &path })?;
339            }
340            let Some(object) = object else {
341                break;
342            };
343            yield object;
344        }
345    })
346}
347
348/// Returns true when a conditional create failed because the object exists.
349/// Only the precondition failure says so: a store may report `AlreadyExists`
350/// for an unrelated path, such as a parent that is a file.
351fn reports_existing_object(error: &object_store::Error) -> bool {
352    error.kind() == ErrorKind::ConditionNotMatch
353}
354
355fn normalize_prefix(prefix: &str) -> Result<String> {
356    let prefix = prefix.trim();
357    ensure!(
358        !prefix.is_empty() && !prefix.starts_with('/'),
359        InvalidWalObjectStoreSnafu {
360            reason: format!("object prefix {prefix:?} is empty or absolute"),
361        }
362    );
363    let prefix = prefix.strip_suffix('/').unwrap_or(prefix);
364    ensure!(
365        !prefix
366            .split('/')
367            .any(|segment| segment.is_empty() || segment == "." || segment == ".."),
368        InvalidWalObjectStoreSnafu {
369            reason: format!("object prefix {prefix:?} has an empty or relative component"),
370        }
371    );
372    Ok(prefix.to_string())
373}
374
375#[cfg(test)]
376mod tests {
377    use std::pin::pin;
378    use std::sync::Arc;
379    use std::sync::atomic::{AtomicUsize, Ordering};
380
381    use common_error::ext::{ErrorExt, RetryHint};
382    use common_test_util::temp_dir::create_temp_dir;
383    use futures::poll;
384    use object_store::layers::mock::{self, MockLayerBuilder, oio};
385    use object_store::secure_fs::SecureFsRoot;
386    use object_store::services::Memory;
387
388    use super::*;
389    use crate::error::Error;
390
391    fn memory_store() -> ObjectStore {
392        ObjectStore::new(Memory::default()).unwrap()
393    }
394
395    /// A reader whose every read fails with a temporary error.
396    struct FailingReader;
397
398    impl oio::Read for FailingReader {
399        async fn open(
400            &self,
401            _range: mock::BytesRange,
402        ) -> mock::Result<(mock::RpRead, Box<dyn oio::ReadStreamDyn>)> {
403            Err(injected_failure("read").set_temporary())
404        }
405
406        async fn read(
407            &self,
408            _range: mock::BytesRange,
409        ) -> mock::Result<(mock::RpRead, mock::Buffer)> {
410            Err(injected_failure("read").set_temporary())
411        }
412    }
413
414    /// A writer that fails to close, so nothing is written.
415    struct FailingWriter;
416
417    impl oio::Write for FailingWriter {
418        async fn write(&mut self, _buffer: mock::Buffer) -> mock::Result<()> {
419            Ok(())
420        }
421
422        async fn close(&mut self) -> mock::Result<mock::Metadata> {
423            Err(injected_failure("write"))
424        }
425
426        async fn abort(&mut self) -> mock::Result<()> {
427            Ok(())
428        }
429    }
430
431    fn injected_failure(operation: &str) -> mock::Error {
432        mock::Error::new(
433            mock::ErrorKind::Unexpected,
434            format!("injected {operation} failure"),
435        )
436    }
437
438    fn with_failing_reads(store: ObjectStore) -> ObjectStore {
439        store.layer(
440            MockLayerBuilder::default()
441                .reader_factory(Arc::new(|_, _, _| Box::new(FailingReader)))
442                .build()
443                .unwrap(),
444        )
445    }
446
447    fn with_failing_writes(store: ObjectStore) -> ObjectStore {
448        store.layer(
449            MockLayerBuilder::default()
450                .writer_factory(Arc::new(|_, _, _| Box::new(FailingWriter)))
451                .build()
452                .unwrap(),
453        )
454    }
455
456    fn memory_io() -> ObjectStoreIo {
457        ObjectStoreIo::new(memory_store(), "datanodes/1/epochs/2/").unwrap()
458    }
459
460    #[test]
461    fn test_io_builds_and_parses_deterministic_object_paths() {
462        let io = ObjectStoreIo::new(memory_store(), " datanodes/1/epochs/2/ ").unwrap();
463        assert_eq!(
464            "datanodes/1/epochs/2/objects/00000000000000000000.wal",
465            io.object_path(0)
466        );
467        assert_eq!(
468            "datanodes/1/epochs/2/objects/00000000000000000042.wal",
469            io.object_path(42)
470        );
471        assert_eq!(
472            "datanodes/1/epochs/2/objects/18446744073709551615.wal",
473            io.object_path(u64::MAX)
474        );
475        assert_eq!(Some(42), io.parse_object_seq(&io.object_path(42)));
476        assert_eq!(
477            None,
478            io.parse_object_seq("datanodes/1/epochs/2/objects/42.wal")
479        );
480    }
481
482    #[test]
483    fn test_io_rejects_invalid_prefixes() {
484        for prefix in [
485            "",
486            " ",
487            "/absolute",
488            ".",
489            "./node",
490            "node/./epoch",
491            "..",
492            "node/../epoch",
493            "node//epoch",
494            "node//",
495            "node///",
496        ] {
497            match ObjectStoreIo::new(memory_store(), prefix) {
498                Err(Error::InvalidWalObjectStore { .. }) => {}
499                Err(error) => panic!("unexpected error for prefix {prefix:?}: {error:?}"),
500                Ok(_) => panic!("expected prefix {prefix:?} to be rejected"),
501            }
502        }
503    }
504
505    #[tokio::test]
506    async fn test_io_puts_gets_and_lists_objects_by_sequence() {
507        let io = memory_io();
508
509        for object_seq in [10, 2, 1] {
510            assert_eq!(
511                PutResult::Created,
512                io.put_if_absent(object_seq, Bytes::from(object_seq.to_string()))
513                    .await
514                    .unwrap()
515            );
516        }
517
518        assert_eq!(Bytes::from_static(b"2"), io.get(2).await.unwrap());
519        let objects = io.list().await.unwrap();
520        assert_eq!(vec![1, 2, 10], object_seqs(objects.clone()));
521        assert_eq!(io.object_path(1), objects[0].path);
522        assert_eq!(
523            vec![1, 1, 2],
524            objects.iter().map(|object| object.size).collect::<Vec<_>>()
525        );
526
527        // A delete of an object that is already gone succeeds.
528        io.delete(2).await.unwrap();
529        io.delete(2).await.unwrap();
530        assert_eq!(vec![1, 10], object_seqs(io.list().await.unwrap()));
531    }
532
533    #[tokio::test]
534    async fn test_io_get_range_reads_a_slice_inside_the_object() {
535        let io = memory_io();
536        io.put_if_absent(3, Bytes::from_static(b"0123456789"))
537            .await
538            .unwrap();
539
540        assert_eq!(
541            Bytes::from_static(b"234"),
542            io.get_range(3, 2, 3).await.unwrap()
543        );
544        assert_eq!(
545            Bytes::from_static(b"89"),
546            io.get_range(3, 8, 2).await.unwrap()
547        );
548        for (object_seq, offset, len) in [(3, 8, 5), (3, 10, 1), (7, 0, 1)] {
549            let error = io.get_range(object_seq, offset, len).await.unwrap_err();
550            assert!(
551                matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(object_seq)),
552                "unexpected error for range {offset}..{len} of object {object_seq}: {error:?}"
553            );
554        }
555        let error = io.get_range(3, u64::MAX, 1).await.unwrap_err();
556        assert!(
557            matches!(error, Error::CorruptedWalObject { .. }),
558            "unexpected error: {error:?}"
559        );
560    }
561
562    #[tokio::test]
563    async fn test_io_same_content_is_an_idempotent_retry() {
564        let io = memory_io();
565        let content = Bytes::from_static(b"immutable");
566
567        assert_eq!(
568            PutResult::Created,
569            io.put_if_absent(7, content.clone()).await.unwrap()
570        );
571        assert_eq!(
572            PutResult::AlreadyPresent,
573            io.put_if_absent(7, content.clone()).await.unwrap()
574        );
575        assert_eq!(content, io.get(7).await.unwrap());
576    }
577
578    #[tokio::test]
579    async fn test_io_rejects_different_content_without_overwriting() {
580        let io = memory_io();
581        let original = Bytes::from_static(b"original");
582        io.put_if_absent(7, original.clone()).await.unwrap();
583
584        let error = io
585            .put_if_absent(7, Bytes::from_static(b"replacement"))
586            .await
587            .unwrap_err();
588
589        assert!(
590            matches!(error, Error::WalObjectConflict { ref path, .. } if path == &io.object_path(7)),
591            "unexpected error: {error:?}"
592        );
593        assert_eq!(original, io.get(7).await.unwrap());
594    }
595
596    #[tokio::test]
597    async fn test_io_retry_of_an_existing_object_reports_the_failed_read() {
598        let store = memory_store();
599        let prefix = "datanodes/1/epochs/2/";
600        let content = Bytes::from_static(b"immutable");
601        ObjectStoreIo::new(store.clone(), prefix)
602            .unwrap()
603            .put_if_absent(7, content.clone())
604            .await
605            .unwrap();
606
607        // The object exists, so the create is refused, and the read that would
608        // settle the retry fails: that failure and its retry hint come back.
609        let io = ObjectStoreIo::new(with_failing_reads(store.clone()), prefix).unwrap();
610        let error = io.put_if_absent(7, content.clone()).await.unwrap_err();
611        assert!(
612            matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(7)),
613            "unexpected error: {error:?}"
614        );
615        assert_eq!(RetryHint::Retryable, error.retry_hint());
616
617        // A create that fails for another reason is reported as the write
618        // failure it is, although the read of the missing object fails too.
619        let io = ObjectStoreIo::new(with_failing_writes(store), prefix).unwrap();
620        let error = io.put_if_absent(8, content).await.unwrap_err();
621        assert!(
622            matches!(error, Error::WalObjectStore { operation: "write", ref path, .. } if path == &io.object_path(8)),
623            "unexpected error: {error:?}"
624        );
625        assert_eq!(RetryHint::NonRetryable, error.retry_hint());
626    }
627
628    #[tokio::test]
629    async fn test_io_create_under_a_file_parent_is_reported_as_the_write_failure() {
630        let temp_dir = create_temp_dir("object_store_wal_io_file_parent");
631        std::fs::write(temp_dir.path().join("datanodes"), []).unwrap();
632        let store = SecureFsRoot::open(temp_dir.path())
633            .unwrap()
634            .build_operator();
635        let io = ObjectStoreIo::new(store, "datanodes/1/epochs/2/").unwrap();
636
637        // The store reports `AlreadyExists` for the parent, not for the
638        // object, so the create failure itself comes back.
639        let error = io
640            .put_if_absent(7, Bytes::from_static(b"wal"))
641            .await
642            .unwrap_err();
643        assert!(
644            matches!(error, Error::WalObjectStore { operation: "write", ref path, .. } if path == &io.object_path(7)),
645            "unexpected error: {error:?}"
646        );
647    }
648
649    #[tokio::test]
650    async fn test_io_get_of_a_missing_object_fails() {
651        let io = memory_io();
652
653        let error = io.get(7).await.unwrap_err();
654        assert!(
655            matches!(error, Error::WalObjectStore { operation: "read", ref path, .. } if path == &io.object_path(7)),
656            "unexpected error: {error:?}"
657        );
658    }
659
660    #[tokio::test]
661    async fn test_io_list_ahead_yields_the_objects_list_returns() {
662        let io = memory_io();
663        assert!(
664            io.list_ahead()
665                .try_collect::<Vec<_>>()
666                .await
667                .unwrap()
668                .is_empty()
669        );
670
671        io.store
672            .write(
673                "datanodes/1/epochs/2/objects/not-a-sequence.wal",
674                "malformed",
675            )
676            .await
677            .unwrap();
678        for object_seq in [10, 2, 1] {
679            io.put_if_absent(object_seq, Bytes::from(object_seq.to_string()))
680                .await
681                .unwrap();
682        }
683        let mut listed = io.list_ahead().try_collect::<Vec<_>>().await.unwrap();
684        listed.sort_unstable_by_key(|object| object.object_seq);
685        assert_eq!(io.list().await.unwrap(), listed);
686    }
687
688    struct CountingLister {
689        inner: oio::Lister,
690        listed: Arc<AtomicUsize>,
691    }
692
693    impl oio::List for CountingLister {
694        async fn next(&mut self) -> mock::Result<Option<oio::Entry>> {
695            self.listed.fetch_add(1, Ordering::SeqCst);
696            self.inner.next().await
697        }
698    }
699
700    #[tokio::test]
701    async fn test_io_listing_waits_while_the_backlog_is_full() {
702        let store = memory_store();
703        let objects = LISTING_AHEAD_OBJECTS + 100;
704        let io = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/2/").unwrap();
705        for object_seq in 0..objects as u64 {
706            io.put_if_absent(object_seq, Bytes::from_static(b"wal"))
707                .await
708                .unwrap();
709        }
710        let listed = Arc::new(AtomicUsize::new(0));
711        let lister_listed = listed.clone();
712        let mock = MockLayerBuilder::default()
713            .lister_factory(Arc::new(move |_, _, inner| {
714                let listed = lister_listed.clone();
715                Box::new(CountingLister { inner, listed })
716            }))
717            .build()
718            .unwrap();
719        let io = ObjectStoreIo::new(store.layer(mock), "datanodes/1/epochs/2/").unwrap();
720        let (listing, mut rx) = io.listing();
721        // Out of the cooperative budget, so it waits only on a full channel.
722        let mut listing = pin!(tokio::task::unconstrained(listing));
723
724        // Nothing takes the objects, so the listing waits once the backlog is
725        // full, with the directory entry and the object waiting to be sent
726        // listed on top of it.
727        for _ in 0..100 {
728            assert!(poll!(listing.as_mut()).is_pending());
729        }
730        let ahead = listed.load(Ordering::SeqCst);
731        assert!(
732            (LISTING_AHEAD_OBJECTS + 1..=LISTING_AHEAD_OBJECTS + 2).contains(&ahead),
733            "{ahead} entries listed"
734        );
735
736        let taken = async {
737            let mut taken = 0;
738            while rx.recv().await.is_some() {
739                taken += 1;
740            }
741            taken
742        };
743        let (listed, taken) = futures::join!(listing, taken);
744        listed.unwrap();
745        assert_eq!(objects, taken);
746    }
747
748    #[tokio::test]
749    async fn test_io_failed_listing_goes_before_the_objects_it_listed() {
750        let io = memory_io();
751        let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS);
752        for object_seq in 0..1000 {
753            let path = io.object_path(object_seq);
754            tx.try_send(ListedObject {
755                object_seq,
756                path,
757                size: 0,
758            })
759            .unwrap();
760        }
761        let listing = AbortOnDropHandle::new(tokio::spawn(async move {
762            drop(tx);
763            Err(injected_failure("list"))
764        }));
765        while !listing.is_finished() {
766            tokio::task::yield_now().await;
767        }
768
769        let mut objects = settle_listing(rx, listing, io.object_prefix.clone());
770        let error = objects.try_next().await.unwrap_err();
771        assert!(
772            matches!(
773                error,
774                Error::WalObjectStore {
775                    operation: "list",
776                    ..
777                }
778            ),
779            "unexpected error: {error:?}"
780        );
781    }
782
783    #[tokio::test]
784    async fn test_io_list_ignores_objects_outside_the_wal_layout() {
785        let io = memory_io();
786        io.store.write("unrelated", "outside").await.unwrap();
787        io.store
788            .write(
789                "datanodes/1/epochs/2/objects/not-a-sequence.wal",
790                "malformed",
791            )
792            .await
793            .unwrap();
794        io.put_if_absent(3, Bytes::from_static(b"wal"))
795            .await
796            .unwrap();
797
798        assert_eq!(vec![3], object_seqs(io.list().await.unwrap()));
799    }
800
801    #[tokio::test]
802    async fn test_io_list_is_isolated_by_prefix() {
803        let store = memory_store();
804        let first = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/1").unwrap();
805        let second = ObjectStoreIo::new(store, "datanodes/1/epochs/2").unwrap();
806        first
807            .put_if_absent(1, Bytes::from_static(b"first"))
808            .await
809            .unwrap();
810        second
811            .put_if_absent(2, Bytes::from_static(b"second"))
812            .await
813            .unwrap();
814
815        assert_eq!(vec![1], object_seqs(first.list().await.unwrap()));
816        assert_eq!(vec![2], object_seqs(second.list().await.unwrap()));
817    }
818
819    #[tokio::test]
820    async fn test_io_equal_sequences_in_different_prefixes_do_not_conflict() {
821        let store = memory_store();
822        let first = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/1").unwrap();
823        let second = ObjectStoreIo::new(store, "datanodes/1/epochs/2").unwrap();
824
825        assert_eq!(
826            PutResult::Created,
827            first
828                .put_if_absent(7, Bytes::from_static(b"first"))
829                .await
830                .unwrap()
831        );
832        assert_eq!(
833            PutResult::Created,
834            second
835                .put_if_absent(7, Bytes::from_static(b"different"))
836                .await
837                .unwrap()
838        );
839        assert_eq!(Bytes::from_static(b"first"), first.get(7).await.unwrap());
840        assert_eq!(
841            Bytes::from_static(b"different"),
842            second.get(7).await.unwrap()
843        );
844    }
845
846    fn object_seqs(objects: Vec<ListedObject>) -> Vec<u64> {
847        objects
848            .into_iter()
849            .map(|object| object.object_seq)
850            .collect()
851    }
852}
log_store/object_store_wal/io.rs

log_store/object_store_wal/
io.rs
