diff --git a/src/log-store/src/object_store_wal/io.rs b/src/log-store/src/object_store_wal/io.rs index f9ed5eee36b..43f018a7ecb 100644 --- a/src/log-store/src/object_store_wal/io.rs +++ b/src/log-store/src/object_store_wal/io.rs @@ -15,9 +15,14 @@ //! Object store access for WAL objects: a deterministic key layout and //! conditional creates that make a retry of the same object a no-op. +use async_stream::try_stream; use bytes::Bytes; -use object_store::{ErrorKind, ObjectStore}; +use futures::TryStreamExt; +use futures::stream::BoxStream; +use object_store::{Entry, ErrorKind, ObjectStore}; use snafu::{OptionExt, ResultExt, ensure}; +use tokio::sync::mpsc; +use tokio_util::task::AbortOnDropHandle; use crate::error::{ CorruptedWalObjectSnafu, InvalidWalObjectStoreSnafu, Result, WalObjectConflictSnafu, @@ -28,6 +33,10 @@ use crate::error::{ /// [`u64::MAX`] so that lexicographic and numeric order agree. const OBJECT_SEQ_WIDTH: usize = 20; const OBJECT_SUFFIX: &str = ".wal"; +/// Objects a listing may hold ahead of the reads that take them, four LIST +/// pages of S3, so the listing stays pages ahead of the reads while what it +/// holds is bounded. +const LISTING_AHEAD_OBJECTS: usize = 4096; /// Outcome of a conditional create. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -48,6 +57,7 @@ pub(crate) struct ListedObject { } /// Reads and writes the WAL objects under one prefix. +#[derive(Clone)] pub(crate) struct ObjectStoreIo { store: ObjectStore, object_prefix: String, @@ -160,20 +170,50 @@ impl ObjectStoreIo { path: self.object_prefix.clone(), })?; let mut objects = entries - .into_iter() - .filter_map(|entry| { - self.parse_object_seq(entry.path()) - .map(|object_seq| ListedObject { - object_seq, - path: entry.path().to_string(), - size: entry.metadata().content_length(), - }) - }) + .iter() + .filter_map(|entry| self.listed_object(entry)) .collect::>(); objects.sort_unstable_by_key(|object| object.object_seq); Ok(objects) } + /// Lists the objects under the prefix like [`list`](Self::list), but from + /// a task of its own and in the order the object store returns them, so + /// the listing goes on while the caller consumes the objects listed so + /// far. A listing that fails, however it fails, ends the stream with an + /// error, before any object it listed that the caller has not taken yet; + /// dropping the stream stops the listing. + pub(crate) fn list_ahead(&self) -> BoxStream<'static, Result> { + let (listing, rx) = self.clone().listing(); + let listing = AbortOnDropHandle::new(common_runtime::spawn_global(listing)); + settle_listing(rx, listing, self.object_prefix.clone()) + } + + /// Returns the listing of the objects under the prefix and the receiver it + /// sends them to; it waits while the receiver holds + /// [`LISTING_AHEAD_OBJECTS`] of them. + fn listing( + self, + ) -> ( + impl Future> + Send + 'static, + mpsc::Receiver, + ) { + let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS); + let listing = async move { + let mut lister = self.store.lister(&self.object_prefix).await?; + while let Some(entry) = lister.try_next().await? { + // The receiver is gone only once the stream was dropped. + if let Some(object) = self.listed_object(&entry) + && tx.send(object).await.is_err() + { + break; + } + } + Ok::<_, object_store::Error>(()) + }; + (listing, rx) + } + pub(crate) fn object_path(&self, object_seq: u64) -> String { format!( "{}{object_seq:0OBJECT_SEQ_WIDTH$}{OBJECT_SUFFIX}", @@ -193,6 +233,15 @@ impl ObjectStoreIo { } } + fn listed_object(&self, entry: &Entry) -> Option { + self.parse_object_seq(entry.path()) + .map(|object_seq| ListedObject { + object_seq, + path: entry.path().to_string(), + size: entry.metadata().content_length(), + }) + } + fn parse_object_seq(&self, path: &str) -> Option { let value = path .strip_prefix(&self.object_prefix)? @@ -204,6 +253,38 @@ impl ObjectStoreIo { } } +/// Yields the objects `listing` sends over `rx`. Once the listing finished it +/// is settled before any further object is yielded, so its failure goes +/// before the objects it listed that the caller has not taken yet. +fn settle_listing( + mut rx: mpsc::Receiver, + listing: AbortOnDropHandle>, + path: String, +) -> BoxStream<'static, Result> { + let mut listing = Some(listing); + Box::pin(try_stream! { + loop { + let object = rx.recv().await; + if let Some(listing) = + listing.take_if(|listing| object.is_none() || listing.is_finished()) + { + listing + .await + .unwrap_or_else(|error| { + let stopped = "the listing task did not complete"; + Err(object_store::Error::new(ErrorKind::Unexpected, stopped) + .set_source(error)) + }) + .context(WalObjectStoreSnafu { operation: "list", path: &path })?; + } + let Some(object) = object else { + break; + }; + yield object; + } + }) +} + /// Returns true when a conditional create failed because the object exists. /// Only the precondition failure says so: a store may report `AlreadyExists` /// for an unrelated path, such as a parent that is a file. @@ -233,10 +314,13 @@ fn normalize_prefix(prefix: &str) -> Result { #[cfg(test)] mod tests { + use std::pin::pin; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use common_error::ext::{ErrorExt, RetryHint}; use common_test_util::temp_dir::create_temp_dir; + use futures::poll; use object_store::layers::mock::{self, MockLayerBuilder, oio}; use object_store::secure_fs::SecureFsRoot; use object_store::services::Memory; @@ -513,6 +597,129 @@ mod tests { ); } + #[tokio::test] + async fn test_io_list_ahead_yields_the_objects_list_returns() { + let io = memory_io(); + assert!( + io.list_ahead() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + + io.store + .write( + "datanodes/1/epochs/2/objects/not-a-sequence.wal", + "malformed", + ) + .await + .unwrap(); + for object_seq in [10, 2, 1] { + io.put_if_absent(object_seq, Bytes::from(object_seq.to_string())) + .await + .unwrap(); + } + let mut listed = io.list_ahead().try_collect::>().await.unwrap(); + listed.sort_unstable_by_key(|object| object.object_seq); + assert_eq!(io.list().await.unwrap(), listed); + } + + struct CountingLister { + inner: oio::Lister, + listed: Arc, + } + + impl oio::List for CountingLister { + async fn next(&mut self) -> mock::Result> { + self.listed.fetch_add(1, Ordering::SeqCst); + self.inner.next().await + } + } + + #[tokio::test] + async fn test_io_listing_waits_while_the_backlog_is_full() { + let store = memory_store(); + let objects = LISTING_AHEAD_OBJECTS + 100; + let io = ObjectStoreIo::new(store.clone(), "datanodes/1/epochs/2/").unwrap(); + for object_seq in 0..objects as u64 { + io.put_if_absent(object_seq, Bytes::from_static(b"wal")) + .await + .unwrap(); + } + let listed = Arc::new(AtomicUsize::new(0)); + let lister_listed = listed.clone(); + let mock = MockLayerBuilder::default() + .lister_factory(Arc::new(move |_, _, inner| { + let listed = lister_listed.clone(); + Box::new(CountingLister { inner, listed }) + })) + .build() + .unwrap(); + let io = ObjectStoreIo::new(store.layer(mock), "datanodes/1/epochs/2/").unwrap(); + let (listing, mut rx) = io.listing(); + // Out of the cooperative budget, so it waits only on a full channel. + let mut listing = pin!(tokio::task::unconstrained(listing)); + + // Nothing takes the objects, so the listing waits once the backlog is + // full, with the directory entry and the object waiting to be sent + // listed on top of it. + for _ in 0..100 { + assert!(poll!(listing.as_mut()).is_pending()); + } + let ahead = listed.load(Ordering::SeqCst); + assert!( + (LISTING_AHEAD_OBJECTS + 1..=LISTING_AHEAD_OBJECTS + 2).contains(&ahead), + "{ahead} entries listed" + ); + + let taken = async { + let mut taken = 0; + while rx.recv().await.is_some() { + taken += 1; + } + taken + }; + let (listed, taken) = futures::join!(listing, taken); + listed.unwrap(); + assert_eq!(objects, taken); + } + + #[tokio::test] + async fn test_io_failed_listing_goes_before_the_objects_it_listed() { + let io = memory_io(); + let (tx, rx) = mpsc::channel(LISTING_AHEAD_OBJECTS); + for object_seq in 0..1000 { + let path = io.object_path(object_seq); + tx.try_send(ListedObject { + object_seq, + path, + size: 0, + }) + .unwrap(); + } + let listing = AbortOnDropHandle::new(tokio::spawn(async move { + drop(tx); + Err(injected_failure("list")) + })); + while !listing.is_finished() { + tokio::task::yield_now().await; + } + + let mut objects = settle_listing(rx, listing, io.object_prefix.clone()); + let error = objects.try_next().await.unwrap_err(); + assert!( + matches!( + error, + Error::WalObjectStore { + operation: "list", + .. + } + ), + "unexpected error: {error:?}" + ); + } + #[tokio::test] async fn test_io_list_ignores_objects_outside_the_wal_layout() { let io = memory_io(); diff --git a/src/log-store/src/object_store_wal/store.rs b/src/log-store/src/object_store_wal/store.rs index d5096c073eb..40d6e75326c 100644 --- a/src/log-store/src/object_store_wal/store.rs +++ b/src/log-store/src/object_store_wal/store.rs @@ -27,8 +27,8 @@ use bytes::Bytes; use common_telemetry::{debug, error, info, warn}; use common_wal::config::object_store::{AckMode, CorruptedSegmentAction, ObjectStoreWalConfig}; use futures::future::{BoxFuture, OptionFuture}; -use futures::stream::FuturesUnordered; -use futures::{StreamExt, TryStreamExt}; +use futures::stream::{BoxStream, FuturesUnordered}; +use futures::{Stream, StreamExt, TryStreamExt}; use object_store::{ErrorKind, ObjectStore}; use snafu::{IntoError, OptionExt, ResultExt, ensure}; use store_api::logstore::entry::{Entry, NaiveEntry}; @@ -3069,12 +3069,12 @@ struct FetchedObject { /// are. Segments are not read, except those of release records, see /// [`read_releases`]; a segment checksum is verified by the read that decodes /// it. Footers are fetched for up to [`RECOVERY_CONCURRENCY`] objects -/// at a time and indexed in sequence order, so the catalog checks the entry -/// ranges of every object against its predecessors like a sequential replay. -/// Only the objects on the chain [`select_chain`] picks are indexed. +/// at a time while the listing goes on, and indexed in sequence order, so the +/// catalog checks the entry ranges of every object against its predecessors +/// like a sequential replay. Only the objects on the chain [`select_chain`] +/// picks are indexed. async fn recover(io: &dyn WalObjectIo) -> Result { - let objects = io.list().await?; - let mut fetched = fetch_footers(io, objects, RECOVERY_CONCURRENCY).await?; + let mut fetched = fetch_footers(io, io.list_ahead(), RECOVERY_CONCURRENCY).await?; let releases = fetched .objects .iter_mut() @@ -3372,18 +3372,18 @@ async fn stale_conflict( .fail() } -/// Fetches and verifies the headers and footers of `objects`, up to -/// `concurrency` objects at a time, and returns them ordered by object sequence -/// whatever the order the fetches complete in. An object the object store no -/// longer has is returned as vanished; any other failure abandons the -/// remaining fetches. +/// Fetches and verifies the headers and footers of `objects` as they are +/// listed, up to `concurrency` objects at a time, and returns them ordered by +/// object sequence whatever the order the fetches complete in. An object the +/// object store no longer has is returned as vanished; a failed listing or +/// any other failure abandons the remaining fetches. async fn fetch_footers( io: &dyn WalObjectIo, - objects: Vec, + objects: impl Stream>, concurrency: usize, ) -> Result { - let outcomes = futures::stream::iter(objects) - .map(|object| async move { + let outcomes = objects + .map_ok(|object| async move { match fetch_footer(io, &object).await { Ok((header, footer, releases)) => Ok(Ok(FetchedObject { object, @@ -3395,7 +3395,7 @@ async fn fetch_footers( Err(error) => Err(error), } }) - .buffer_unordered(concurrency) + .try_buffer_unordered(concurrency) .try_collect::>() .await?; let mut objects = Vec::with_capacity(outcomes.len()); @@ -3549,6 +3549,16 @@ pub(crate) trait WalObjectIo: Send + Sync { async fn list(&self) -> Result>; + /// See [`ObjectStoreIo::list_ahead`]; by default the objects + /// [`list`](Self::list) returns. + fn list_ahead(&self) -> BoxStream<'_, Result> { + Box::pin(try_stream! { + for object in self.list().await? { + yield object; + } + }) + } + fn object_path(&self, object_seq: u64) -> String; /// Describes the object store the objects live in, for messages. @@ -3577,6 +3587,10 @@ impl WalObjectIo for ObjectStoreIo { ObjectStoreIo::list(self).await } + fn list_ahead(&self) -> BoxStream<'_, Result> { + ObjectStoreIo::list_ahead(self) + } + fn object_path(&self, object_seq: u64) -> String { ObjectStoreIo::object_path(self, object_seq) } @@ -5174,7 +5188,7 @@ mod tests { let mut fetch = { let io = io.clone(); tokio::spawn(async move { - fetch_footers(io.as_ref(), objects, RECOVERY_CONCURRENCY) + fetch_footers(io.as_ref(), listing(objects), RECOVERY_CONCURRENCY) .await .unwrap() .objects @@ -5261,6 +5275,170 @@ mod tests { store.stop().await.unwrap(); } + fn listing(objects: Vec) -> impl Stream> { + futures::stream::iter(objects.into_iter().map(Ok)) + } + + const LISTING_PAGE: usize = 4; + + /// Serves the entries of a listing in pages of [`LISTING_PAGE`], each + /// only once `pages` allows it, fails, or panics with `panics`, after + /// `fail_after` entries, and records when it is dropped. + struct PagedLister { + inner: oio::Lister, + served: usize, + pages: watch::Receiver, + fail_after: Option, + panics: bool, + dropped: Arc, + } + + impl oio::List for PagedLister { + async fn next(&mut self) -> mock::Result> { + if self.fail_after == Some(self.served) { + assert!(!self.panics, "injected listing panic"); + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected listing failure", + )); + } + let page = self.served / LISTING_PAGE; + if self.pages.wait_for(|pages| *pages > page).await.is_err() { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "listing abandoned", + )); + } + self.served += 1; + self.inner.next().await + } + } + + impl Drop for PagedLister { + fn drop(&mut self) { + self.dropped.store(true, Ordering::SeqCst); + } + } + + /// The prefix of `object_store` listed by [`PagedLister`]s, whose reads + /// go through `read`. + fn paged_listing( + object_store: &ObjectStore, + pages: watch::Receiver, + fail_after: Option, + panics: bool, + read: mock::MockReaderFactory, + ) -> (ObjectStoreIo, Arc) { + let dropped = Arc::new(AtomicBool::new(false)); + let lister_dropped = dropped.clone(); + let mock = MockLayerBuilder::default() + .lister_factory(Arc::new(move |_, _, inner| { + Box::new(PagedLister { + inner, + served: 0, + pages: pages.clone(), + fail_after, + panics, + dropped: lister_dropped.clone(), + }) + })) + .reader_factory(read) + .build() + .unwrap(); + let io = ObjectStoreIo::new(object_store.clone().layer(mock), PREFIX).unwrap(); + (io, dropped) + } + + #[tokio::test] + async fn test_store_recovery_reads_footers_while_the_listing_goes_on() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE + 1, 3).await; + // Every read allows one more page, so recovery completes only if it + // reads footers before the listing ends. + let (pages, listed_pages) = watch::channel(1); + let (io, _) = paged_listing( + &object_store, + listed_pages, + None, + false, + Arc::new(move |_, _, inner| { + pages.send_modify(|pages| *pages += 1); + inner + }), + ); + + let recovered = timeout(WAIT, recover(&io)).await.unwrap().unwrap(); + let expected = recover_by_decoding(&ObjectStoreIo::new(object_store, PREFIX).unwrap()) + .await + .unwrap(); + assert_same_recovery(&expected, &recovered); + assert_eq!(3 * LISTING_PAGE as u64 + 1, recovered.recovered_objects); + } + + #[tokio::test] + async fn test_store_recovery_fails_on_a_listing_that_fails_part_way() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE, 3).await; + for (panics, reason) in [ + (false, "injected listing failure"), + (true, "the listing task did not complete"), + ] { + let (_pages, listed_pages) = watch::channel(usize::MAX); + let (io, _) = paged_listing( + &object_store, + listed_pages, + Some(LISTING_PAGE + 1), + panics, + Arc::new(|_, _, inner| inner), + ); + + let error = timeout(WAIT, recover(&io)).await.unwrap().unwrap_err(); + assert!( + matches!(&error, Error::WalObjectStore { operation: "list", error, .. } + if error.to_string().contains(reason)), + "unexpected error: {error:?}" + ); + } + } + + #[tokio::test] + async fn test_store_recovery_stops_the_listing_when_a_footer_read_fails() { + let _serialized = SERIALIZED.lock().await; + let object_store = memory_store(); + populate(&object_store, 3 * LISTING_PAGE, 3).await; + // Only the first page is ever listed, so the listing waits for good + // unless the failed read stops it. + let (_pages, listed_pages) = watch::channel(1); + let (io, dropped) = paged_listing( + &object_store, + listed_pages, + None, + false, + Arc::new(|_, _, _| Box::new(RefusedReader(ErrorKind::PermissionDenied))), + ); + + let error = timeout(WAIT, recover(&io)).await.unwrap().unwrap_err(); + assert!( + matches!( + &error, + Error::WalObjectStore { + operation: "read", + .. + } + ), + "unexpected error: {error:?}" + ); + timeout(WAIT, async { + while !dropped.load(Ordering::SeqCst) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + } + #[tokio::test] async fn test_store_recovery_reads_only_header_and_tail_of_large_object() { let _serialized = SERIALIZED.lock().await; @@ -8880,7 +9058,7 @@ mod tests { } async fn recover_listed(io: &ObjectStoreIo, listed: Vec) -> Result { - finish_recovery(fetch_footers(io, listed, RECOVERY_CONCURRENCY).await?) + finish_recovery(fetch_footers(io, listing(listed), RECOVERY_CONCURRENCY).await?) } #[tokio::test] @@ -10125,6 +10303,10 @@ mod tests { self.inner.list().await } + fn list_ahead(&self) -> BoxStream<'_, Result> { + self.inner.list_ahead() + } + fn object_path(&self, object_seq: u64) -> String { self.inner.object_path(object_seq) }