use std::sync::atomic::AtomicU64;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use alloy::dyn_abi::Eip712Domain;
use alloy::primitives::{Address, U256};
use decdn_bao_range::AlignedRange;
use decdn_incentive::PoolId;
use iroh::{Endpoint, EndpointAddr};
use crate::sink::{PullReader, StashedFault};
use crate::{
Connections, PoolContext, PoolLedger, PullDeadlines, UpstreamPull, UpstreamPullHeader,
VoucherProgress,
};
pub type SourceFuture<'a, T> =
core::pin::Pin<Box<dyn core::future::Future<Output = anyhow::Result<T>> + Send + 'a>>;
#[doc(hidden)]
pub type SourceStream<'a, T> = core::pin::Pin<Box<dyn futures_util::Stream<Item = T> + Send + 'a>>;
#[doc(hidden)]
pub trait BaoRangeReader: iroh_io::AsyncStreamReader + StashedFault + Send {}
impl<T> BaoRangeReader for T where T: iroh_io::AsyncStreamReader + StashedFault + Send {}
#[doc(hidden)]
pub trait BlobSource: Send + Sync {
type Reader: BaoRangeReader;
fn open(
&self,
hash: [u8; 32],
range: AlignedRange,
) -> SourceFuture<'_, (UpstreamPullHeader, Self::Reader)>;
fn finish(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress>;
fn stop(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress>;
fn max_blob_size_bytes(&self) -> u64 {
0
}
}
#[doc(hidden)]
pub type IngestFuture<'a, R> =
core::pin::Pin<Box<dyn core::future::Future<Output = anyhow::Result<(R, IngestEnd)>> + 'a>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[doc(hidden)]
pub enum IngestEnd {
Drained,
Stopped,
}
#[doc(hidden)]
pub trait IngestStore: decdn_bao_range::RangedStore {
fn ingest_stream<'a, R>(
&'a self,
range: &'a AlignedRange,
reader: R,
on_progress: Option<&'a (dyn Fn(u64) + Send + Sync)>,
claimed_total: u64,
stop_at: Option<&'a AtomicU64>,
) -> IngestFuture<'a, R>
where
R: BaoRangeReader + 'a;
fn flush_present_record(&self) -> SourceFuture<'_, ()>;
fn proven(&self) -> Option<u64> {
Some(self.total_bytes())
}
fn set_bound(&self, bound: u64) {
let _ = bound;
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PoolReplaced {
pub closed: PoolId,
pub opened: PoolId,
}
impl std::fmt::Display for PoolReplaced {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"pool {} no longer accepts funds, so new pool {} replaces it; the remaining work \
runs again against the new pool",
self.closed, self.opened
)
}
}
impl std::error::Error for PoolReplaced {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Recovery {
ToppedUp(U256),
Replaced(PoolReplaced),
Unavailable,
}
pub trait Funder: Send + Sync {
fn recover(&self, remaining: U256) -> SourceFuture<'_, Recovery>;
fn pool_spent(&self) -> Option<U256> {
None
}
}
fn micros_now() -> u64 {
u64::try_from(
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| d.as_micros()),
)
.unwrap_or(u64::MAX)
}
pub struct PeerSource<'a> {
endpoint: &'a Endpoint,
target: EndpointAddr,
ctx: Arc<Mutex<PoolContext>>,
ledger: Arc<PoolLedger>,
slash_domain: &'a Eip712Domain,
expected_signer: Address,
namespace_id: [u8; 32],
max_blob_size_bytes: u64,
max_rate_per_mb: u64,
deadlines: PullDeadlines,
dial_runtime: Option<tokio::runtime::Handle>,
connections: Option<Connections>,
}
impl std::fmt::Debug for PeerSource<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PeerSource")
.field("target", &self.target)
.field("expected_signer", &self.expected_signer)
.field("namespace_id", &self.namespace_id)
.field("max_blob_size_bytes", &self.max_blob_size_bytes)
.field("max_rate_per_mb", &self.max_rate_per_mb)
.field("deadlines", &self.deadlines)
.finish_non_exhaustive()
}
}
impl<'a> PeerSource<'a> {
#[must_use]
#[allow(clippy::too_many_arguments)]
pub const fn new(
endpoint: &'a Endpoint,
target: EndpointAddr,
ctx: Arc<Mutex<PoolContext>>,
ledger: Arc<PoolLedger>,
slash_domain: &'a Eip712Domain,
expected_signer: Address,
namespace_id: [u8; 32],
max_blob_size_bytes: u64,
max_rate_per_mb: u64,
deadlines: PullDeadlines,
connections: Option<Connections>,
) -> Self {
Self {
endpoint,
target,
ctx,
ledger,
slash_domain,
expected_signer,
namespace_id,
max_blob_size_bytes,
max_rate_per_mb,
deadlines,
dial_runtime: None,
connections,
}
}
#[must_use]
#[doc(hidden)]
pub fn with_dial_runtime(mut self, runtime: tokio::runtime::Handle) -> Self {
self.dial_runtime = Some(runtime);
self
}
#[doc(hidden)]
pub fn open_whole(&self, hash: [u8; 32]) -> SourceFuture<'_, (UpstreamPullHeader, PullReader)> {
Box::pin(async move {
let pull = self.open_pull(hash, 0, 0).await?;
Ok((pull.0, PullReader::new(pull.1)))
})
}
async fn open_pull(
&self,
hash: [u8; 32],
byte_offset: u64,
byte_len: u64,
) -> anyhow::Result<(UpstreamPullHeader, UpstreamPull)> {
let ctx = {
self.ctx
.lock()
.map_err(|_| anyhow::anyhow!("pool context lock poisoned"))?
.clone()
};
if let Some(connections) = &self.connections {
self.open_shared(connections, &ctx, hash, byte_offset, byte_len)
.await
} else {
crate::open_progressive_pull(
self.endpoint,
self.target.clone(),
&ctx,
Arc::clone(&self.ledger),
self.slash_domain,
self.expected_signer,
hash,
self.namespace_id,
byte_offset,
micros_now(),
self.max_blob_size_bytes,
self.max_rate_per_mb,
self.deadlines,
byte_len,
self.dial_runtime.as_ref(),
)
.await
}
}
async fn open_shared(
&self,
connections: &Connections,
ctx: &PoolContext,
hash: [u8; 32],
byte_offset: u64,
byte_len: u64,
) -> anyhow::Result<(UpstreamPullHeader, UpstreamPull)> {
let open = self.deadlines.open;
let warm = tokio::time::timeout(open, connections.get(&self.target))
.await
.map_err(|_| anyhow::Error::new(crate::PullTimeout { after: open }))??;
let conn = warm.connection().clone();
crate::open_progressive_pull_on(
&conn,
crate::connection::ConnOwner::Shared {
warm,
map: connections.clone(),
},
ctx,
Arc::clone(&self.ledger),
self.slash_domain,
self.expected_signer,
hash,
self.namespace_id,
byte_offset,
micros_now(),
self.max_blob_size_bytes,
self.max_rate_per_mb,
self.deadlines,
byte_len,
)
.await
}
}
impl BlobSource for PeerSource<'_> {
type Reader = PullReader;
fn open(
&self,
hash: [u8; 32],
range: AlignedRange,
) -> SourceFuture<'_, (UpstreamPullHeader, Self::Reader)> {
Box::pin(async move {
tracing::debug!(
peer = %self.target.id,
hash = %blake3::Hash::from_bytes(hash).to_hex(),
byte_offset = range.fetch_start(),
byte_len = range.fetch_len(),
"opening a paid leg"
);
let (header, pull) = self
.open_pull(hash, range.fetch_start(), range.fetch_len())
.await?;
Ok((header, PullReader::new(pull)))
})
}
fn finish(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
Box::pin(async move { reader.into_inner().finish().await })
}
fn stop(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
Box::pin(async move { reader.into_inner().stop().await })
}
fn max_blob_size_bytes(&self) -> u64 {
self.max_blob_size_bytes
}
}
pub(crate) const PRIMED_MAX_IDLE: Duration = Duration::from_secs(2);
struct Primed<R> {
hash: [u8; 32],
range: AlignedRange,
header: UpstreamPullHeader,
reader: R,
at: tokio::time::Instant,
}
#[doc(hidden)]
pub struct PrimedSource<S: BlobSource> {
inner: S,
primed: Mutex<Option<Primed<S::Reader>>>,
}
impl<S: BlobSource> std::fmt::Debug for PrimedSource<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let primed = self.slot().is_some();
f.debug_struct("PrimedSource")
.field("primed", &primed)
.finish_non_exhaustive()
}
}
impl<S: BlobSource> PrimedSource<S> {
fn slot(&self) -> std::sync::MutexGuard<'_, Option<Primed<S::Reader>>> {
self.primed
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[must_use]
pub const fn new(inner: S) -> Self {
Self {
inner,
primed: Mutex::new(None),
}
}
pub fn prime(
&self,
hash: [u8; 32],
range: AlignedRange,
header: UpstreamPullHeader,
reader: S::Reader,
opened_at: tokio::time::Instant,
) {
let parked = Primed {
hash,
range,
header,
reader,
at: opened_at,
};
let replaced = self.slot().replace(parked);
if replaced.is_some() {
tracing::debug!("a primed pull no open took is replaced and closed");
}
}
pub fn clear(&self) {
if self.slot().take().is_some() {
tracing::debug!("a primed pull no open took is closed");
}
}
pub const fn inner(&self) -> &S {
&self.inner
}
fn take(
&self,
hash: [u8; 32],
range: &AlignedRange,
) -> Option<(UpstreamPullHeader, S::Reader)> {
let mut slot = self.slot();
let parked = slot.take()?;
if parked.at.elapsed() > PRIMED_MAX_IDLE {
tracing::debug!("a primed pull went stale before its open and is closed");
return None;
}
if parked.hash == hash && parked.range == *range {
return Some((parked.header, parked.reader));
}
*slot = Some(parked);
None
}
}
impl<S: BlobSource> BlobSource for PrimedSource<S> {
type Reader = S::Reader;
fn open(
&self,
hash: [u8; 32],
range: AlignedRange,
) -> SourceFuture<'_, (UpstreamPullHeader, Self::Reader)> {
match self.take(hash, &range) {
Some(adopted) => Box::pin(async move { Ok(adopted) }),
None => self.inner.open(hash, range),
}
}
fn finish(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
self.inner.finish(reader)
}
fn stop(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
self.inner.stop(reader)
}
fn max_blob_size_bytes(&self) -> u64 {
self.inner.max_blob_size_bytes()
}
}
#[cfg(any(test, feature = "test-util"))]
mod doubles {
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use super::{Recovery, SourceFuture, StashedFault};
use crate::{UpstreamPullHeader, VoucherProgress};
use alloy::primitives::U256;
use bao_tree::io::outboard::PreOrderMemOutboard;
use bytes::Bytes;
use decdn_bao_range::{AlignedRange, IROH_BLOCK_SIZE, encode_verified_range};
const SCRIPTED_RATE_PER_MB: u64 = 1;
const SCRIPTED_INTERVAL_BYTES: u64 = 1024 * 1024;
#[cfg(test)]
pub(crate) fn ctx_with(provider: u8, deposit: U256) -> crate::PoolContext {
use alloy::primitives::{Address, B256};
use alloy::signers::local::PrivateKeySigner;
let signer = PrivateKeySigner::random();
crate::PoolContext {
pool_id: B256::ZERO,
provider: Address::repeat_byte(provider),
deposit,
client_signer: Arc::new(signer),
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,
}
}
type FaultFn = Arc<dyn Fn() -> anyhow::Error + Send + Sync>;
type LegTimes = (u64, tokio::time::Instant, Option<tokio::time::Instant>);
#[derive(Clone)]
pub struct ScriptedSource {
root: [u8; 32],
blob: Bytes,
outboard: Bytes,
fault: Option<(usize, FaultFn)>,
faults_left: Option<Arc<AtomicU32>>,
signed_size: Option<u64>,
refuse_open: Option<(std::ops::Range<usize>, FaultFn)>,
refuse_blocks: Option<(Vec<u32>, FaultFn)>,
fault_gates: Vec<(usize, tokio::sync::watch::Receiver<bool>)>,
timeline: Arc<Mutex<Vec<LegTimes>>>,
opened: Arc<Mutex<Vec<(u64, u64)>>>,
delivered: Arc<AtomicU64>,
gate: Option<tokio::sync::watch::Receiver<bool>>,
first_read_stall: Option<Duration>,
finish_stall: Option<Duration>,
stall_after: Option<(u64, Duration)>,
ledger: Option<Arc<crate::PoolLedger>>,
stops: Arc<AtomicU32>,
read_delay: Option<Duration>,
}
impl std::fmt::Debug for ScriptedSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ScriptedSource")
.field("root", &blake3::Hash::from_bytes(self.root))
.field("total_bytes", &self.total_bytes())
.field("fault_after", &self.fault.as_ref().map(|(after, _)| *after))
.field("opened_ranges", &self.opened_ranges())
.finish_non_exhaustive()
}
}
impl ScriptedSource {
pub fn new(blob: impl Into<Bytes>) -> anyhow::Result<Self> {
let blob = blob.into();
let ob = PreOrderMemOutboard::create(&blob, IROH_BLOCK_SIZE);
Ok(Self {
root: *ob.root.as_bytes(),
blob,
outboard: ob.data.into(),
fault: None,
faults_left: None,
signed_size: None,
refuse_open: None,
refuse_blocks: None,
fault_gates: Vec::new(),
timeline: Arc::new(Mutex::new(Vec::new())),
opened: Arc::new(Mutex::new(Vec::new())),
delivered: Arc::new(AtomicU64::new(0)),
gate: None,
first_read_stall: None,
finish_stall: None,
stall_after: None,
ledger: None,
stops: Arc::new(AtomicU32::new(0)),
read_delay: None,
})
}
#[cfg(test)]
#[must_use]
pub(crate) const fn slow_finish(mut self, stall: Duration) -> Self {
self.finish_stall = Some(stall);
self
}
#[cfg(test)]
#[must_use]
pub(crate) const fn stall_after(mut self, after_bytes: u64, stall: Duration) -> Self {
self.stall_after = Some((after_bytes, stall));
self
}
#[cfg(test)]
#[must_use]
pub(crate) const fn throttled(mut self, per_read: Duration) -> Self {
self.read_delay = Some(per_read);
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn gated_on(mut self, gate: tokio::sync::watch::Receiver<bool>) -> Self {
self.gate = Some(gate);
self
}
#[cfg(test)]
#[must_use]
pub(crate) const fn slow_to_start(mut self, stall: Duration) -> Self {
self.first_read_stall = Some(stall);
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn delivered_bytes(&self) -> u64 {
self.delivered.load(Ordering::SeqCst)
}
#[must_use]
pub fn paying(mut self, ledger: Arc<crate::PoolLedger>) -> Self {
self.ledger = Some(ledger);
self
}
async fn settle(
&self,
reader: &ScriptedReader,
wire: u64,
) -> anyhow::Result<VoucherProgress> {
if let Some(leg) = reader.leg
&& let Ok(mut timeline) = self.timeline.lock()
&& let Some(entry) = timeline.get_mut(leg)
{
entry.2 = Some(tokio::time::Instant::now());
}
let Some(ledger) = &self.ledger else {
return Ok(VoucherProgress::default());
};
if wire > 0 {
ledger
.issue(
wire,
SCRIPTED_RATE_PER_MB,
crate::EpochAction::Keep,
|_next, _chain| async { Ok(()) },
)
.await?;
}
Ok(VoucherProgress::from_cumulative(
ledger.committed(),
U256::ZERO,
))
}
#[cfg(test)]
#[must_use]
pub(crate) fn stopped_pulls(&self) -> u32 {
self.stops.load(Ordering::SeqCst)
}
#[must_use]
pub const fn root(&self) -> [u8; 32] {
self.root
}
#[must_use]
pub const fn total_bytes(&self) -> u64 {
self.blob.len() as u64
}
#[must_use]
pub(crate) fn opened_ranges(&self) -> Vec<(u64, u64)> {
self.opened.lock().map(|o| o.clone()).unwrap_or_default()
}
#[cfg(test)]
#[must_use]
pub(crate) fn refusing_open(
mut self,
nth: usize,
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
self.refuse_open = Some((nth..nth.saturating_add(1), Arc::new(make)));
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn holding_fault(
mut self,
nth: usize,
gate: tokio::sync::watch::Receiver<bool>,
) -> Self {
self.fault_gates.push((nth, gate));
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn refusing_opens_from(
mut self,
first: usize,
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
self.refuse_open = Some((first..usize::MAX, Arc::new(make)));
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn refusing_blocks(
mut self,
blocks: &[u32],
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
self.refuse_blocks = Some((blocks.to_vec(), Arc::new(make)));
self
}
#[cfg(test)]
#[must_use]
pub(crate) fn timeline(&self) -> Vec<LegTimes> {
self.timeline.lock().map(|t| t.clone()).unwrap_or_default()
}
#[cfg(test)]
#[must_use]
pub(crate) fn opened_bytes(&self) -> u64 {
self.opened
.lock()
.map_or(0, |o| o.iter().map(|(_, len)| *len).sum())
}
#[must_use]
pub fn with_fault_after(
mut self,
wire_bytes: usize,
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
self.fault = Some((wire_bytes, Arc::new(make)));
self
}
#[must_use]
pub fn fault_once_after(
self,
wire_bytes: usize,
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
self.fault_times_after(1, wire_bytes, make)
}
#[must_use]
pub fn fault_times_after(
self,
times: u32,
wire_bytes: usize,
make: impl Fn() -> anyhow::Error + Send + Sync + 'static,
) -> Self {
let mut this = self.with_fault_after(wire_bytes, make);
this.faults_left = Some(Arc::new(AtomicU32::new(times)));
this
}
#[must_use]
pub const fn signing_size(mut self, total_bytes: u64) -> Self {
self.signed_size = Some(total_bytes);
self
}
fn wire_for(&self, range: &AlignedRange) -> anyhow::Result<Bytes> {
let s = usize::try_from(range.fetch_start())?;
let e = usize::try_from(range.fetch_end())?;
let data = self
.blob
.get(s..e)
.ok_or_else(|| anyhow::anyhow!("scripted range out of bounds"))?;
let combined = encode_verified_range(self.root, range, data, self.outboard.clone())?;
let wire = combined
.get(8..)
.ok_or_else(|| anyhow::anyhow!("combined wire shorter than its 8-byte header"))?;
Ok(Bytes::copy_from_slice(wire))
}
}
impl super::BlobSource for ScriptedSource {
type Reader = ScriptedReader;
fn open(
&self,
hash: [u8; 32],
range: AlignedRange,
) -> SourceFuture<'_, (UpstreamPullHeader, Self::Reader)> {
Box::pin(async move {
if hash != self.root {
anyhow::bail!("scripted source opened for a foreign hash");
}
let index = self.opened.lock().map_or(0, |mut log| {
log.push((range.fetch_start(), range.fetch_len()));
log.len().saturating_sub(1)
});
let leg = self.timeline.lock().ok().map(|mut t| {
t.push((range.fetch_start(), tokio::time::Instant::now(), None));
t.len().saturating_sub(1)
});
let fault_gate = self
.fault_gates
.iter()
.find(|(nth, _)| *nth == index)
.map(|(_, gate)| gate.clone());
if let Some((refused, make)) = &self.refuse_open
&& refused.contains(&index)
{
if let Some(mut gate) = fault_gate {
let _ = gate.wait_for(|open| *open).await;
}
return Err(make());
}
if let Some((blocks, make)) = &self.refuse_blocks {
let block_bytes = decdn_protocol::discovery_block_bytes();
let first = range.fetch_start() / block_bytes;
let last = range.fetch_end().saturating_sub(1) / block_bytes;
if blocks
.iter()
.any(|&b| (first..=last).contains(&u64::from(b)))
{
return Err(make());
}
}
let served = decdn_bao_range::align_range_clamped(
range.fetch_start(),
range.fetch_len(),
self.total_bytes(),
)?;
let mut wire = self.wire_for(&served)?;
let mut fault = None;
if let Some((after, make)) = &self.fault
&& *after < wire.len()
&& self.faults_left.as_ref().is_none_or(|left| {
left.try_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
.is_ok()
})
{
wire = wire.slice(..*after);
fault = Some(make());
}
let fault_gate = fault.as_ref().and(fault_gate);
let header = UpstreamPullHeader {
total_bytes: self.signed_size.unwrap_or_else(|| self.total_bytes()),
rate_per_mb: SCRIPTED_RATE_PER_MB,
interval_bytes: SCRIPTED_INTERVAL_BYTES,
};
let wire_len = wire.len() as u64;
Ok((
header,
ScriptedReader {
wire,
fault,
leg,
wire_len,
delivered: Arc::clone(&self.delivered),
gate: self.gate.clone(),
fault_gate,
first_read_stall: self.first_read_stall,
stall_after: self.stall_after,
read_so_far: 0,
read_delay: self.read_delay,
},
))
})
}
fn finish(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
Box::pin(async move {
if let Some(stall) = self.finish_stall {
tokio::time::sleep(stall).await;
}
let wire = reader.wire_len;
self.settle(&reader, wire).await
})
}
fn stop(&self, reader: Self::Reader) -> SourceFuture<'_, VoucherProgress> {
Box::pin(async move {
self.stops.fetch_add(1, Ordering::SeqCst);
let wire = reader.wire_len.saturating_sub(reader.wire.len() as u64);
self.settle(&reader, wire).await
})
}
}
#[derive(Debug)]
pub struct ScriptedReader {
wire: Bytes,
fault: Option<anyhow::Error>,
leg: Option<usize>,
gate: Option<tokio::sync::watch::Receiver<bool>>,
fault_gate: Option<tokio::sync::watch::Receiver<bool>>,
wire_len: u64,
delivered: Arc<AtomicU64>,
first_read_stall: Option<Duration>,
stall_after: Option<(u64, Duration)>,
read_so_far: u64,
read_delay: Option<Duration>,
}
impl ScriptedReader {
async fn pass_fault_gate(&mut self) {
if let Some(mut gate) = self.fault_gate.take() {
let _ = gate.wait_for(|open| *open).await;
}
}
}
impl iroh_io::AsyncStreamReader for ScriptedReader {
async fn read_bytes(&mut self, len: usize) -> std::io::Result<Bytes> {
if let Some(mut gate) = self.gate.take() {
let _ = gate.wait_for(|open| *open).await;
}
if let Some(stall) = self.first_read_stall.take() {
tokio::time::sleep(stall).await;
}
if let Some(delay) = self.read_delay {
tokio::time::sleep(delay).await;
}
if let Some((after, stall)) = self.stall_after
&& self.read_so_far >= after
{
self.stall_after = None;
tokio::time::sleep(stall).await;
}
tokio::task::yield_now().await;
if self.wire.len() < len {
self.pass_fault_gate().await;
}
let take = self.wire.len().min(len);
let chunk = self.wire.split_to(take);
self.delivered
.fetch_add(chunk.len() as u64, Ordering::SeqCst);
self.read_so_far = self.read_so_far.saturating_add(chunk.len() as u64);
Ok(chunk)
}
async fn read<const L: usize>(&mut self) -> std::io::Result<[u8; L]> {
if self.wire.len() < L {
self.pass_fault_gate().await;
return Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"scripted reader exhausted before a fixed-size bao read",
));
}
let got = self.wire.split_to(L);
let mut out = [0u8; L];
out.copy_from_slice(&got);
Ok(out)
}
}
impl StashedFault for ScriptedReader {
fn take_fault(&mut self) -> Option<anyhow::Error> {
self.fault.take()
}
}
#[derive(Debug)]
pub struct FakeFunder {
script: Mutex<Vec<Recovery>>,
calls: Mutex<Vec<U256>>,
pool_spent: Option<U256>,
}
impl FakeFunder {
#[must_use]
pub fn new(outcome: Recovery) -> Self {
Self::scripted(vec![outcome])
}
#[must_use]
pub fn scripted(outcomes: Vec<Recovery>) -> Self {
let mut script = outcomes;
script.reverse();
Self {
script: Mutex::new(script),
calls: Mutex::new(Vec::new()),
pool_spent: None,
}
}
#[must_use]
pub const fn with_pool_spent(mut self, spent: U256) -> Self {
self.pool_spent = Some(spent);
self
}
#[must_use]
pub fn calls(&self) -> Vec<U256> {
self.calls.lock().map(|c| c.clone()).unwrap_or_default()
}
}
impl super::Funder for FakeFunder {
fn recover(&self, remaining: U256) -> SourceFuture<'_, Recovery> {
Box::pin(async move {
if let Ok(mut calls) = self.calls.lock() {
calls.push(remaining);
}
let outcome = match self.script.lock() {
Ok(mut script) if script.len() > 1 => script.pop(),
Ok(script) => script.last().copied(),
Err(_) => None,
};
Ok(outcome.unwrap_or(Recovery::Unavailable))
})
}
fn pool_spent(&self) -> Option<U256> {
self.pool_spent
}
}
}
#[cfg(test)]
mod flush_double;
#[cfg(test)]
pub(crate) use flush_double::FlushCountingStore;
#[cfg(any(test, feature = "test-util"))]
pub use doubles::{FakeFunder, ScriptedReader, ScriptedSource};
#[cfg(test)]
pub(crate) use doubles::ctx_with;
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::cast_possible_truncation)] mod tests;