use crate::error::classify;
use crate::lane_ctx::{LaneLabel, PoisonKind};
use crate::split::SplitRange;
use crate::trim::RangeTrim;
use crate::unit::{Origin, quoted};
use bytes::Bytes;
use futures_util::StreamExt as _;
use futures_util::future::{Either, select};
use object_store::path::Path;
use object_store::{GetOptions, GetRange, ObjectStore, ObjectStoreExt as _};
use spate_core::error::{ErrorClass, SourceError};
use std::future::Future;
use std::pin::{Pin, pin};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::mpsc;
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct ObjectEntry {
pub(crate) key: String,
pub(crate) size: u64,
pub(crate) etag: Option<String>,
pub(crate) last_modified_ms: i64,
}
#[derive(Debug)]
pub(crate) enum ChunkMsg {
ObjectStart {
ordinal: u32,
key: String,
last_modified_ms: i64,
},
Chunk(Bytes),
ObjectEnd,
LaneFailed(SplitFailure),
}
#[derive(Debug)]
pub(crate) enum SplitFailure {
Poison(PoisonKind, String),
Fatal(SourceError),
Overwritten(Box<Overwrite>),
}
impl std::fmt::Display for SplitFailure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SplitFailure::Poison(_, reason) => write!(f, "split poisoned: {reason}"),
SplitFailure::Fatal(e) => write!(f, "{e}"),
SplitFailure::Overwritten(o) => write!(f, "\"{}\" was overwritten", o.key),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct Overwrite {
pub(crate) key: String,
pub(crate) etag: Option<String>,
pub(crate) size: u64,
pub(crate) last_modified_ms: i64,
}
#[cfg(test)]
pub(crate) async fn list_all(
store: &Arc<dyn ObjectStore>,
prefix: Option<&Path>,
) -> Result<Vec<ObjectEntry>, object_store::Error> {
let mut entries = Vec::new();
let mut stream = store.list(prefix);
while let Some(meta) = stream.next().await {
let meta = meta?;
entries.push(ObjectEntry {
key: meta.location.to_string(),
size: meta.size,
etag: meta.e_tag,
last_modified_ms: meta.last_modified.timestamp_millis(),
});
}
entries.sort_unstable_by(|a, b| a.key.cmp(&b.key));
Ok(entries)
}
pub(crate) struct FetcherParams {
pub(crate) split: LaneLabel,
pub(crate) store: Arc<dyn ObjectStore>,
pub(crate) slice: Arc<Vec<ObjectEntry>>,
pub(crate) start_ordinal: u32,
pub(crate) resume_etag: Option<String>,
pub(crate) chunk_bytes: usize,
pub(crate) range_bytes: usize,
pub(crate) tx: mpsc::Sender<ChunkMsg>,
pub(crate) pause: Arc<AtomicBool>,
pub(crate) retry_base: Duration,
pub(crate) retries: Option<spate_core::metrics::Counter>,
pub(crate) range: Option<SplitRange>,
pub(crate) bytes_read: Option<spate_core::metrics::Counter>,
pub(crate) origin: Origin,
}
pub(crate) const MAX_ATTEMPTS: u32 = 8;
pub(crate) const BACKOFF_CAP: Duration = Duration::from_secs(5);
const PAUSE_POLL: Duration = Duration::from_millis(25);
pub(crate) async fn run_fetcher(params: FetcherParams) {
let FetcherParams {
split,
store,
slice,
start_ordinal,
resume_etag,
chunk_bytes,
range_bytes,
tx,
pause,
retry_base,
retries,
range,
bytes_read,
origin,
} = params;
let mut ahead = ReadAhead::default();
for (ordinal, entry) in slice.iter().enumerate().skip(start_ordinal as usize) {
let ordinal = ordinal as u32;
let pinned_etag = if ordinal == start_ordinal && resume_etag.is_some() {
resume_etag.clone()
} else {
entry.etag.clone()
};
let start = ChunkMsg::ObjectStart {
ordinal,
key: entry.key.clone(),
last_modified_ms: entry.last_modified_ms,
};
if !ahead.alongside(gated_send(&pause, &tx, start)).await {
return; }
let next = match range {
Some(_) => None,
None => slice
.iter()
.skip(ordinal as usize + 1)
.find(|later| later.size > 0)
.and_then(|later| {
let etag = later.etag.clone()?;
let path = object_path(&split, later).ok()?;
Some(WindowRead {
split: Arc::clone(&split),
store: Arc::clone(&store),
path,
entry: later.clone(),
etag,
start: 0,
end: (range_bytes.max(1) as u64).min(later.size),
retry_base,
retries: retries.clone(),
origin: origin.clone(),
})
}),
};
match stream_object(
&split,
&store,
entry,
pinned_etag.as_deref(),
range,
range_bytes,
chunk_bytes,
&tx,
&pause,
retry_base,
retries.as_ref(),
bytes_read.as_ref(),
&origin,
&mut ahead,
next,
)
.await
{
Ok(true) => {}
Ok(false) => return, Err(e) => {
let _ = tx.send(ChunkMsg::LaneFailed(e)).await;
return;
}
}
if ahead.alongside(tx.send(ChunkMsg::ObjectEnd)).await.is_err() {
return;
}
}
}
async fn pause_gate(pause: &AtomicBool, tx: &mpsc::Sender<ChunkMsg>) -> Result<(), ()> {
while pause.load(Ordering::Relaxed) {
if tx.is_closed() {
return Err(());
}
tokio::time::sleep(PAUSE_POLL).await;
}
Ok(())
}
async fn gated_send(pause: &AtomicBool, tx: &mpsc::Sender<ChunkMsg>, msg: ChunkMsg) -> bool {
pause_gate(pause, tx).await.is_ok() && tx.send(msg).await.is_ok()
}
type Window = Result<Bytes, SplitFailure>;
struct WindowRead {
split: LaneLabel,
store: Arc<dyn ObjectStore>,
path: Path,
entry: ObjectEntry,
etag: String,
start: u64,
end: u64,
retry_base: Duration,
retries: Option<spate_core::metrics::Counter>,
origin: Origin,
}
impl WindowRead {
async fn fetch(self) -> Window {
let mut attempt: u32 = 0;
loop {
let options = GetOptions {
if_match: Some(self.etag.clone()),
range: Some(GetRange::Bounded(self.start..self.end)),
version: version_of(&self.origin),
..Default::default()
};
let fetched = match self.store.get_opts(&self.path, options).await {
Ok(result) => result.bytes().await,
Err(e) => Err(e),
};
match fetched {
Ok(bytes) if bytes.is_empty() => {
return Err(SplitFailure::Poison(
PoisonKind::Undecodable,
format!(
"split {}: ranged read of \"{}\" at byte {} returned no bytes",
self.split, self.entry.key, self.start
),
));
}
Ok(bytes) => return Ok(bytes),
Err(e) if crate::error::is_overwrite(&e, &self.origin) => {
return Err(head_after_overwrite(
&self.split,
&self.store,
&self.path,
&self.entry,
&self.etag,
self.start,
self.retry_base,
self.retries.as_ref(),
)
.await);
}
Err(e) => {
retry_or_fail(
&self.split,
&self.entry,
&e,
self.start,
&mut attempt,
self.retry_base,
self.retries.as_ref(),
)
.await?;
}
}
}
}
}
#[derive(Default)]
struct ReadAhead(Option<Ahead>);
enum Ahead {
Reading(Pin<Box<dyn Future<Output = Window> + Send>>),
Read(Window),
}
impl ReadAhead {
fn start(&mut self, read: WindowRead) {
debug_assert!(self.0.is_none(), "a second window read ahead");
self.0 = Some(Ahead::Reading(Box::pin(read.fetch())));
}
async fn alongside<T>(&mut self, work: impl Future<Output = T>) -> T {
let work = pin!(work);
let Some(Ahead::Reading(reading)) = &mut self.0 else {
return work.await;
};
match select(work, reading.as_mut()).await {
Either::Left((out, _)) => out,
Either::Right((window, work)) => {
self.0 = Some(Ahead::Read(window));
work.await
}
}
}
async fn take(&mut self) -> Option<Window> {
Some(match self.0.take()? {
Ahead::Reading(reading) => reading.await,
Ahead::Read(window) => window,
})
}
}
#[expect(
clippy::too_many_arguments,
reason = "internal seam between run_fetcher and the read loops"
)]
async fn stream_object(
split: &LaneLabel,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
pinned_etag: Option<&str>,
range: Option<SplitRange>,
range_bytes: usize,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
bytes_read: Option<&spate_core::metrics::Counter>,
origin: &Origin,
ahead: &mut ReadAhead,
next: Option<WindowRead>,
) -> Result<bool, SplitFailure> {
if entry.size == 0 && range.is_none() {
return Ok(true);
}
match pinned_etag {
Some(etag) => {
stream_object_ranged(
split,
store,
entry,
etag,
range.map(|range| RangeTrim::new(range, entry.size)),
range_bytes,
chunk_bytes,
tx,
pause,
retry_base,
retries,
bytes_read,
origin,
ahead,
next,
)
.await
}
None if range.is_some() => Err(SplitFailure::Fatal(SourceError::Client {
class: ErrorClass::Fatal,
reason: format!(
"split {split}: ranged read of \"{}\" has no ETag pin — internal wiring bug",
entry.key
),
})),
None => {
stream_object_streaming(
split,
store,
entry,
chunk_bytes,
tx,
pause,
retry_base,
retries,
origin,
)
.await
}
}
}
#[expect(
clippy::too_many_arguments,
reason = "internal read loop, mirrors stream_object_streaming"
)]
async fn stream_object_ranged(
split: &LaneLabel,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
etag: &str,
mut trim: Option<RangeTrim>,
range_bytes: usize,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
bytes_read: Option<&spate_core::metrics::Counter>,
origin: &Origin,
ahead: &mut ReadAhead,
mut next: Option<WindowRead>,
) -> Result<bool, SplitFailure> {
let path = object_path(split, entry)?;
let window = range_bytes.max(1) as u64;
let window_end = |trim: &Option<RangeTrim>, at: u64| match trim {
Some(trim) => trim.next_window(at, window, chunk_bytes as u64),
None if at < entry.size => Some(at.saturating_add(window).min(entry.size)),
None => None,
};
let read_at = |start: u64, end: u64| WindowRead {
split: Arc::clone(split),
store: Arc::clone(store),
path: path.clone(),
entry: entry.clone(),
etag: etag.to_owned(),
start,
end,
retry_base,
retries: retries.cloned(),
origin: origin.clone(),
};
let mut delivered: u64 = trim.as_ref().map_or(0, RangeTrim::first_byte);
let mut buffered = match ahead.take().await {
Some(window) => window?,
None => match window_end(&trim, delivered) {
Some(end) => read_at(delivered, end).fetch().await?,
None => return Ok(true),
},
};
loop {
let read = buffered.len() as u64;
let mut bytes = match &mut trim {
Some(trim) => {
let owned = trim.feed(delivered, buffered);
if let Some(counter) = bytes_read {
counter.increment(read - owned.len() as u64);
}
owned
}
None => buffered,
};
let following = delivered + read;
let own_next = window_end(&trim, following);
if pause_gate(pause, tx).await.is_err() {
return Ok(false);
}
match own_next {
Some(end) => ahead.start(read_at(following, end)),
None => {
if let Some(read) = next.take() {
ahead.start(read);
}
}
}
while !bytes.is_empty() {
let take = bytes.len().min(chunk_bytes);
let chunk = bytes.split_to(take);
if !ahead
.alongside(gated_send(pause, tx, ChunkMsg::Chunk(chunk)))
.await
{
return Ok(false);
}
}
if own_next.is_none() {
return Ok(true);
}
delivered = following;
let Some(window) = ahead.take().await else {
unreachable!("the object's next window was started above");
};
buffered = window?;
}
}
#[expect(
clippy::too_many_arguments,
reason = "internal read loop, mirrors stream_object_ranged"
)]
async fn stream_object_streaming(
split: &LaneLabel,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
origin: &Origin,
) -> Result<bool, SplitFailure> {
let path = object_path(split, entry)?;
let mut delivered: u64 = 0;
let mut attempt: u32 = 0;
'attempts: loop {
let options = GetOptions {
version: version_of(origin),
..Default::default()
};
let result = match store.get_opts(&path, options).await {
Ok(r) => r,
Err(e) => {
retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await?;
continue 'attempts;
}
};
let mut stream = result.into_stream();
loop {
match stream.next().await {
Some(Ok(mut bytes)) => {
delivered += bytes.len() as u64;
while !bytes.is_empty() {
let take = bytes.len().min(chunk_bytes);
let chunk = bytes.split_to(take);
if pause_gate(pause, tx).await.is_err() {
return Ok(false);
}
if tx.send(ChunkMsg::Chunk(chunk)).await.is_err() {
return Ok(false);
}
}
attempt = 0;
}
Some(Err(e)) => {
if delivered > 0 {
if crate::error::is_pipeline_fatal(&e) {
return Err(SplitFailure::Fatal(SourceError::Client {
class: ErrorClass::Fatal,
reason: format!(
"split {split}: reading \"{}\" failed mid-object at byte \
{delivered}: {}",
entry.key,
crate::error::bounded_reason(&e)
),
}));
}
return Err(SplitFailure::Poison(
PoisonKind::Undecodable,
format!(
"split {split}: read of \"{}\" failed mid-object at byte \
{delivered} and the store reports no ETag to pin a \
resumed read to: {}",
entry.key,
crate::error::bounded_reason(&e)
),
));
}
retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await?;
continue 'attempts;
}
None => return Ok(true),
}
}
}
}
fn version_of(origin: &Origin) -> Option<String> {
match origin {
Origin::Notification { version, .. } => version.clone(),
Origin::Lister => None,
}
}
#[expect(
clippy::too_many_arguments,
reason = "internal seam beside the read loops, which carry the same context"
)]
async fn head_after_overwrite(
split: &LaneLabel,
store: &Arc<dyn ObjectStore>,
path: &Path,
entry: &ObjectEntry,
pinned: &str,
delivered: u64,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
) -> SplitFailure {
let mut attempt = 0;
loop {
match store.head(path).await {
Ok(meta) => {
let etag = meta.e_tag.map(quoted);
if etag.as_deref() == Some(pinned) {
return SplitFailure::Poison(
PoisonKind::EtagDrift,
format!(
"split {split}: the store refused reading \"{}\" at byte \
{delivered} under its ETag pin, and the object still reports the \
ETag {pinned}",
entry.key
),
);
}
return SplitFailure::Overwritten(Box::new(Overwrite {
key: entry.key.clone(),
etag,
size: meta.size,
last_modified_ms: meta.last_modified.timestamp_millis(),
}));
}
Err(e) => {
if let Err(failure) = retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await
{
let headed = |reason: String| {
format!("{reason} (the HEAD after the store refused the pinned read)")
};
return match failure {
SplitFailure::Poison(kind, reason) => {
SplitFailure::Poison(kind, headed(reason))
}
SplitFailure::Fatal(SourceError::Client { class, reason }) => {
SplitFailure::Fatal(SourceError::Client {
class,
reason: headed(reason),
})
}
other => other,
};
}
}
}
}
}
fn object_path(split: &LaneLabel, entry: &ObjectEntry) -> Result<Path, SplitFailure> {
match Path::parse(&entry.key) {
Ok(path) if path.as_ref() == entry.key => Ok(path),
parsed => {
let why = parsed.map_or_else(
|e| e.to_string(),
|path| format!("the object store would read \"{path}\" in its place"),
);
Err(SplitFailure::Poison(
PoisonKind::NotFound,
format!(
"split {split}: \"{}\" is not an object key: {why}",
entry.key
),
))
}
}
}
async fn retry_or_fail(
split: &LaneLabel,
entry: &ObjectEntry,
e: &object_store::Error,
delivered: u64,
attempt: &mut u32,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
) -> Result<(), SplitFailure> {
if classify(e) != ErrorClass::Retryable {
let reason = format!(
"split {split}: reading \"{}\" failed at byte {delivered}: {}",
entry.key,
crate::error::bounded_reason(e)
);
return Err(if crate::error::is_pipeline_fatal(e) {
SplitFailure::Fatal(SourceError::Client {
class: ErrorClass::Fatal,
reason,
})
} else {
SplitFailure::Poison(crate::error::poison_kind(e), reason)
});
}
*attempt += 1;
if *attempt >= MAX_ATTEMPTS {
return Err(SplitFailure::Poison(
PoisonKind::RetriesExhausted,
format!(
"split {split}: reading \"{}\" still failing at byte {delivered} after \
{MAX_ATTEMPTS} attempts: {}",
entry.key,
crate::error::bounded_reason(e)
),
));
}
let backoff = retry_base
.saturating_mul(1 << (*attempt - 1).min(16))
.min(BACKOFF_CAP);
tracing::warn!(
split = %split,
key = %entry.key,
attempt = *attempt,
delivered,
error = %crate::error::bounded_reason(e),
"transient object read failure; backing off and resuming with a ranged get"
);
tokio::time::sleep(backoff).await;
if let Some(c) = retries {
c.increment(1);
}
Ok(())
}
#[expect(
clippy::disallowed_methods,
reason = "a fetcher task ends when the test aborts it or the test's runtime drops"
)]
#[cfg(test)]
mod tests {
use super::*;
use object_store::memory::InMemory;
use object_store::{GetResult, GetResultPayload, PutPayload};
use std::fmt;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::AtomicU64;
use std::sync::mpsc as std_mpsc;
fn entry(key: &str, size: u64, etag: Option<&str>) -> ObjectEntry {
ObjectEntry {
key: key.into(),
size,
etag: etag.map(str::to_owned),
last_modified_ms: 0,
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RangeKind {
Full,
Offset(u64),
Bounded(u64, u64),
}
fn range_kind(range: &Option<GetRange>) -> RangeKind {
match range {
None => RangeKind::Full,
Some(GetRange::Offset(o)) => RangeKind::Offset(*o),
Some(GetRange::Bounded(r)) => RangeKind::Bounded(r.start, r.end),
Some(GetRange::Suffix(_)) => unreachable!("the fetcher never issues a suffix range"),
}
}
#[derive(Debug)]
struct FlakyStore {
inner: InMemory,
fail_gets: AtomicU32,
fail_past: u64,
fail_past_gets: AtomicU32,
cut_streams: AtomicU32,
cut_after: usize,
cut_fatal: bool,
error_text: Option<String>,
gets: AtomicU32,
ranges: std::sync::Mutex<Vec<RangeKind>>,
bodies_drained: Arc<AtomicU32>,
requests: std::sync::Mutex<Vec<(String, Option<String>)>>,
hold_on: u32,
release: Arc<tokio::sync::Semaphore>,
empty_on: u32,
held: AtomicU32,
held_dropped: Arc<AtomicU32>,
count_live: bool,
live: Arc<AtomicU32>,
max_live: Arc<AtomicU32>,
}
struct DropCount(Arc<AtomicU32>);
impl Drop for DropCount {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
struct LiveWindow {
bytes: Bytes,
live: Arc<AtomicU32>,
}
impl AsRef<[u8]> for LiveWindow {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}
impl Drop for LiveWindow {
fn drop(&mut self) {
self.live.fetch_sub(1, Ordering::Relaxed);
}
}
impl FlakyStore {
fn new(inner: InMemory) -> FlakyStore {
FlakyStore {
inner,
fail_gets: AtomicU32::new(0),
fail_past: u64::MAX,
fail_past_gets: AtomicU32::new(0),
cut_streams: AtomicU32::new(0),
cut_after: 0,
cut_fatal: false,
error_text: None,
gets: AtomicU32::new(0),
ranges: std::sync::Mutex::new(Vec::new()),
bodies_drained: Arc::new(AtomicU32::new(0)),
requests: std::sync::Mutex::new(Vec::new()),
hold_on: 0,
empty_on: 0,
release: Arc::new(tokio::sync::Semaphore::new(0)),
held: AtomicU32::new(0),
held_dropped: Arc::new(AtomicU32::new(0)),
count_live: false,
live: Arc::new(AtomicU32::new(0)),
max_live: Arc::new(AtomicU32::new(0)),
}
}
fn requests(&self) -> Vec<(String, Option<String>)> {
self.requests.lock().unwrap().clone()
}
fn bodies_drained(&self) -> u32 {
self.bodies_drained.load(Ordering::Relaxed)
}
fn track_drain(&self, result: GetResult) -> GetResult {
let GetResult {
payload,
meta,
range,
attributes,
extensions,
} = result;
let stream = match payload {
GetResultPayload::Stream(s) => s,
#[allow(unreachable_patterns)]
_ => unreachable!("InMemory yields streams"),
};
let stream = if self.count_live {
let live = Arc::clone(&self.live);
let max_live = Arc::clone(&self.max_live);
stream
.map(move |item| {
item.map(|bytes| {
let now = live.fetch_add(1, Ordering::Relaxed) + 1;
max_live.fetch_max(now, Ordering::Relaxed);
Bytes::from_owner(LiveWindow {
bytes,
live: Arc::clone(&live),
})
})
})
.boxed()
} else {
stream
};
let drained = Arc::clone(&self.bodies_drained);
let marker = futures_util::stream::poll_fn(move |_| {
drained.fetch_add(1, Ordering::Relaxed);
std::task::Poll::Ready(None::<object_store::Result<Bytes>>)
});
GetResult {
payload: GetResultPayload::Stream(Box::pin(stream.chain(marker))),
meta,
range,
attributes,
extensions,
}
}
fn recorded_ranges(&self) -> Vec<RangeKind> {
self.ranges.lock().unwrap().clone()
}
fn generic(what: &str) -> object_store::Error {
object_store::Error::Generic {
store: "flaky",
source: what.to_owned().into(),
}
}
fn cut_error(fatal: bool, text: Option<&str>) -> object_store::Error {
if fatal {
object_store::Error::PermissionDenied {
path: "k".into(),
source: text.unwrap_or("injected permission loss").to_owned().into(),
}
} else {
Self::generic(text.unwrap_or("injected stream cut"))
}
}
}
impl fmt::Display for FlakyStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "FlakyStore")
}
}
#[async_trait::async_trait]
impl ObjectStore for FlakyStore {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: object_store::PutOptions,
) -> object_store::Result<object_store::PutResult> {
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: object_store::PutMultipartOptions,
) -> object_store::Result<Box<dyn object_store::MultipartUpload>> {
self.inner.put_multipart_opts(location, opts).await
}
async fn get_opts(
&self,
location: &Path,
options: GetOptions,
) -> object_store::Result<GetResult> {
let nth = self.gets.fetch_add(1, Ordering::Relaxed) + 1;
self.ranges.lock().unwrap().push(range_kind(&options.range));
self.requests
.lock()
.unwrap()
.push((location.to_string(), options.if_match.clone()));
if nth == self.hold_on {
self.held.fetch_add(1, Ordering::Relaxed);
let _dropped = DropCount(Arc::clone(&self.held_dropped));
self.release.acquire().await.unwrap().forget();
}
if self
.fail_gets
.try_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
.is_ok()
{
return Err(Self::generic(
self.error_text.as_deref().unwrap_or("injected get failure"),
));
}
if let Some(GetRange::Bounded(r)) = &options.range
&& r.end > self.fail_past
&& self
.fail_past_gets
.try_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
.is_ok()
{
return Err(Self::generic(
self.error_text
.as_deref()
.unwrap_or("injected get failure past the range end"),
));
}
let result = self.inner.get_opts(location, options).await?;
if nth == self.empty_on {
return Ok(GetResult {
payload: GetResultPayload::Stream(Box::pin(futures_util::stream::empty())),
..result
});
}
if self
.cut_streams
.try_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
.is_ok()
{
let cut_after = self.cut_after;
let cut_fatal = self.cut_fatal;
let error_text = self.error_text.clone();
let GetResult {
payload,
meta,
range,
attributes,
extensions,
} = result;
let inner_stream = match payload {
GetResultPayload::Stream(s) => s,
#[allow(unreachable_patterns)]
_ => unreachable!("InMemory yields streams"),
};
let mut sent = 0usize;
let cut = inner_stream.flat_map(move |item| {
let out: Vec<object_store::Result<Bytes>> = match item {
Ok(bytes) => {
if sent >= cut_after {
vec![Err(Self::cut_error(cut_fatal, error_text.as_deref()))]
} else {
let take = bytes.len().min(cut_after - sent);
sent += take;
let mut out = vec![Ok(bytes.slice(0..take))];
if sent >= cut_after {
out.push(Err(Self::cut_error(
cut_fatal,
error_text.as_deref(),
)));
}
out
}
}
Err(e) => vec![Err(e)],
};
futures_util::stream::iter(out)
});
return Ok(GetResult {
payload: GetResultPayload::Stream(Box::pin(cut)),
meta,
range,
attributes,
extensions,
});
}
Ok(self.track_drain(result))
}
fn delete_stream(
&self,
locations: futures_util::stream::BoxStream<'static, object_store::Result<Path>>,
) -> futures_util::stream::BoxStream<'static, object_store::Result<Path>> {
self.inner.delete_stream(locations)
}
fn list(
&self,
prefix: Option<&Path>,
) -> futures_util::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>>
{
self.inner.list(prefix)
}
async fn list_with_delimiter(
&self,
prefix: Option<&Path>,
) -> object_store::Result<object_store::ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy_opts(
&self,
from: &Path,
to: &Path,
options: object_store::CopyOptions,
) -> object_store::Result<()> {
self.inner.copy_opts(from, to, options).await
}
}
async fn seeded(objects: &[(&str, &[u8])]) -> InMemory {
let store = InMemory::new();
for (key, body) in objects {
store
.put(&Path::from(*key), PutPayload::from(body.to_vec()))
.await
.unwrap();
}
store
}
async fn listed(store: &Arc<dyn ObjectStore>) -> Vec<ObjectEntry> {
list_all(store, None).await.unwrap()
}
async fn collect_fetch(
store: Arc<dyn ObjectStore>,
slice: Vec<ObjectEntry>,
start_ordinal: u32,
chunk_bytes: usize,
range_bytes: usize,
) -> Vec<ChunkMsg> {
let (tx, mut rx) = mpsc::channel(64);
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal,
resume_etag: None,
chunk_bytes,
range_bytes,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
let task = tokio::spawn(run_fetcher(params));
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
}
task.await.unwrap();
msgs
}
fn assembled(msgs: &[ChunkMsg]) -> Vec<(String, Vec<u8>)> {
let mut out: Vec<(String, Vec<u8>)> = Vec::new();
for m in msgs {
match m {
ChunkMsg::ObjectStart { key, .. } => out.push((key.clone(), Vec::new())),
ChunkMsg::Chunk(b) => out.last_mut().unwrap().1.extend_from_slice(b),
ChunkMsg::ObjectEnd => {}
ChunkMsg::LaneFailed(e) => panic!("lane failed: {e}"),
}
}
out
}
const ENCODED_KEYS: [&str; 4] = [
"dir/a b.ndjson",
"dir/a%20b.ndjson",
"dir/x#y[1].ndjson",
"dir/\u{e9}.ndjson",
];
async fn seeded_verbatim(keys: &[&str]) -> Arc<dyn ObjectStore> {
let store = InMemory::new();
for key in keys {
store
.put(
&Path::parse(*key).unwrap(),
PutPayload::from(key.as_bytes().to_vec()),
)
.await
.unwrap();
}
Arc::new(store)
}
fn verbatim_bodies() -> Vec<(String, Vec<u8>)> {
let mut want: Vec<_> = ENCODED_KEYS
.iter()
.map(|k| ((*k).to_owned(), k.as_bytes().to_vec()))
.collect();
want.sort();
want
}
#[tokio::test(start_paused = true)]
async fn the_next_members_get_starts_before_the_lane_reads() {
let store = Arc::new(FlakyStore::new(
seeded(&[("a.ndjson", b"a\n"), ("b.ndjson", b"b\n")]).await,
));
let slice = listed(&(Arc::clone(&store) as Arc<dyn ObjectStore>)).await;
assert!(slice.iter().all(|e| e.etag.is_some()), "the ranged path");
let (tx, _rx) = mpsc::channel(crate::lane_ctx::LANE_HANDOFF_CHUNKS);
let params = FetcherParams {
split: "s3-test".into(),
store: Arc::clone(&store) as Arc<dyn ObjectStore>,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
let task = tokio::spawn(run_fetcher(params));
park().await;
assert_eq!(
store.gets.load(Ordering::Relaxed),
2,
"the second member's GET started with nothing read"
);
task.abort();
}
#[tokio::test]
async fn ranged_read_fetches_the_listed_key_verbatim() {
let store = seeded_verbatim(&ENCODED_KEYS).await;
let slice = listed(&store).await;
assert!(slice.iter().all(|e| e.etag.is_some()));
let got = assembled(&collect_fetch(store, slice, 0, 4, 64).await);
assert_eq!(got, verbatim_bodies());
}
#[tokio::test]
async fn streaming_read_fetches_the_listed_key_verbatim() {
let store = seeded_verbatim(&ENCODED_KEYS).await;
let mut slice = listed(&store).await;
for entry in &mut slice {
entry.etag = None;
}
let got = assembled(&collect_fetch(store, slice, 0, 4, 64).await);
assert_eq!(got, verbatim_bodies());
}
#[tokio::test]
async fn an_unparseable_key_poisons_the_split() {
for etag in [Some("\"e\""), None] {
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let slice = vec![entry("dir//x.ndjson", 1, etag)];
let msgs = collect_fetch(store, slice, 0, 4, 64).await;
let poison = msgs.iter().find_map(|m| match m {
ChunkMsg::LaneFailed(SplitFailure::Poison(kind, reason)) => {
Some((*kind, reason.clone()))
}
_ => None,
});
let (kind, reason) = poison.expect("an unparseable key is split poison");
assert_eq!(kind, PoisonKind::NotFound);
assert!(
reason.contains("\"dir//x.ndjson\" is not an object key"),
"{reason}"
);
}
}
#[tokio::test]
async fn a_key_the_store_would_rewrite_poisons_the_split() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"other\n")]).await);
let msgs = collect_fetch(store, vec![entry("/p/a", 6, None)], 0, 4, 64).await;
assert!(
!msgs.iter().any(|m| matches!(m, ChunkMsg::Chunk(_))),
"no bytes of p/a are read for /p/a"
);
let reason = msgs.iter().find_map(|m| match m {
ChunkMsg::LaneFailed(SplitFailure::Poison(PoisonKind::NotFound, reason)) => {
Some(reason.clone())
}
_ => None,
});
let reason = reason.expect("a rewritten key is not_found split poison");
assert!(reason.contains("\"/p/a\" is not an object key"), "{reason}");
assert!(reason.contains("\"p/a\" in its place"), "{reason}");
}
#[tokio::test]
async fn streams_objects_in_order_with_bounded_chunks() {
let store: Arc<dyn ObjectStore> =
Arc::new(seeded(&[("p/a", b"aaaaaaaaaa"), ("p/b", b"bb")]).await);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 4, 64).await;
for m in &msgs {
if let ChunkMsg::Chunk(b) = m {
assert!(b.len() <= 4, "chunk over the bound: {}", b.len());
}
}
assert_eq!(
assembled(&msgs),
vec![
("p/a".to_string(), b"aaaaaaaaaa".to_vec()),
("p/b".to_string(), b"bb".to_vec())
]
);
}
#[tokio::test]
async fn start_ordinal_skips_committed_objects() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", b"first"), ("p/b", b"second")]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 1, 64, 64).await;
assert_eq!(
assembled(&msgs),
vec![("p/b".to_string(), b"second".to_vec())]
);
let keys: Vec<String> = flaky.requests().into_iter().map(|(k, _)| k).collect();
assert_eq!(keys, ["p/b"], "a skipped member is never read");
}
#[tokio::test]
async fn transient_get_failures_are_retried() {
let flaky = FlakyStore::new(seeded(&[("p/a", b"payload")]).await);
flaky.fail_gets.store(2, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"payload".to_vec())]
);
}
#[tokio::test]
async fn mid_stream_cut_resumes_with_a_ranged_get_without_gap_or_dup() {
let body = b"0123456789abcdefghij";
let flaky = FlakyStore {
cut_after: 7,
..FlakyStore::new(seeded(&[("p/a", body)]).await)
};
flaky.cut_streams.store(1, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 4, 8).await;
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), body.to_vec())],
"a window that fails mid-read is retried cleanly, without gap or dup"
);
}
#[tokio::test]
async fn missing_object_poisons_the_split() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"x")]).await);
let slice = vec![entry("p/ghost", 1, Some("\"e\""))];
let (tx, mut rx) = mpsc::channel(8);
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
tokio::spawn(run_fetcher(params));
let mut saw_poison = false;
while let Some(m) = rx.recv().await {
if let ChunkMsg::LaneFailed(failure) = m {
let SplitFailure::Poison(kind, reason) = failure else {
panic!("a missing object is split poison, not pipeline-fatal: {failure}");
};
assert_eq!(kind, PoisonKind::NotFound);
assert!(reason.contains("p/ghost"), "{reason}");
saw_poison = true;
}
}
assert!(saw_poison);
}
#[tokio::test]
async fn stale_etag_pin_poisons_the_split() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"new content")]).await);
let mut slice = listed(&store).await;
slice[0].etag = Some("\"stale\"".into());
let (tx, mut rx) = mpsc::channel(8);
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
tokio::spawn(run_fetcher(params));
let mut saw_poison = false;
while let Some(m) = rx.recv().await {
if let ChunkMsg::LaneFailed(failure) = m {
assert!(
matches!(failure, SplitFailure::Poison(PoisonKind::EtagDrift, _)),
"an overwrite under the pin is split poison: {failure}"
);
saw_poison = true;
}
}
assert!(saw_poison, "a failed precondition must poison the split");
}
#[tokio::test(start_paused = true)]
async fn pause_halts_fetching_and_resume_continues() {
let flaky = Arc::new(FlakyStore::new(seeded(&[("p/a", b"abcdefgh")]).await));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(1); let pause = Arc::new(AtomicBool::new(true));
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 2,
range_bytes: 64,
tx,
pause: Arc::clone(&pause),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
tokio::spawn(run_fetcher(params));
tokio::time::sleep(Duration::from_millis(80)).await;
assert!(
rx.try_recv().is_err(),
"paused fetcher must not deliver messages"
);
assert_eq!(flaky.gets.load(Ordering::Relaxed), 0, "nor read ahead");
pause.store(false, Ordering::Relaxed);
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
}
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"abcdefgh".to_vec())]
);
}
#[tokio::test]
async fn dropping_the_lane_ends_the_fetcher() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", &[b'x'; 4096][..])]).await);
let slice = listed(&store).await;
let (tx, rx) = mpsc::channel(1);
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 8,
range_bytes: 4096,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
let task = tokio::spawn(run_fetcher(params));
drop(rx);
tokio::time::timeout(Duration::from_secs(5), task)
.await
.expect("fetcher must end promptly when the lane is dropped")
.unwrap();
}
#[tokio::test]
async fn pinned_reads_walk_the_object_in_bounded_contiguous_windows() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", b"0123456789abcdefghij")]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
assert!(
slice[0].etag.is_some(),
"InMemory reports an ETag → pinned path"
);
let msgs = collect_fetch(Arc::clone(&store), slice, 0, 4, 8).await;
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"0123456789abcdefghij".to_vec())]
);
assert_eq!(
flaky.recorded_ranges(),
vec![
RangeKind::Bounded(0, 8),
RangeKind::Bounded(8, 16),
RangeKind::Bounded(16, 20),
],
"windows must be bounded, contiguous, and cover the object exactly"
);
}
const DEPTH: usize = crate::lane_ctx::LANE_HANDOFF_CHUNKS;
fn fetcher(
store: Arc<dyn ObjectStore>,
slice: Vec<ObjectEntry>,
chunk_bytes: usize,
range_bytes: usize,
tx: mpsc::Sender<ChunkMsg>,
) -> FetcherParams {
FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes,
range_bytes,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
}
}
async fn park() {
tokio::time::sleep(Duration::from_millis(80)).await;
}
async fn drain(rx: &mut mpsc::Receiver<ChunkMsg>) -> Vec<ChunkMsg> {
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
}
msgs
}
fn debugged(msgs: &[ChunkMsg]) -> Vec<String> {
msgs.iter().map(|m| format!("{m:?}")).collect()
}
#[tokio::test(start_paused = true)]
async fn a_blocked_lane_holds_two_windows_and_fetches_no_further() {
let flaky = Arc::new(FlakyStore::new(seeded(&[("p/a", &[b'x'; 48][..])]).await));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, _rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 16, tx)));
park().await;
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(0, 16), RangeKind::Bounded(16, 32)],
"the second window is read while the first is handed on"
);
assert_eq!(flaky.bodies_drained(), 2);
task.abort();
}
#[tokio::test(start_paused = true)]
async fn the_next_members_get_starts_while_the_last_window_is_queued() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", &[b'a'; 16][..]), ("p/b", &[b'b'; 12][..])]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
assert_ne!(slice[0].etag, slice[1].etag);
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(
Arc::clone(&store),
slice.clone(),
2,
16,
tx,
)));
park().await;
assert_eq!(
flaky.requests(),
vec![
("p/a".to_string(), slice[0].etag.clone()),
("p/b".to_string(), slice[1].etag.clone()),
]
);
let msgs = drain(&mut rx).await;
task.await.unwrap();
let unblocked = collect_fetch(store, slice, 0, 2, 16).await;
assert_eq!(debugged(&msgs), debugged(&unblocked));
}
#[tokio::test(start_paused = true)]
async fn a_zero_byte_member_is_skipped_by_the_read_ahead() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[
("p/a", &[b'a'; 16][..]),
("p/b", b""),
("p/c", &[b'c'; 12][..]),
])
.await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(
Arc::clone(&store),
slice.clone(),
2,
16,
tx,
)));
park().await;
let keys: Vec<String> = flaky.requests().into_iter().map(|(k, _)| k).collect();
assert_eq!(keys, ["p/a", "p/c"]);
let msgs = drain(&mut rx).await;
task.await.unwrap();
let unblocked = collect_fetch(store, slice, 0, 2, 16).await;
assert_eq!(debugged(&msgs), debugged(&unblocked));
assert_eq!(
assembled(&msgs),
vec![
("p/a".to_string(), vec![b'a'; 16]),
("p/b".to_string(), Vec::new()),
("p/c".to_string(), vec![b'c'; 12]),
]
);
}
#[tokio::test(start_paused = true)]
async fn a_paused_lane_finishes_its_read_ahead_and_fetches_no_further() {
let mut flaky = FlakyStore::new(seeded(&[("p/a", &[b'x'; 48][..])]).await);
flaky.hold_on = 2;
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(DEPTH);
let params = fetcher(store, slice, 2, 16, tx);
let pause = Arc::clone(¶ms.pause);
let task = tokio::spawn(run_fetcher(params));
park().await;
assert!(matches!(
rx.recv().await,
Some(ChunkMsg::ObjectStart { .. })
));
assert!(matches!(rx.recv().await, Some(ChunkMsg::Chunk(_))));
pause.store(true, Ordering::Relaxed);
park().await;
assert_eq!(flaky.bodies_drained(), 1, "the read-ahead is still held");
flaky.release.add_permits(1);
park().await;
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(0, 16), RangeKind::Bounded(16, 32)]
);
assert_eq!(
flaky.bodies_drained(),
2,
"the read-ahead finished while the fetcher waited out the pause"
);
task.abort();
}
#[tokio::test(start_paused = true)]
async fn a_blocked_ranged_read_stops_at_the_owned_bytes() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", b"aa\nbbbb\ncc\r\n \ndddddd\nee\nfff")]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(1);
let mut params = fetcher(store, slice, 2, 4, tx);
params.range = Some(SplitRange::new(8, 21, b'\n'));
let task = tokio::spawn(run_fetcher(params));
park().await;
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(7, 11), RangeKind::Bounded(11, 15)]
);
let msgs = drain(&mut rx).await;
task.await.unwrap();
assert_eq!(assembled(&msgs)[0].1, b"cc\r\n \ndddddd\n");
assert_eq!(
flaky.recorded_ranges(),
vec![
RangeKind::Bounded(7, 11),
RangeKind::Bounded(11, 15),
RangeKind::Bounded(15, 19),
RangeKind::Bounded(19, 23),
]
);
}
#[tokio::test(start_paused = true)]
async fn a_lane_holds_at_most_three_windows() {
let mut flaky = FlakyStore::new(seeded(&[("p/a", &[b'x'; 80][..])]).await);
flaky.count_live = true;
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 16, tx)));
loop {
park().await;
if rx.recv().await.is_none() {
break;
}
}
task.await.unwrap();
assert_eq!(flaky.recorded_ranges().len(), 5);
assert_eq!(flaky.live.load(Ordering::Relaxed), 0);
assert_eq!(flaky.max_live.load(Ordering::Relaxed), 3);
}
#[tokio::test(start_paused = true)]
async fn a_failed_read_ahead_reaches_the_lane_after_the_earlier_window() {
for (origin, overwritten) in [(notified(None), true), (Origin::Lister, false)] {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(
HeadStore::new(inner, 2).head_answers(vec![Ok(meta(Some("\"moved\""), 30))]),
);
let (tx, mut rx) = mpsc::channel(DEPTH);
let mut params = fetcher(store, vec![entry], 2, 10, tx);
params.origin = origin;
let task = tokio::spawn(run_fetcher(params));
park().await;
let mut msgs = drain(&mut rx).await;
task.await.unwrap();
let failure = match msgs.pop() {
Some(ChunkMsg::LaneFailed(failure)) => failure,
other => panic!("the read did not fail: {other:?}"),
};
match failure {
SplitFailure::Overwritten(_) => assert!(overwritten),
SplitFailure::Poison(PoisonKind::EtagDrift, _) => assert!(!overwritten),
other => panic!("unexpected failure: {other}"),
}
assert!(matches!(msgs[0], ChunkMsg::ObjectStart { .. }));
assert_eq!(msgs.len(), 6, "ObjectStart and five chunks");
assert_eq!(assembled(&msgs)[0].1, b"0123456789");
}
}
#[tokio::test(start_paused = true)]
async fn a_failed_member_read_ahead_follows_that_members_start() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", &[b'a'; 16][..])]).await);
let mut slice = listed(&store).await;
slice.push(entry("p/gone", 12, Some("\"gone\"")));
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 16, tx)));
park().await;
let msgs = drain(&mut rx).await;
task.await.unwrap();
let tail: Vec<String> = debugged(&msgs[msgs.len() - 3..])
.into_iter()
.map(|m| m.split([' ', '(']).next().unwrap().to_owned())
.collect();
assert_eq!(tail, ["ObjectEnd", "ObjectStart", "LaneFailed"]);
assert!(matches!(
msgs.last(),
Some(ChunkMsg::LaneFailed(SplitFailure::Poison(
PoisonKind::NotFound,
_
)))
));
}
#[tokio::test(start_paused = true)]
async fn each_window_gets_its_own_retry_budget() {
for (window_two_failures, completes) in [(MAX_ATTEMPTS - 1, true), (MAX_ATTEMPTS, false)] {
let mut flaky = FlakyStore::new(seeded(&[("p/a", b"0123456789abcdef")]).await);
flaky.fail_gets = AtomicU32::new(MAX_ATTEMPTS - 1);
flaky.fail_past = 8;
flaky.fail_past_gets = AtomicU32::new(window_two_failures);
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 8, tx)));
let msgs = drain(&mut rx).await;
task.await.unwrap();
if completes {
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"0123456789abcdef".to_vec())]
);
assert_eq!(flaky.gets.load(Ordering::Relaxed), 2 * MAX_ATTEMPTS);
} else {
assert!(matches!(
msgs.last(),
Some(ChunkMsg::LaneFailed(SplitFailure::Poison(
PoisonKind::RetriesExhausted,
_
)))
));
assert_eq!(assembled(&msgs[..msgs.len() - 1])[0].1, b"01234567");
assert_eq!(flaky.gets.load(Ordering::Relaxed), 2 * MAX_ATTEMPTS);
}
}
}
#[tokio::test(start_paused = true)]
async fn dropping_the_lane_mid_read_ahead_ends_the_fetcher() {
for paused in [false, true] {
let mut flaky = FlakyStore::new(seeded(&[("p/a", &[b'x'; 48][..])]).await);
flaky.hold_on = 2;
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(DEPTH);
let params = fetcher(store, slice, 2, 16, tx);
params.pause.store(paused, Ordering::Relaxed);
let pause = Arc::clone(¶ms.pause);
let task = tokio::spawn(run_fetcher(params));
park().await;
if paused {
assert_eq!(
flaky.held.load(Ordering::Relaxed),
0,
"paused before any GET"
);
pause.store(false, Ordering::Relaxed);
park().await;
pause.store(true, Ordering::Relaxed);
assert!(rx.recv().await.is_some());
park().await;
}
assert_eq!(flaky.held.load(Ordering::Relaxed), 1);
drop(rx);
tokio::time::timeout(Duration::from_secs(5), task)
.await
.expect("the fetcher ends when its lane is dropped")
.unwrap();
assert_eq!(flaky.held_dropped.load(Ordering::Relaxed), 1);
}
}
#[tokio::test(start_paused = true)]
async fn an_empty_ranged_read_poisons_the_unit() {
let mut flaky = FlakyStore::new(seeded(&[("p/a", b"0123456789")]).await);
flaky.empty_on = 1;
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 4, 8).await;
assert!(
matches!(
msgs.last(),
Some(ChunkMsg::LaneFailed(SplitFailure::Poison(PoisonKind::Undecodable, reason)))
if reason.contains("returned no bytes")
),
"{msgs:?}"
);
}
#[tokio::test(start_paused = true)]
async fn a_pause_during_the_read_ahead_wait_opens_no_further_get() {
let mut flaky = FlakyStore::new(seeded(&[("p/a", &[b'x'; 48][..])]).await);
flaky.hold_on = 2;
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, _rx) = mpsc::channel(DEPTH);
let params = fetcher(store, slice, 16, 16, tx);
let pause = Arc::clone(¶ms.pause);
let task = tokio::spawn(run_fetcher(params));
park().await;
assert_eq!(flaky.held.load(Ordering::Relaxed), 1);
pause.store(true, Ordering::Relaxed);
flaky.release.add_permits(1);
park().await;
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(0, 16), RangeKind::Bounded(16, 32)],
"a paused lane opened a GET"
);
task.abort();
}
#[tokio::test(start_paused = true)]
async fn a_read_ahead_drains_while_the_member_boundary_waits() {
for depth in [9, 10] {
let mut flaky = FlakyStore::new(
seeded(&[("p/a", &[b'a'; 16][..]), ("p/b", &[b'b'; 12][..])]).await,
);
flaky.hold_on = 2;
let flaky = Arc::new(flaky);
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(depth);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 16, tx)));
park().await;
assert_eq!(flaky.held.load(Ordering::Relaxed), 1, "depth {depth}");
flaky.release.add_permits(1);
park().await;
assert_eq!(flaky.bodies_drained(), 2, "depth {depth}");
let msgs = drain(&mut rx).await;
task.await.unwrap();
assert_eq!(assembled(&msgs)[1].1, vec![b'b'; 12]);
}
}
#[tokio::test(start_paused = true)]
async fn a_member_without_an_etag_is_not_read_ahead() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", &[b'a'; 16][..]), ("p/b", &[b'b'; 12][..])]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let mut slice = listed(&store).await;
slice[1].etag = None;
let (tx, mut rx) = mpsc::channel(DEPTH);
let task = tokio::spawn(run_fetcher(fetcher(store, slice, 2, 16, tx)));
park().await;
assert_eq!(flaky.recorded_ranges(), vec![RangeKind::Bounded(0, 16)]);
let msgs = drain(&mut rx).await;
task.await.unwrap();
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(0, 16), RangeKind::Full]
);
assert_eq!(assembled(&msgs)[1].1, vec![b'b'; 12]);
}
fn expected_messages(
slice: &[ObjectEntry],
bodies: &[Vec<u8>],
chunk_bytes: usize,
range_bytes: usize,
) -> Vec<String> {
let mut msgs = Vec::new();
for (ordinal, (entry, body)) in slice.iter().zip(bodies).enumerate() {
msgs.push(ChunkMsg::ObjectStart {
ordinal: ordinal as u32,
key: entry.key.clone(),
last_modified_ms: entry.last_modified_ms,
});
for window in body.chunks(range_bytes) {
for chunk in window.chunks(chunk_bytes) {
msgs.push(ChunkMsg::Chunk(Bytes::copy_from_slice(chunk)));
}
}
msgs.push(ChunkMsg::ObjectEnd);
}
debugged(&msgs)
}
proptest::proptest! {
#![proptest_config(proptest::prelude::ProptestConfig::with_cases(64))]
#[test]
fn a_unit_streams_its_members_in_order_under_any_windows(
sizes in proptest::collection::vec(0usize..48, 1..7),
range_bytes in 1usize..20,
chunk_bytes in 1usize..8,
depth in 1usize..6,
park_every in 1usize..4,
) {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.unwrap();
let bodies: Vec<Vec<u8>> = sizes
.iter()
.enumerate()
.map(|(i, n)| (0..*n).map(|b| b'a' + ((b + i) % 26) as u8).collect())
.collect();
let (msgs, slice, ranges) = rt.block_on(async {
let keys: Vec<String> = (0..bodies.len()).map(|i| format!("p/{i}")).collect();
let objects: Vec<(&str, &[u8])> = keys
.iter()
.zip(&bodies)
.map(|(k, b)| (k.as_str(), b.as_slice()))
.collect();
let flaky = Arc::new(FlakyStore::new(seeded(&objects).await));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(depth);
let task = tokio::spawn(run_fetcher(fetcher(
store,
slice.clone(),
chunk_bytes,
range_bytes,
tx,
)));
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
if msgs.len() % park_every == 0 {
park().await;
}
}
task.await.unwrap();
(msgs, slice, flaky.recorded_ranges())
});
proptest::prop_assert_eq!(
debugged(&msgs),
expected_messages(&slice, &bodies, chunk_bytes, range_bytes)
);
let want: Vec<RangeKind> = bodies
.iter()
.flat_map(|body| {
let size = body.len() as u64;
(0..size)
.step_by(range_bytes)
.map(move |w| RangeKind::Bounded(w, (w + range_bytes as u64).min(size)))
})
.collect();
proptest::prop_assert_eq!(ranges, want);
}
}
async fn no_etag_mid_stream_failure(cut_fatal: bool) -> SplitFailure {
no_etag_mid_stream_failure_with(cut_fatal, None).await
}
async fn no_etag_mid_stream_failure_with(
cut_fatal: bool,
error_text: Option<String>,
) -> SplitFailure {
let flaky = FlakyStore {
cut_after: 4,
cut_fatal,
error_text,
..FlakyStore::new(seeded(&[("p/a", b"0123456789")]).await)
};
flaky.cut_streams.store(1, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let mut slice = listed(&store).await;
slice[0].etag = None; let msgs = collect_fetch(store, slice, 0, 64, 64).await;
match msgs.into_iter().last() {
Some(ChunkMsg::LaneFailed(failure)) => failure,
other => panic!("a mid-stream break without an ETag must fail the lane: {other:?}"),
}
}
#[tokio::test]
async fn mid_stream_permission_loss_without_etag_is_pipeline_fatal() {
let failure = no_etag_mid_stream_failure(true).await;
let SplitFailure::Fatal(e) = failure else {
panic!("PermissionDenied mid-stream must be pipeline-fatal, got: {failure}");
};
assert!(e.to_string().contains("p/a"), "{e}");
}
#[tokio::test]
async fn mid_stream_generic_failure_without_etag_poisons_the_split() {
let failure = no_etag_mid_stream_failure(false).await;
let SplitFailure::Poison(kind, reason) = failure else {
panic!("a generic mid-stream break is split poison, got: {failure}");
};
assert_eq!(kind, PoisonKind::Undecodable);
assert!(reason.contains("no ETag"), "{reason}");
}
#[tokio::test]
async fn an_object_without_an_etag_uses_the_streaming_fallback() {
let flaky = Arc::new(FlakyStore::new(seeded(&[("p/a", b"hello world")]).await));
let store: Arc<dyn ObjectStore> = flaky.clone();
let mut slice = listed(&store).await;
slice[0].etag = None; let msgs = collect_fetch(Arc::clone(&store), slice, 0, 4, 4).await;
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"hello world".to_vec())]
);
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Full],
"the fallback issues a single un-ranged GET, never a bounded window"
);
}
#[tokio::test(start_paused = true)]
async fn a_paused_fetcher_exits_when_its_lane_drops_the_receiver() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"aaaa")]).await);
let slice = listed(&store).await;
let (tx, rx) = mpsc::channel(8);
let params = FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(true)),
retry_base: Duration::from_millis(1),
retries: None,
range: None,
bytes_read: None,
origin: Origin::Lister,
};
let task = tokio::spawn(run_fetcher(params));
tokio::time::sleep(Duration::from_millis(80)).await;
drop(rx);
tokio::time::timeout(Duration::from_secs(5), task)
.await
.expect("a paused fetcher exits once the receiver is gone")
.unwrap();
}
#[tokio::test]
async fn a_tls_rejection_on_a_read_is_fatal_and_named() {
let addr =
spate_test::tls_alert_server(b"", u8::from(rustls::AlertDescription::HandshakeFailure));
let store: Arc<dyn ObjectStore> = Arc::new(crate::test_servers::tls_store_at(
&format!("https://{addr}"),
&crate::test_servers::TestCa::new("any"),
));
let slice = vec![ObjectEntry {
key: "p/a".to_owned(),
size: 10,
etag: Some("\"e\"".to_owned()),
last_modified_ms: 0,
}];
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
let Some(ChunkMsg::LaneFailed(SplitFailure::Fatal(e))) = msgs.into_iter().last() else {
panic!("expected a pipeline-fatal lane failure");
};
assert!(e.to_string().contains("HandshakeFailure"), "{e}");
}
fn large_body(code: &str) -> String {
format!(
"<Error><Code>{code}</Code><Message>{}</Message></Error>",
"x".repeat(64 * 1024)
)
}
fn assert_cut(reason: &str) {
assert!(reason.len() < 2048, "{} bytes", reason.len());
assert!(reason.contains('…'), "{reason}");
}
async fn read_answered(status: &'static str, body: String, key: &str) -> SplitFailure {
let url = crate::test_servers::sequence_server(vec![(status, body)]);
let store: Arc<dyn ObjectStore> = Arc::new(crate::test_servers::store_at(&url));
let slice = vec![entry(key, 10, Some("\"e\""))];
match collect_fetch(store, slice, 0, 64, 64)
.await
.into_iter()
.last()
{
Some(ChunkMsg::LaneFailed(failure)) => failure,
other => panic!("a refused read must fail the lane: {other:?}"),
}
}
#[tokio::test]
async fn a_refused_read_with_a_large_body_is_fatal_and_cut() {
let failure = read_answered("403 Forbidden", large_body("AccessDenied"), "p/a").await;
let SplitFailure::Fatal(e) = failure else {
panic!("a 403 read must be fatal, got: {failure}");
};
let reason = e.to_string();
assert!(reason.contains("p/a"), "{reason}");
assert!(reason.contains("403 Forbidden"), "{reason}");
assert!(reason.contains("AccessDenied"), "{reason}");
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
}
#[tokio::test]
async fn a_missing_object_with_a_large_body_poisons_and_cut() {
let failure = read_answered("404 Not Found", large_body("NoSuchKey"), "p/a").await;
let SplitFailure::Poison(PoisonKind::NotFound, reason) = failure else {
panic!("a 404 read must poison as not_found, got: {failure}");
};
assert!(reason.contains("p/a"), "{reason}");
assert!(reason.contains("NoSuchKey"), "{reason}");
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
}
#[tokio::test]
async fn a_multibyte_error_body_is_cut_on_a_character() {
let body = format!(
"<Error><Message>{}</Message></Error>",
"é".repeat(64 * 1024)
);
let failure = read_answered("403 Forbidden", body, "p/a").await;
let SplitFailure::Fatal(e) = failure else {
panic!("a 403 read must be fatal, got: {failure}");
};
let reason = e.to_string();
let text = reason
.split_once("403 Forbidden: ")
.map(|(_, text)| text)
.expect("the store's body follows the status");
assert_eq!(text.chars().count(), 512 + 1, "{text}");
assert!(text.ends_with("é…"), "{text}");
}
#[tokio::test]
async fn a_long_key_refusal_keeps_the_status() {
let key = format!("p/{}", "k".repeat(1000));
let failure = read_answered("403 Forbidden", large_body("AccessDenied"), &key).await;
let SplitFailure::Fatal(e) = failure else {
panic!("a 403 read must be fatal, got: {failure}");
};
let reason = e.to_string();
assert!(reason.contains(&format!("\"{key}\"")), "{reason}");
assert!(
reason.contains("lacked the necessary privileges"),
"{reason}"
);
assert!(reason.contains("403 Forbidden: <Error>"), "{reason}");
assert!(reason.ends_with('…'), "{reason}");
assert!(
reason.len() < 2048 + 2 * key.len(),
"{} bytes",
reason.len()
);
}
#[tokio::test]
async fn a_tls_rejection_on_a_long_key_is_still_named() {
let addr =
spate_test::tls_alert_server(b"", u8::from(rustls::AlertDescription::HandshakeFailure));
let store: Arc<dyn ObjectStore> = Arc::new(crate::test_servers::tls_store_at(
&format!("https://{addr}"),
&crate::test_servers::TestCa::new("any"),
));
let slice = vec![entry(&format!("p/{}", "k".repeat(1000)), 10, Some("\"e\""))];
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
let Some(ChunkMsg::LaneFailed(SplitFailure::Fatal(e))) = msgs.into_iter().last() else {
panic!("expected a pipeline-fatal lane failure");
};
assert!(e.to_string().contains("HandshakeFailure"), "{e}");
}
#[tokio::test(start_paused = true)]
async fn a_read_failing_past_the_budget_with_a_large_error_is_cut() {
let flaky = FlakyStore {
error_text: Some("x".repeat(64 * 1024)),
..FlakyStore::new(seeded(&[("p/a", b"0123456789")]).await)
};
flaky.fail_gets.store(MAX_ATTEMPTS, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
let Some(ChunkMsg::LaneFailed(SplitFailure::Poison(PoisonKind::RetriesExhausted, reason))) =
msgs.into_iter().last()
else {
panic!("a read failing past its budget must poison as retries_exhausted");
};
assert!(reason.contains("attempts"), "{reason}");
assert!(reason.contains("xxxx"), "{reason}");
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
}
#[tokio::test(start_paused = true)]
async fn a_read_of_a_long_key_past_the_budget_keeps_the_status() {
let text = format!(
"Error performing GET http://127.0.0.1/b/p/{} in 1ms - Server returned non-2xx \
status code: 503 Service Unavailable: <Error><Code>SlowDown</Code><Message>{}",
"%E6%95%B0".repeat(60),
"x".repeat(64 * 1024)
);
let flaky = FlakyStore {
error_text: Some(text),
..FlakyStore::new(seeded(&[("p/a", b"0123456789")]).await)
};
flaky.fail_gets.store(MAX_ATTEMPTS, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
let Some(ChunkMsg::LaneFailed(SplitFailure::Poison(PoisonKind::RetriesExhausted, reason))) =
msgs.into_iter().last()
else {
panic!("a read failing past its budget must poison as retries_exhausted");
};
assert!(reason.contains("503 Service Unavailable"), "{reason}");
assert!(reason.contains("SlowDown"), "{reason}");
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
}
#[tokio::test]
async fn a_mid_object_failure_with_a_large_error_is_cut() {
let text = Some("x".repeat(64 * 1024));
let SplitFailure::Fatal(e) = no_etag_mid_stream_failure_with(true, text.clone()).await
else {
panic!("PermissionDenied mid-stream must be pipeline-fatal");
};
let reason = e.to_string();
assert!(
reason.contains("p/a") && reason.contains("mid-object"),
"{reason}"
);
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
let SplitFailure::Poison(PoisonKind::Undecodable, reason) =
no_etag_mid_stream_failure_with(false, text).await
else {
panic!("a generic mid-stream break is split poison");
};
assert!(
reason.contains("p/a") && reason.contains("mid-object"),
"{reason}"
);
assert!(reason.ends_with('…'), "{reason}");
assert_cut(&reason);
}
#[test]
fn a_retried_read_logs_its_store_error_cut() {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.start_paused(true)
.build()
.unwrap();
let lines = spate_test::capture_logs(tracing::Level::WARN, || {
rt.block_on(async {
let flaky = FlakyStore {
error_text: Some("x".repeat(64 * 1024)),
..FlakyStore::new(seeded(&[("p/a", b"0123456789")]).await)
};
flaky.fail_gets.store(2, Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = listed(&store).await;
let msgs = collect_fetch(store, slice, 0, 64, 64).await;
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"0123456789".to_vec())]
);
});
});
let retried: Vec<_> = lines
.iter()
.filter(|l| l.contains("transient object read failure"))
.collect();
assert_eq!(retried.len(), 2, "{lines:?}");
for line in retried {
assert!(line.contains("xxxx…"), "{line}");
assert!(line.len() < 2048, "{} bytes", line.len());
}
}
#[derive(Clone, Copy, Debug)]
struct ReadPlan {
range_bytes: usize,
chunk_bytes: usize,
head_failures: u32,
tail_failures: u32,
cut_after: Option<usize>,
}
impl ReadPlan {
fn windows(range_bytes: usize, chunk_bytes: usize) -> ReadPlan {
ReadPlan {
range_bytes,
chunk_bytes,
head_failures: 0,
tail_failures: 0,
cut_after: None,
}
}
}
fn runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap()
}
fn read_through_lane(
rt: &tokio::runtime::Runtime,
object: &[u8],
range: Option<SplitRange>,
resume: u64,
plan: ReadPlan,
) -> Vec<(i64, Vec<u8>)> {
use crate::lane::S3Lane;
use crate::lane_ctx::SplitTracker;
use crate::offset::Position;
use spate_core::clock::tokio::SystemClock;
use spate_core::source::{PayloadBatch as _, SourceLane as _};
let flaky = FlakyStore {
cut_after: plan.cut_after.unwrap_or(0),
fail_past: range.map_or(u64::MAX, |r| r.end),
..FlakyStore::new(rt.block_on(seeded(&[("p/big.ndjson", object)])))
};
flaky.fail_gets.store(plan.head_failures, Ordering::Relaxed);
flaky
.fail_past_gets
.store(plan.tail_failures, Ordering::Relaxed);
flaky
.cut_streams
.store(u32::from(plan.cut_after.is_some()), Ordering::Relaxed);
let store: Arc<dyn ObjectStore> = Arc::new(flaky);
let slice = rt.block_on(listed(&store));
let (tx, rx) = mpsc::channel(4);
let fetcher = spate_core::task::spawn_on(
"s3-fetcher",
run_fetcher(FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: plan.chunk_bytes,
range_bytes: plan.range_bytes,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range,
bytes_read: None,
origin: Origin::Lister,
}),
rt.handle(),
);
let tracker = Arc::new(SplitTracker::new());
let (poison_tx, poison_rx) = std_mpsc::channel();
let mut lane = S3Lane::new(
spate_core::source::LaneId(0),
spate_core::record::PartitionId(0),
rx,
fetcher,
rt.handle().clone(),
Arc::new(SystemClock),
spate_core::checkpoint::Checkpointer::new().handle(),
crate::config::Compression::None,
Arc::new(|| Box::new(crate::testutil::TestLineFramer::new(1 << 20))),
Some(Position {
ordinal: 0,
record: resume,
}),
Arc::clone(&tracker),
poison_tx,
spate_core::feed::wake_channel().0,
None,
);
let deadline = std::time::Instant::now() + Duration::from_secs(30);
let mut out = Vec::new();
loop {
match lane
.poll(1024, Duration::from_millis(20))
.expect("the lane polls")
{
Some(mut batch) => {
while let Some(p) = batch.next_payload() {
out.push((p.offset, p.bytes.to_vec()));
}
}
None if tracker.terminal().is_some() => return out,
None => {
if let Ok(report) = poison_rx.try_recv() {
panic!("the split poisoned: {}", report.reason);
}
assert!(
std::time::Instant::now() < deadline,
"the lane never reached end of input"
);
}
}
}
}
fn bytes_of(records: Vec<(i64, Vec<u8>)>) -> Vec<Vec<u8>> {
records.into_iter().map(|(_, bytes)| bytes).collect()
}
fn framed_whole(object: &[u8]) -> Vec<Vec<u8>> {
let mut framer = crate::framer::ObjectFramer::new(Arc::new(|| {
Box::new(crate::testutil::TestLineFramer::new(1 << 20))
}));
framer.begin_object(crate::framer::Codec::Plain).unwrap();
framer.push_chunk(Bytes::copy_from_slice(object)).unwrap();
framer.finish_object().unwrap();
std::iter::from_fn(|| framer.pop_record()).collect()
}
fn arb_lines_object() -> impl proptest::strategy::Strategy<Value = Vec<u8>> {
use proptest::prelude::*;
(
proptest::collection::vec(
prop_oneof![
4 => "[a-z]{1,6}".prop_map(|s| format!("{s}\n")),
1 => "[a-z]{1,6}".prop_map(|s| format!("{s}\r\n")),
1 => "[ \t]{0,3}".prop_map(|s| format!("{s}\n")),
1 => "[a-z]{30,60}".prop_map(|s| format!("{s}\n")),
],
1..30,
),
proptest::option::of("[a-z]{1,6}"),
)
.prop_map(|(lines, last)| {
let mut object = lines.concat();
object.push_str(last.as_deref().unwrap_or_default());
object.into_bytes()
})
}
proptest::proptest! {
#![proptest_config(proptest::prelude::ProptestConfig::with_cases(48))]
#[test]
fn ranges_tiling_an_object_emit_each_record_once(
object in arb_lines_object(),
pieces in 1u64..10,
range_bytes in 1usize..32,
chunk_bytes in 1usize..16,
head_failures in 0u32..2,
tail_failures in 0u32..3,
cut_after in proptest::option::of(0usize..24),
) {
let rt = runtime();
let size = object.len() as u64;
let n = pieces.min(size);
let bound = |i: u64| (u128::from(i) * u128::from(size) / u128::from(n)) as u64;
let plan = ReadPlan { range_bytes, chunk_bytes, head_failures, tail_failures, cut_after };
let mut joined = Vec::new();
for i in 0..n {
let range = SplitRange::new(bound(i), bound(i + 1), b'\n');
joined.extend(bytes_of(read_through_lane(&rt, &object, Some(range), 0, plan)));
}
proptest::prop_assert_eq!(joined, framed_whole(&object));
}
}
#[test]
fn a_resumed_range_emits_the_same_tail_under_any_windows() {
let rt = runtime();
let object = b"aa\nbbbb\ncc\r\n \ndddddd\nee\nfff";
let ranges = [
(0, 3),
(3, 8),
(3, 7),
(5, 14),
(8, 21),
(15, 20),
(14, 27),
(16, 27),
];
let reference = ReadPlan::windows(4, 2);
let others = [
ReadPlan {
head_failures: 1,
tail_failures: 1,
..ReadPlan::windows(5, 3)
},
ReadPlan::windows(13, 1),
ReadPlan::windows(64, 64),
];
for (start, end) in ranges {
let range = Some(SplitRange::new(start, end, b'\n'));
let full = read_through_lane(&rt, object, range, 0, reference);
for plan in others {
for k in 0..=full.len() {
let resumed = read_through_lane(&rt, object, range, k as u64, plan);
assert_eq!(
resumed,
full[k..],
"range [{start}, {end}) resumed at {k} under {plan:?}"
);
}
}
}
let read = |start, end| {
bytes_of(read_through_lane(
&rt,
object,
Some(SplitRange::new(start, end, b'\n')),
0,
reference,
))
};
assert_eq!(read(3, 7), vec![b"bbbb".to_vec()]);
assert_eq!(read(5, 14), vec![b"cc".to_vec()]);
assert_eq!(read(8, 21), vec![b"cc".to_vec(), b"dddddd".to_vec()]);
assert!(read(15, 20).is_empty());
assert_eq!(
read(14, 27),
vec![b"dddddd".to_vec(), b"ee".to_vec(), b"fff".to_vec()]
);
assert_eq!(read(16, 27), vec![b"ee".to_vec(), b"fff".to_vec()]);
}
#[tokio::test]
async fn a_ranged_read_counts_the_bytes_it_trims() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", b"aa\nbbbb\ncc\r\n \ndddddd\nee\nfff")]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let trimmed = Arc::new(AtomicU64::new(0));
let (tx, mut rx) = mpsc::channel(64);
let task = tokio::spawn(run_fetcher(FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 2,
range_bytes: 4,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
range: Some(SplitRange::new(8, 21, b'\n')),
bytes_read: Some(spate_core::metrics::Counter::from_arc(Arc::clone(&trimmed))),
origin: Origin::Lister,
}));
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
}
task.await.unwrap();
let forwarded = assembled(&msgs);
assert_eq!(forwarded[0].1, b"cc\r\n \ndddddd\n");
let fetched: u64 = flaky
.recorded_ranges()
.iter()
.map(|r| match r {
RangeKind::Bounded(start, end) => end - start,
other => panic!("a ranged read issues bounded windows, got {other:?}"),
})
.sum();
assert_eq!(fetched, 16, "windows 7..11, 11..15, 15..19 and 19..23");
assert_eq!(
trimmed.load(Ordering::Relaxed) + forwarded[0].1.len() as u64,
fetched
);
}
type Seen = (Option<String>, Option<String>, bool);
#[derive(Debug)]
struct HeadStore {
inner: InMemory,
precondition_on: u32,
range_refused: bool,
gets: AtomicU32,
heads: std::sync::Mutex<std::collections::VecDeque<object_store::Result<ObjectMeta>>>,
seen: std::sync::Mutex<Vec<Seen>>,
}
use object_store::ObjectMeta;
impl HeadStore {
fn new(inner: InMemory, precondition_on: u32) -> HeadStore {
HeadStore {
inner,
precondition_on,
range_refused: false,
gets: AtomicU32::new(0),
heads: std::sync::Mutex::new(std::collections::VecDeque::new()),
seen: std::sync::Mutex::new(Vec::new()),
}
}
fn refusing_the_range(mut self) -> HeadStore {
self.range_refused = true;
self
}
fn head_answers(self, answers: Vec<object_store::Result<ObjectMeta>>) -> HeadStore {
*self.heads.lock().unwrap() = answers.into();
self
}
fn seen(&self) -> Vec<Seen> {
self.seen.lock().unwrap().clone()
}
fn heads(&self) -> usize {
self.seen().iter().filter(|(_, _, head)| *head).count()
}
}
impl fmt::Display for HeadStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "HeadStore")
}
}
#[async_trait::async_trait]
impl ObjectStore for HeadStore {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: object_store::PutOptions,
) -> object_store::Result<object_store::PutResult> {
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: object_store::PutMultipartOptions,
) -> object_store::Result<Box<dyn object_store::MultipartUpload>> {
self.inner.put_multipart_opts(location, opts).await
}
async fn get_opts(
&self,
location: &Path,
options: GetOptions,
) -> object_store::Result<GetResult> {
self.seen.lock().unwrap().push((
options.if_match.clone(),
options.version.clone(),
options.head,
));
if options.head {
let scripted = self.heads.lock().unwrap().pop_front();
if let Some(answer) = scripted {
let meta = answer?;
return Ok(GetResult {
payload: GetResultPayload::Stream(Box::pin(futures_util::stream::empty())),
range: 0..meta.size,
meta,
attributes: Default::default(),
extensions: Default::default(),
});
}
return self.inner.get_opts(location, options).await;
}
if self.gets.fetch_add(1, Ordering::Relaxed) + 1 == self.precondition_on {
if self.range_refused {
return Err(object_store::Error::Generic {
store: "head",
source: "Server returned non-2xx status code: 416 Range Not Satisfiable: "
.into(),
});
}
return Err(object_store::Error::Precondition {
path: location.to_string(),
source: "overwritten".into(),
});
}
self.inner.get_opts(location, options).await
}
fn delete_stream(
&self,
locations: futures_util::stream::BoxStream<'static, object_store::Result<Path>>,
) -> futures_util::stream::BoxStream<'static, object_store::Result<Path>> {
self.inner.delete_stream(locations)
}
fn list(
&self,
prefix: Option<&Path>,
) -> futures_util::stream::BoxStream<'static, object_store::Result<ObjectMeta>> {
self.inner.list(prefix)
}
async fn list_with_delimiter(
&self,
prefix: Option<&Path>,
) -> object_store::Result<object_store::ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy_opts(
&self,
from: &Path,
to: &Path,
options: object_store::CopyOptions,
) -> object_store::Result<()> {
self.inner.copy_opts(from, to, options).await
}
}
fn notified(version: Option<&str>) -> Origin {
Origin::Notification {
version: version.map(str::to_owned),
re_enqueues: 0,
}
}
fn meta(etag: Option<&str>, size: u64) -> ObjectMeta {
ObjectMeta {
location: Path::from("k"),
last_modified: chrono::DateTime::from_timestamp_millis(1_760_000_000_000).unwrap(),
size,
e_tag: etag.map(str::to_owned),
version: None,
}
}
async fn thirty_bytes() -> (InMemory, ObjectEntry) {
let store = seeded(&[("k", b"0123456789abcdefghijABCDEFGHIJ")]).await;
let listed = list_all(&(Arc::new(store.clone()) as Arc<dyn ObjectStore>), None)
.await
.unwrap();
(store, listed[0].clone())
}
async fn fetch_from(
store: Arc<HeadStore>,
entry: ObjectEntry,
origin: Origin,
retries: Option<spate_core::metrics::Counter>,
) -> Vec<ChunkMsg> {
let (tx, mut rx) = mpsc::channel(64);
let task = tokio::spawn(run_fetcher(FetcherParams {
split: "s3-test".into(),
store,
slice: Arc::new(vec![entry]),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 10,
tx,
pause: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries,
range: None,
bytes_read: None,
origin,
}));
let mut msgs = Vec::new();
while let Some(m) = rx.recv().await {
msgs.push(m);
}
task.await.unwrap();
msgs
}
fn last_failure(msgs: Vec<ChunkMsg>) -> SplitFailure {
match msgs.into_iter().last() {
Some(ChunkMsg::LaneFailed(failure)) => failure,
other => panic!("the read did not fail: {other:?}"),
}
}
#[tokio::test]
async fn a_notification_unit_with_a_version_id_reads_that_version() {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(HeadStore::new(inner.clone(), 0));
let msgs = fetch_from(
Arc::clone(&store),
entry.clone(),
notified(Some("v1")),
None,
)
.await;
assert_eq!(assembled(&msgs)[0].1.len(), 30);
let seen = store.seen();
assert_eq!(seen.len(), 3, "{seen:?}");
for (if_match, version, head) in seen {
assert_eq!(if_match, entry.etag);
assert_eq!(version.as_deref(), Some("v1"));
assert!(!head);
}
let store = Arc::new(HeadStore::new(inner, 0));
let unpinned = ObjectEntry {
etag: None,
..entry
};
fetch_from(Arc::clone(&store), unpinned, notified(Some("v1")), None).await;
assert_eq!(store.seen(), [(None, Some("v1".to_owned()), false)]);
}
#[tokio::test]
async fn a_null_or_absent_version_id_reads_by_key() {
let (inner, entry) = thirty_bytes().await;
let scope = crate::unit::Scope {
bucket: Some("bucket".to_owned()),
prefix: Path::from(""),
};
let packing = crate::split::Packing {
target_bytes: 1 << 20,
compression: crate::config::Compression::None,
delimiter: Some(b'\n'),
};
let mut origins = vec![Origin::Lister];
for version in [
None,
Some(serde_json::Value::Null),
Some("".into()),
Some("null".into()),
] {
let mut object = serde_json::json!({ "key": "k", "size": 30 });
if let Some(version) = version {
object["versionId"] = version;
}
let body = serde_json::json!({ "Records": [{
"eventName": "ObjectCreated:Put",
"s3": { "bucket": { "name": "bucket" }, "object": object },
}]});
let crate::unit::Decoded::Read(cursors) =
crate::unit::decode(&serde_json::to_vec(&body).unwrap(), &scope, packing)
else {
panic!("the notification is not read");
};
origins.extend(cursors.into_iter().flatten().map(|part| part.origin));
}
assert_eq!(origins.len(), 5);
for origin in origins {
let store = Arc::new(HeadStore::new(inner.clone(), 0));
fetch_from(Arc::clone(&store), entry.clone(), origin.clone(), None).await;
let seen = store.seen();
assert!(!seen.is_empty());
assert!(
seen.iter().all(|(_, version, _)| version.is_none()),
"{origin:?}: {seen:?}"
);
}
}
#[tokio::test]
async fn a_412_on_a_notification_read_heads_the_key_and_reports_the_overwrite() {
let (inner, entry) = thirty_bytes().await;
let store =
Arc::new(HeadStore::new(inner, 2).head_answers(vec![Ok(meta(Some("new"), 99))]));
let msgs = fetch_from(Arc::clone(&store), entry, notified(None), None).await;
let chunks: Vec<&[u8]> = msgs
.iter()
.filter_map(|m| match m {
ChunkMsg::Chunk(b) => Some(b.as_ref()),
_ => None,
})
.collect();
assert_eq!(chunks, [b"0123456789".as_slice()]);
assert_eq!(store.heads(), 1);
let SplitFailure::Overwritten(overwrite) = last_failure(msgs) else {
panic!("not an overwrite");
};
assert_eq!(
*overwrite,
Overwrite {
key: "k".to_owned(),
etag: Some("\"new\"".to_owned()),
size: 99,
last_modified_ms: 1_760_000_000_000,
}
);
}
#[tokio::test]
async fn a_412_on_a_lister_or_versioned_read_poisons_without_a_head() {
for origin in [Origin::Lister, notified(Some("v1"))] {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(HeadStore::new(inner, 1));
let msgs = fetch_from(Arc::clone(&store), entry, origin.clone(), None).await;
assert_eq!(store.heads(), 0, "{origin:?}");
assert!(
matches!(
last_failure(msgs),
SplitFailure::Poison(PoisonKind::EtagDrift, _)
),
"{origin:?}"
);
}
}
#[tokio::test]
async fn a_head_reporting_the_pinned_etag_poisons_the_read() {
let (inner, entry) = thirty_bytes().await;
let pinned = entry.etag.clone().unwrap();
let store = Arc::new(
HeadStore::new(inner, 1)
.head_answers(vec![Ok(meta(Some(pinned.trim_matches('"')), 30))]),
);
let entry = ObjectEntry {
etag: Some(crate::unit::quoted(pinned.clone())),
..entry
};
let msgs = fetch_from(Arc::clone(&store), entry, notified(None), None).await;
assert_eq!(store.heads(), 1);
assert!(matches!(
last_failure(msgs),
SplitFailure::Poison(PoisonKind::EtagDrift, _)
));
}
#[test]
fn a_split_failure_displays_its_cause() {
let poison = SplitFailure::Poison(PoisonKind::NotFound, "\"k\" is gone".into());
assert_eq!(poison.to_string(), "split poisoned: \"k\" is gone");
let overwrite = SplitFailure::Overwritten(Box::new(Overwrite {
key: "data/k".into(),
etag: None,
size: 1,
last_modified_ms: 0,
}));
assert_eq!(overwrite.to_string(), "\"data/k\" was overwritten");
let fatal = SplitFailure::Fatal(SourceError::Client {
class: ErrorClass::Fatal,
reason: "denied".into(),
});
assert!(fatal.to_string().contains("denied"), "{fatal}");
}
#[tokio::test]
async fn a_head_reporting_no_etag_reports_an_unpinned_overwrite() {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(HeadStore::new(inner, 1).head_answers(vec![Ok(meta(None, 7))]));
let msgs = fetch_from(store, entry, notified(None), None).await;
let SplitFailure::Overwritten(overwrite) = last_failure(msgs) else {
panic!("not an overwrite");
};
assert_eq!((overwrite.etag, overwrite.size), (None, 7));
}
#[tokio::test]
async fn a_head_after_a_412_classifies_404_and_403_as_a_read_does() {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(HeadStore::new(inner.clone(), 2).head_answers(vec![Err(
object_store::Error::NotFound {
path: "k".into(),
source: "gone".into(),
},
)]));
let msgs = fetch_from(store, entry.clone(), notified(None), None).await;
let SplitFailure::Poison(PoisonKind::NotFound, reason) = last_failure(msgs) else {
panic!("a 404 on the HEAD must poison as not_found");
};
assert!(reason.contains("at byte 10"), "{reason}");
assert!(reason.contains("the HEAD after"), "{reason}");
let store = Arc::new(HeadStore::new(inner, 2).head_answers(vec![Err(
object_store::Error::PermissionDenied {
path: "k".into(),
source: "denied".into(),
},
)]));
let msgs = fetch_from(store, entry, notified(None), None).await;
let SplitFailure::Fatal(e) = last_failure(msgs) else {
panic!("a 403 on the HEAD must be fatal");
};
assert!(e.to_string().contains("the HEAD after"), "{e}");
}
#[tokio::test]
async fn a_failed_head_with_a_large_error_keeps_its_suffix() {
let big = || "x".repeat(64 * 1024);
let suffix = "(the HEAD after the store refused the pinned read)";
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(HeadStore::new(inner.clone(), 2).head_answers(vec![Err(
object_store::Error::NotFound {
path: "k".into(),
source: big().into(),
},
)]));
let msgs = fetch_from(store, entry.clone(), notified(None), None).await;
let SplitFailure::Poison(PoisonKind::NotFound, reason) = last_failure(msgs) else {
panic!("a 404 on the HEAD must poison as not_found");
};
assert!(reason.ends_with(suffix), "{reason}");
assert_cut(&reason);
let store = Arc::new(HeadStore::new(inner, 2).head_answers(vec![Err(
object_store::Error::PermissionDenied {
path: "k".into(),
source: big().into(),
},
)]));
let msgs = fetch_from(store, entry, notified(None), None).await;
let SplitFailure::Fatal(e) = last_failure(msgs) else {
panic!("a 403 on the HEAD must be fatal");
};
let reason = e.to_string();
assert!(reason.ends_with(suffix), "{reason}");
assert_cut(&reason);
}
#[tokio::test]
async fn a_416_on_a_notification_read_heads_the_key_and_reports_the_overwrite() {
let (inner, entry) = thirty_bytes().await;
let store = Arc::new(
HeadStore::new(inner, 2)
.refusing_the_range()
.head_answers(vec![Ok(meta(Some("new"), 5))]),
);
let msgs = fetch_from(Arc::clone(&store), entry, notified(None), None).await;
let SplitFailure::Overwritten(overwrite) = last_failure(msgs) else {
panic!("a 416 on a notification read must report an overwrite");
};
assert_eq!(overwrite.size, 5);
assert_eq!(store.heads(), 1);
assert_eq!(store.seen().len(), 3, "two GETs and the HEAD, no retry");
}
#[tokio::test(start_paused = true)]
async fn a_retryable_head_after_a_412_is_retried() {
let transient = || object_store::Error::Generic {
store: "head",
source: "503".into(),
};
let (inner, entry) = thirty_bytes().await;
let retried = Arc::new(AtomicU64::new(0));
let store = Arc::new(
HeadStore::new(inner.clone(), 1)
.head_answers(vec![Err(transient()), Ok(meta(Some("new"), 30))]),
);
let msgs = fetch_from(
Arc::clone(&store),
entry.clone(),
notified(None),
Some(spate_core::metrics::Counter::from_arc(Arc::clone(&retried))),
)
.await;
assert!(matches!(last_failure(msgs), SplitFailure::Overwritten(_)));
assert_eq!(store.heads(), 2);
assert_eq!(retried.load(Ordering::Relaxed), 1);
let failing = (0..MAX_ATTEMPTS).map(|_| Err(transient())).collect();
let store = Arc::new(HeadStore::new(inner, 1).head_answers(failing));
let msgs = fetch_from(Arc::clone(&store), entry, notified(None), None).await;
assert!(matches!(
last_failure(msgs),
SplitFailure::Poison(PoisonKind::RetriesExhausted, _)
));
assert_eq!(store.heads(), MAX_ATTEMPTS as usize);
}
}