use crate::error::classify;
use crate::split_ctx::PoisonKind;
use bytes::Bytes;
use futures_util::StreamExt as _;
use object_store::path::Path;
use object_store::{GetOptions, GetRange, ObjectStore};
use spate_core::coordination::SplitId;
use spate_core::error::{ErrorClass, SourceError};
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),
}
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}"),
}
}
}
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: SplitId,
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) stop: Arc<AtomicBool>,
pub(crate) retry_base: Duration,
pub(crate) retries: Option<spate_core::metrics::Counter>,
}
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,
stop,
retry_base,
retries,
} = params;
for (ordinal, entry) in slice.iter().enumerate().skip(start_ordinal as usize) {
if stop.load(Ordering::Relaxed) {
return;
}
let ordinal = ordinal as u32;
let pinned_etag = if ordinal == start_ordinal && resume_etag.is_some() {
resume_etag.clone()
} else {
entry.etag.clone()
};
if pause_gate(&pause, Some(&stop), &tx).await.is_err() {
return;
}
let started = tx
.send(ChunkMsg::ObjectStart {
ordinal,
key: entry.key.clone(),
last_modified_ms: entry.last_modified_ms,
})
.await;
if started.is_err() {
return; }
match stream_object(
&split,
&store,
entry,
pinned_etag.as_deref(),
range_bytes,
chunk_bytes,
&tx,
&pause,
retry_base,
retries.as_ref(),
)
.await
{
Ok(true) => {}
Ok(false) => return, Err(e) => {
let _ = tx.send(ChunkMsg::LaneFailed(e)).await;
return;
}
}
if tx.send(ChunkMsg::ObjectEnd).await.is_err() {
return;
}
}
}
async fn pause_gate(
pause: &AtomicBool,
stop: Option<&AtomicBool>,
tx: &mpsc::Sender<ChunkMsg>,
) -> Result<(), ()> {
while pause.load(Ordering::Relaxed) {
if tx.is_closed() {
return Err(());
}
if let Some(stop) = stop
&& stop.load(Ordering::Relaxed)
{
return Err(());
}
tokio::time::sleep(PAUSE_POLL).await;
}
Ok(())
}
#[expect(
clippy::too_many_arguments,
reason = "internal seam between run_fetcher and the read loops"
)]
async fn stream_object(
split: &SplitId,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
pinned_etag: Option<&str>,
range_bytes: usize,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
) -> Result<bool, SplitFailure> {
match pinned_etag {
Some(etag) => {
stream_object_ranged(
split,
store,
entry,
etag,
range_bytes,
chunk_bytes,
tx,
pause,
retry_base,
retries,
)
.await
}
None => {
stream_object_streaming(
split,
store,
entry,
chunk_bytes,
tx,
pause,
retry_base,
retries,
)
.await
}
}
}
#[expect(
clippy::too_many_arguments,
reason = "internal read loop, mirrors stream_object_streaming"
)]
async fn stream_object_ranged(
split: &SplitId,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
etag: &str,
range_bytes: usize,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
) -> Result<bool, SplitFailure> {
let path = Path::from(entry.key.as_str());
let window = range_bytes.max(1) as u64;
let mut delivered: u64 = 0;
let mut attempt: u32 = 0;
while delivered < entry.size {
let end = delivered.saturating_add(window).min(entry.size);
let options = GetOptions {
if_match: Some(etag.to_owned()),
range: Some(GetRange::Bounded(delivered..end)),
..Default::default()
};
let buffered = match store.get_opts(&path, options).await {
Ok(result) => match result.bytes().await {
Ok(bytes) => bytes,
Err(e) => {
retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await?;
continue;
}
},
Err(e) => {
retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await?;
continue;
}
};
let read = buffered.len() as u64;
if read == 0 {
return Err(SplitFailure::Poison(
PoisonKind::Undecodable,
format!(
"split {split}: ranged read of \"{}\" at byte {delivered} returned no bytes",
entry.key
),
));
}
let mut bytes = buffered;
while !bytes.is_empty() {
let take = bytes.len().min(chunk_bytes);
let chunk = bytes.split_to(take);
if pause_gate(pause, None, tx).await.is_err() {
return Ok(false);
}
if tx.send(ChunkMsg::Chunk(chunk)).await.is_err() {
return Ok(false);
}
}
delivered += read;
attempt = 0;
}
Ok(true)
}
#[expect(
clippy::too_many_arguments,
reason = "internal read loop, mirrors stream_object_ranged"
)]
async fn stream_object_streaming(
split: &SplitId,
store: &Arc<dyn ObjectStore>,
entry: &ObjectEntry,
chunk_bytes: usize,
tx: &mpsc::Sender<ChunkMsg>,
pause: &AtomicBool,
retry_base: Duration,
retries: Option<&spate_core::metrics::Counter>,
) -> Result<bool, SplitFailure> {
let path = Path::from(entry.key.as_str());
let mut delivered: u64 = 0;
let mut attempt: u32 = 0;
'attempts: loop {
if delivered >= entry.size {
return Ok(true);
}
let options = GetOptions::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, None, 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}: {e}",
entry.key
),
}));
}
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: {e}",
entry.key
),
));
}
retry_or_fail(
split,
entry,
&e,
delivered,
&mut attempt,
retry_base,
retries,
)
.await?;
continue 'attempts;
}
None => return Ok(true),
}
}
}
}
async fn retry_or_fail(
split: &SplitId,
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}: {e}",
entry.key
);
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 let Some(c) = retries {
c.increment(1);
}
if *attempt >= MAX_ATTEMPTS {
return Err(SplitFailure::Poison(
PoisonKind::RetriesExhausted,
format!(
"split {split}: reading \"{}\" still failing at byte {delivered} after \
{MAX_ATTEMPTS} attempts: {e}",
entry.key
),
));
}
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 = %e,
"transient object read failure; backing off and resuming with a ranged get"
);
tokio::time::sleep(backoff).await;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use object_store::memory::InMemory;
use object_store::{GetResult, GetResultPayload, ObjectStoreExt as _, PutPayload};
use std::fmt;
use std::sync::atomic::AtomicU32;
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,
cut_streams: AtomicU32,
cut_after: usize,
cut_fatal: bool,
gets: AtomicU32,
ranges: std::sync::Mutex<Vec<RangeKind>>,
bodies_drained: Arc<AtomicU32>,
}
impl FlakyStore {
fn new(inner: InMemory) -> FlakyStore {
FlakyStore {
inner,
fail_gets: AtomicU32::new(0),
cut_streams: AtomicU32::new(0),
cut_after: 0,
cut_fatal: false,
gets: AtomicU32::new(0),
ranges: std::sync::Mutex::new(Vec::new()),
bodies_drained: Arc::new(AtomicU32::new(0)),
}
}
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 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) -> object_store::Error {
if fatal {
object_store::Error::PermissionDenied {
path: "k".into(),
source: "injected permission loss".into(),
}
} else {
Self::generic("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> {
self.gets.fetch_add(1, Ordering::Relaxed);
self.ranges.lock().unwrap().push(range_kind(&options.range));
if self
.fail_gets
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
.is_ok()
{
return Err(Self::generic("injected get failure"));
}
let result = self.inner.get_opts(location, options).await?;
if self
.cut_streams
.fetch_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 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))]
} 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)));
}
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: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal,
resume_etag: None,
chunk_bytes,
range_bytes,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
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
}
#[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 store: Arc<dyn ObjectStore> =
Arc::new(seeded(&[("p/a", b"first"), ("p/b", b"second")]).await);
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())]
);
}
#[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: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
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: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
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]
async fn pause_halts_fetching_and_resume_continues() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"abcdefgh")]).await);
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(1); let pause = Arc::new(AtomicBool::new(true));
let params = FetcherParams {
split: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 2,
range_bytes: 64,
tx,
pause: Arc::clone(&pause),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
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"
);
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: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 8,
range_bytes: 4096,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
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"
);
}
#[tokio::test]
async fn a_backpressured_lane_buffers_one_window_and_fetches_no_further() {
let flaky = Arc::new(FlakyStore::new(
seeded(&[("p/a", b"0123456789abcdef")]).await,
));
let store: Arc<dyn ObjectStore> = flaky.clone();
let slice = listed(&store).await;
let (tx, _rx) = mpsc::channel(1); let params = FetcherParams {
split: SplitId::new("s3-test").unwrap(),
store: Arc::clone(&store),
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 4,
range_bytes: 8,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(false)),
retry_base: Duration::from_millis(1),
retries: None,
};
let task = tokio::spawn(run_fetcher(params));
tokio::time::sleep(Duration::from_millis(80)).await;
assert_eq!(
flaky.recorded_ranges(),
vec![RangeKind::Bounded(0, 8)],
"one window fetched: the second is not started while the \
hand-off is blocked"
);
assert_eq!(
flaky.bodies_drained(),
1,
"the first window's body must be fully drained (connection \
released) before the fetcher blocks on the hand-off"
);
task.abort();
}
async fn no_etag_mid_stream_failure(cut_fatal: bool) -> SplitFailure {
let flaky = FlakyStore {
cut_after: 4,
cut_fatal,
..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]
async fn a_preset_stop_ends_the_fetcher_before_reading_anything() {
let store: Arc<dyn ObjectStore> =
Arc::new(seeded(&[("p/a", b"aaaa"), ("p/b", b"bbbb")]).await);
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(64);
let params = FetcherParams {
split: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::new(AtomicBool::new(false)),
stop: Arc::new(AtomicBool::new(true)),
retry_base: Duration::from_millis(1),
retries: None,
};
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();
assert!(
msgs.is_empty(),
"a fetcher that starts stopped reads nothing: {msgs:?}"
);
}
#[tokio::test]
async fn a_paused_fetcher_notices_the_stop_and_closes_cleanly() {
let store: Arc<dyn ObjectStore> = Arc::new(seeded(&[("p/a", b"aaaa")]).await);
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(8);
let pause = Arc::new(AtomicBool::new(true));
let stop = Arc::new(AtomicBool::new(false));
let params = FetcherParams {
split: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 64,
range_bytes: 64,
tx,
pause: Arc::clone(&pause),
stop: Arc::clone(&stop),
retry_base: Duration::from_millis(1),
retries: None,
};
let task = tokio::spawn(run_fetcher(params));
tokio::time::sleep(Duration::from_millis(80)).await;
assert!(rx.try_recv().is_err(), "a paused fetcher delivers nothing");
stop.store(true, Ordering::Relaxed);
let mut msgs = Vec::new();
let drained = tokio::time::timeout(Duration::from_secs(5), async {
while let Some(m) = rx.recv().await {
msgs.push(m);
}
})
.await;
assert!(
drained.is_ok(),
"a paused fetcher must notice the stop promptly"
);
task.await.unwrap();
assert!(msgs.is_empty(), "no object was ever started: {msgs:?}");
}
#[tokio::test]
async fn a_stop_finishes_the_open_object_then_cuts_at_the_boundary() {
let store: Arc<dyn ObjectStore> =
Arc::new(seeded(&[("p/a", b"0123456789"), ("p/b", b"bbbb")]).await);
let slice = listed(&store).await;
let (tx, mut rx) = mpsc::channel(1); let pause = Arc::new(AtomicBool::new(false));
let stop = Arc::new(AtomicBool::new(false));
let params = FetcherParams {
split: SplitId::new("s3-test").unwrap(),
store,
slice: Arc::new(slice),
start_ordinal: 0,
resume_etag: None,
chunk_bytes: 2,
range_bytes: 64,
tx,
pause: Arc::clone(&pause),
stop: Arc::clone(&stop),
retry_base: Duration::from_millis(1),
retries: None,
};
let task = tokio::spawn(run_fetcher(params));
let m0 = rx.recv().await.unwrap();
assert!(matches!(m0, ChunkMsg::ObjectStart { ordinal: 0, .. }));
let m1 = rx.recv().await.unwrap();
assert!(matches!(m1, ChunkMsg::Chunk(_)), "mid the first object");
pause.store(true, Ordering::Relaxed);
stop.store(true, Ordering::Relaxed);
tokio::time::sleep(Duration::from_millis(40)).await;
pause.store(false, Ordering::Relaxed); let mut msgs = vec![m0, m1];
while let Some(m) = rx.recv().await {
msgs.push(m);
}
task.await.unwrap();
assert!(
!msgs.iter().any(|m| matches!(m, ChunkMsg::LaneFailed(_))),
"a boundary stop is a clean channel close, never a failure"
);
assert_eq!(
assembled(&msgs),
vec![("p/a".to_string(), b"0123456789".to_vec())],
"the open object finishes; the boundary stop precedes the next object"
);
}
}