use std::sync::{Arc, Mutex};
use std::time::Duration;
use super::DownloadTarget;
use super::Downloader;
use crate::source::{BlobSource, FakeFunder, Recovery, ScriptedSource};
use crate::{
ClientRangedStore, Cumulative, GaveUp, PoolContext, PoolLedger, ProgressClock, StaticSources,
StopPolicy, StreamCandidate,
};
use alloy::primitives::{Address, B256, U256};
use alloy::signers::local::PrivateKeySigner;
fn ctx_for(provider: u8) -> PoolContext {
PoolContext {
pool_id: B256::ZERO,
provider: Address::repeat_byte(provider),
deposit: U256::from(u128::MAX),
client_signer: Arc::new(PrivateKeySigner::random()),
voucher_domain: decdn_incentive::bind_node_id_domain(1, Address::ZERO),
prior_bytes_delivered: U256::ZERO,
prior_amount: U256::ZERO,
client_binding: None,
capability: None,
}
}
fn candidate<S>(source: S, ledger: Arc<PoolLedger>, provider: u8) -> StreamCandidate<S> {
StreamCandidate {
source,
ctx: Arc::new(Mutex::new(ctx_for(provider))),
ledger,
coverage: None,
lease: crate::LaneLease::default(),
widen: None,
}
}
fn funder() -> FakeFunder {
FakeFunder::new(Recovery::ToppedUp(U256::from(u128::MAX)))
}
fn downloader<S: BlobSource>(
candidates: Vec<StreamCandidate<S>>,
) -> anyhow::Result<Downloader<StaticSources<S>, FakeFunder>> {
let sources = StaticSources::new(candidates)?;
let holders = sources.holders();
Ok(Downloader::new(sources, funder()).holders(holders))
}
#[tokio::test]
async fn a_downloader_streams_at_most_max_lanes() -> anyhow::Result<()> {
let blob = payload(4 * 1024 * 1024);
let ledger_a = Arc::new(PoolLedger::new(Cumulative::default()));
let ledger_b = Arc::new(PoolLedger::new(Cumulative::default()));
let src_a = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger_a));
let src_b = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger_b));
let (root, total) = (src_a.root(), src_a.total_bytes());
let (probe_a, probe_b) = (src_a.clone(), src_b.clone());
let sources = StaticSources::new(vec![
candidate(src_a, ledger_a, 0xA1),
candidate(src_b, ledger_b, 0xB2),
])?;
let holders = sources.holders();
let downloader = Downloader::new(sources, funder())
.holders(holders)
.max_lanes(1);
let dir = tempfile::tempdir()?;
downloader
.fetch_to_dir(&[(root, total)], dir.path(), None)
.await?;
anyhow::ensure!(
(probe_a.delivered_bytes() > 0) != (probe_b.delivered_bytes() > 0),
"exactly one lane streamed (a={}, b={})",
probe_a.delivered_bytes(),
probe_b.delivered_bytes()
);
Ok(())
}
fn payload(len: usize) -> Vec<u8> {
let mut out = vec![0u8; len];
let mut x: u32 = 0x1234_5678;
for b in &mut out {
x ^= x << 13;
x ^= x >> 17;
x ^= x << 5;
*b = x.to_le_bytes()[0];
}
out
}
#[tokio::test]
async fn fetch_to_dir_writes_a_blake3_identical_file() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let paths = downloader
.fetch_to_dir(&[(root, total)], dir.path(), None)
.await?;
anyhow::ensure!(
paths.len() == 1,
"one entry → one file, got {}",
paths.len()
);
let path = paths
.first()
.ok_or_else(|| anyhow::anyhow!("no path returned"))?;
let got = std::fs::read(path)?;
anyhow::ensure!(
got == blob,
"written file must be byte-identical to the blob"
);
anyhow::ensure!(
blake3::hash(&got).as_bytes() == &root,
"written file must be BLAKE3-identical to the entry hash"
);
anyhow::ensure!(
path.file_name().and_then(|n| n.to_str())
== Some(blake3::Hash::from_bytes(root).to_hex().as_str()),
"file is named by its content-address hex"
);
Ok(())
}
#[tokio::test]
async fn fetch_to_dir_reports_progress_up_to_the_blobs_total() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let max_pos = Arc::new(std::sync::atomic::AtomicU64::new(0));
let saw_blob_total = Arc::new(std::sync::atomic::AtomicBool::new(false));
let mp = Arc::clone(&max_pos);
let sbt = Arc::clone(&saw_blob_total);
let on_progress = move |pos: u64, tot: u64| {
mp.fetch_max(pos, std::sync::atomic::Ordering::SeqCst);
if tot == total {
sbt.store(true, std::sync::atomic::Ordering::SeqCst);
}
};
downloader
.fetch_to_dir(&[(root, total)], dir.path(), Some(&on_progress))
.await?;
anyhow::ensure!(
saw_blob_total.load(std::sync::atomic::Ordering::SeqCst),
"progress must report the blob's content total in the total slot"
);
anyhow::ensure!(
max_pos.load(std::sync::atomic::Ordering::SeqCst) == total,
"final progress must reach the blob's total content bytes"
);
Ok(())
}
#[tokio::test]
async fn fetch_to_paths_writes_each_blob_to_its_named_dest() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("my-model.safetensors");
let paths = downloader
.fetch_to_paths(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &dest,
ranges: None,
}],
None,
)
.await?;
anyhow::ensure!(
paths == vec![dest.clone()],
"returns the caller's destination path, not a hex name"
);
let got = std::fs::read(&dest)?;
anyhow::ensure!(got == blob, "written file is byte-identical to the blob");
anyhow::ensure!(
blake3::hash(&got).as_bytes() == &root,
"written file is BLAKE3-identical to the entry hash"
);
Ok(())
}
#[tokio::test]
async fn a_ranged_target_is_promoted_only_when_its_store_completes() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let total = u64::try_from(blob.len())?;
let dir = tempfile::tempdir()?;
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let part = dir.path().join("part.bin");
let half = [(0, total / 2)];
let paths = downloader(vec![candidate(source, ledger, 0xA1)])?
.fetch_to_paths(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &part,
ranges: Some(&half),
}],
None,
)
.await?;
anyhow::ensure!(paths.is_empty(), "an incomplete target is not returned");
anyhow::ensure!(!part.exists(), "an incomplete target is not promoted");
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let whole = dir.path().join("whole.bin");
let all = [(0, total)];
let paths = downloader(vec![candidate(source, ledger, 0xA1)])?
.fetch_to_paths(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &whole,
ranges: Some(&all),
}],
None,
)
.await?;
anyhow::ensure!(
paths == vec![whole.clone()],
"the promoted dest is returned"
);
anyhow::ensure!(
std::fs::read(&whole)? == blob,
"the promoted file is the blob"
);
Ok(())
}
#[tokio::test]
async fn fetch_to_paths_threads_a_shared_ledger_registry() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("shared-ledger.bin");
let registry = crate::LaneLedgers::new();
let stop = StopPolicy::new(true, None, Arc::new(ProgressClock::new()));
let paths = downloader
.fetch_to_paths_shared(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &dest,
ranges: None,
}],
®istry,
None,
&stop,
)
.await?;
anyhow::ensure!(paths == vec![dest.clone()]);
anyhow::ensure!(std::fs::read(&dest)? == blob);
Ok(())
}
#[tokio::test]
async fn two_holders_stripe_a_blake3_identical_file() -> anyhow::Result<()> {
let blob = payload(4 * 1024 * 1024);
let ledger_a = Arc::new(PoolLedger::new(Cumulative::default()));
let ledger_b = Arc::new(PoolLedger::new(Cumulative::default()));
let src_a = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger_a));
let src_b = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger_b));
let root = src_a.root();
let total = src_a.total_bytes();
let probe_a = src_a.clone();
let probe_b = src_b.clone();
let downloader = downloader(vec![
candidate(src_a, ledger_a, 0xA1),
candidate(src_b, ledger_b, 0xB2),
])?;
let dir = tempfile::tempdir()?;
let paths = downloader
.fetch_to_dir(&[(root, total)], dir.path(), None)
.await?;
let path = paths
.first()
.ok_or_else(|| anyhow::anyhow!("no path returned"))?;
anyhow::ensure!(
std::fs::read(path)? == blob,
"striped file must be identical"
);
anyhow::ensure!(
probe_a.delivered_bytes() > 0 && probe_b.delivered_bytes() > 0,
"both holders must have contributed (a={}, b={})",
probe_a.delivered_bytes(),
probe_b.delivered_bytes()
);
Ok(())
}
#[tokio::test]
async fn a_pre_seeded_prefix_is_not_refetched() -> anyhow::Result<()> {
let blob = payload(2 * 1024 * 1024);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = source.total_bytes();
let probe = source.clone();
let dir = tempfile::tempdir()?;
let stem = blake3::Hash::from_bytes(root).to_hex();
let seeded_prefix = total / 2;
ClientRangedStore::seed_checkpointed_prefix(dir.path(), stem.as_str(), &blob, seeded_prefix)?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let paths = downloader
.fetch_to_dir(&[(root, total)], dir.path(), None)
.await?;
let path = paths
.first()
.ok_or_else(|| anyhow::anyhow!("no path returned"))?;
anyhow::ensure!(
std::fs::read(path)? == blob,
"resumed file must be identical"
);
let opened = probe.opened_ranges();
anyhow::ensure!(
opened.iter().all(|&(start, _)| start >= seeded_prefix),
"the seeded prefix must not be re-opened, opened={opened:?}"
);
anyhow::ensure!(
probe.opened_bytes() < total,
"a dedup'd download must fetch fewer than the whole blob's bytes ({} of {total})",
probe.opened_bytes()
);
Ok(())
}
async fn fetch_ranges(
blob: &[u8],
dest: &std::path::Path,
claim: u64,
ranges: Option<&[(u64, u64)]>,
) -> anyhow::Result<(ScriptedSource, Vec<std::path::PathBuf>)> {
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.to_vec())?.paying(Arc::clone(&ledger));
let probe = source.clone();
let paths = downloader(vec![candidate(source, ledger, 0xA1)])?
.fetch_to_paths(
&[DownloadTarget {
hash: probe.root(),
total_bytes: claim,
dest,
ranges,
}],
None,
)
.await?;
Ok((probe, paths))
}
#[tokio::test]
async fn a_resumed_record_with_a_bound_past_the_claim_finishes() -> anyhow::Result<()> {
let blob = payload(3 * 1024 * 1024);
let total = u64::try_from(blob.len())?;
let held = 1024 * 1024;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let (_, paths) = fetch_ranges(&blob, &dest, total, Some(&[(0, held)])).await?;
anyhow::ensure!(paths.is_empty(), "the first run leaves a partial");
let (probe, paths) = fetch_ranges(&blob, &dest, held, None).await?;
anyhow::ensure!(paths == vec![dest.clone()], "the rerun promotes");
anyhow::ensure!(std::fs::read(&dest)? == blob, "the file is the blob");
let opened = probe.opened_ranges();
anyhow::ensure!(
opened.iter().all(|&(start, _)| start >= held),
"no held byte is opened again: {opened:?}"
);
Ok(())
}
#[tokio::test]
async fn a_resumed_record_with_a_proven_size_past_the_claim_finishes() -> anyhow::Result<()> {
let blob = payload(3 * 1024 * 1024);
let total = u64::try_from(blob.len())?;
let held = 1024 * 1024;
let group = decdn_bao_range::CHUNK_GROUP_BYTES;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let first = [(0, held), (total - group, group)];
let (_, paths) = fetch_ranges(&blob, &dest, total, Some(&first)).await?;
anyhow::ensure!(paths.is_empty(), "the first run leaves a partial");
let (probe, paths) = fetch_ranges(&blob, &dest, held, None).await?;
anyhow::ensure!(paths == vec![dest.clone()], "the rerun promotes");
anyhow::ensure!(std::fs::read(&dest)? == blob, "the file is the blob");
let opened = probe.opened_ranges();
anyhow::ensure!(
opened
.iter()
.all(|&(start, len)| start >= held && start + len <= total - group),
"only the missing middle is opened: {opened:?}"
);
Ok(())
}
#[tokio::test]
async fn a_finalize_hash_mismatch_fetches_again_and_completes() -> anyhow::Result<()> {
let blob = payload(2 * 1024 * 1024);
let total = u64::try_from(blob.len())?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let mut drifted = blob.clone();
if let Some(b) = drifted.get_mut(1000) {
*b ^= 0xFF;
}
ClientRangedStore::seed_checkpointed_prefix(dir.path(), "model.bin", &drifted, total)?;
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let probe = source.clone();
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let target = DownloadTarget {
hash: probe.root(),
total_bytes: total,
dest: &dest,
ranges: None,
};
let paths = downloader.fetch_to_paths(&[target], None).await?;
anyhow::ensure!(paths == vec![dest.clone()], "the fetch promotes");
anyhow::ensure!(std::fs::read(&dest)? == blob, "the file is the blob");
anyhow::ensure!(
probe.opened_bytes() >= total,
"the second pass fetched the blob again: {}",
probe.opened_bytes()
);
anyhow::ensure!(
downloader.refetched_targets() == 1,
"the downloader reports the target it fetched again"
);
Ok(())
}
#[tokio::test(start_paused = true)]
async fn a_lone_stalled_source_recovers_after_it_cools() -> anyhow::Result<()> {
let blob = payload(12 * 1024 * 1024);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?
.stall_after(6 * 1024 * 1024, Duration::from_hours(1))
.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let stop = StopPolicy::new(false, None, Arc::new(ProgressClock::new()));
downloader
.fetch_to_paths_until(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &dest,
ranges: None,
}],
None,
&stop,
)
.await?;
anyhow::ensure!(std::fs::read(&dest)? == blob, "the file must be identical");
Ok(())
}
#[tokio::test(start_paused = true)]
async fn a_download_with_no_progress_gives_up() -> anyhow::Result<()> {
let blob = payload(2 * 1024 * 1024);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?
.paying(Arc::clone(&ledger))
.with_fault_after(0, || anyhow::anyhow!("connection reset"));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let limit = Duration::from_secs(30);
let stop = StopPolicy::new(false, Some(limit), Arc::new(ProgressClock::new()));
let fetched = downloader
.fetch_to_paths_until(
&[DownloadTarget {
hash: root,
total_bytes: total,
dest: &dest,
ranges: None,
}],
None,
&stop,
)
.await;
let Err(err) = fetched else {
anyhow::bail!("a download with no progress must give up");
};
anyhow::ensure!(
err.downcast_ref::<GaveUp>() == Some(&GaveUp { idle: limit }),
"{err:#}"
);
anyhow::ensure!(!dest.exists(), "a stopped target is not promoted");
Ok(())
}
#[tokio::test(start_paused = true)]
async fn a_dropped_download_records_the_landed_prefix() -> anyhow::Result<()> {
const STALL: u64 = 6 * 1024 * 1024;
let blob = payload(12 * 1024 * 1024);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?
.stall_after(STALL, Duration::from_hours(1))
.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let downloader = downloader(vec![candidate(source, ledger, 0xA1)])?;
let dir = tempfile::tempdir()?;
let dest = dir.path().join("model.bin");
let targets = [DownloadTarget {
hash: root,
total_bytes: total,
dest: &dest,
ranges: None,
}];
let dropped = tokio::time::timeout(
crate::driver::PRESENT_RECORD_FLUSH_INTERVAL / 5,
downloader.fetch_to_paths(&targets, None),
)
.await;
anyhow::ensure!(
dropped.is_err(),
"the stalled download must still be running"
);
let partial = ClientRangedStore::open(dir.path(), "model.bin", root)?;
let recorded = crate::driver::ranges_content_len(
&decdn_bao_range::RangedStore::present_ranges(&partial).await?,
total,
);
let landed = ClientRangedStore::checkpointed_len(STALL);
anyhow::ensure!(
recorded >= landed && recorded < total,
"the landed prefix is recorded for a resume: {recorded}, want at least {landed}"
);
Ok(())
}
#[tokio::test]
async fn empty_static_sources_are_a_clear_error() -> anyhow::Result<()> {
let Err(err) = downloader::<ScriptedSource>(Vec::new()) else {
anyhow::bail!("empty static sources must be rejected");
};
anyhow::ensure!(
err.to_string().contains("at least one candidate"),
"the error must name the empty candidate set, got: {err}"
);
Ok(())
}
#[tokio::test(start_paused = true)]
async fn a_downloader_with_no_holder_discovers_them() -> anyhow::Result<()> {
let blob = payload(1_500_000);
let ledger = Arc::new(PoolLedger::new(Cumulative::default()));
let source = ScriptedSource::new(blob.clone())?.paying(Arc::clone(&ledger));
let root = source.root();
let total = u64::try_from(blob.len())?;
let sources = StaticSources::new(vec![candidate(source, ledger, 0xA1)])?;
let downloader = Downloader::new(sources, funder());
let dir = tempfile::tempdir()?;
let paths = downloader
.fetch_to_dir(&[(root, total)], dir.path(), None)
.await?;
let path = paths
.first()
.ok_or_else(|| anyhow::anyhow!("no path returned"))?;
anyhow::ensure!(std::fs::read(path)? == blob);
Ok(())
}