use std::future::Future;
use std::io;
use std::path::Path;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use anyhow::Context as _;
use bytes::Bytes;
use tokio::io::{AsyncRead, ReadBuf};
use tokio::sync::Notify;
use decdn_protocol::Coverage;
use crate::credential::CredentialSlot;
use crate::driver::{PacingWait, WaitReason};
use crate::health::PeerHealth;
use crate::pacer::{DownstreamFrontier, PULL_WINDOW_FLOOR, WindowPacer};
use crate::recovery::RecoveryGate;
use crate::scheduler::{AcquireEnv, AcquireTarget, ConsumptionPacing, LaneLease, acquire};
use crate::sink::BlobCache;
use crate::source::Funder;
use crate::source_set::{Holder, SourceProvider, SourceSet};
use crate::stop::{ProgressClock, StopPolicy};
use crate::{ClientRangedStore, PoolContext, PoolLedger, PullConfig, RangedStore};
const MAX_READ_CHUNK: u64 = 1024 * 1024;
struct StreamState {
store: ClientRangedStore,
cursor: AtomicU64,
hint: u64,
consumed: Notify,
progressed: Notify,
outcome: Mutex<Option<Result<(), Arc<anyhow::Error>>>>,
}
impl StreamState {
fn finish(&self, result: anyhow::Result<()>) {
if let Ok(mut slot) = self.outcome.lock()
&& slot.is_none()
{
*slot = Some(result.map_err(Arc::new));
}
self.progressed.notify_waiters();
}
fn outcome(&self) -> Option<Result<(), Arc<anyhow::Error>>> {
match self.outcome.lock() {
Ok(slot) => slot.clone(),
Err(_) => Some(Err(Arc::new(anyhow::anyhow!(
"stream outcome lock poisoned"
)))),
}
}
}
struct ConsumedWait {
state: Arc<StreamState>,
clock: Arc<ProgressClock>,
}
impl ConsumedWait {
async fn unread_ahead(&self, cursor: u64) -> bool {
present_frontier(&self.state.store)
.await
.is_ok_and(|frontier| frontier > cursor)
}
}
impl PacingWait for ConsumedWait {
fn wait(
&self,
observed: DownstreamFrontier,
_reason: WaitReason,
) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
Box::pin(async move {
self.state.progressed.notify_waiters();
loop {
let consumed = self.state.consumed.notified();
let landed = self.state.progressed.notified();
tokio::pin!(consumed, landed);
consumed.as_mut().enable();
landed.as_mut().enable();
let cursor = self.state.cursor.load(Ordering::SeqCst);
if cursor > observed.served_paid {
return;
}
let _hold = self.unread_ahead(cursor).await.then(|| self.clock.hold());
tokio::select! {
() = consumed => {}
() = landed => {}
}
}
})
}
}
#[non_exhaustive]
pub struct StreamCandidate<S> {
pub source: S,
pub ctx: Arc<Mutex<PoolContext>>,
pub ledger: Arc<PoolLedger>,
pub coverage: Option<Coverage>,
#[doc(hidden)]
pub lease: LaneLease,
#[doc(hidden)]
pub widen: Option<crate::LaneWiden>,
}
impl<S> StreamCandidate<S> {
#[must_use]
pub fn new(source: S, ctx: Arc<Mutex<PoolContext>>, ledger: Arc<PoolLedger>) -> Self {
Self {
source,
ctx,
ledger,
coverage: None,
lease: LaneLease::default(),
widen: None,
}
}
}
impl<S> std::fmt::Debug for StreamCandidate<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StreamCandidate").finish_non_exhaustive()
}
}
struct DriveInputs<P, F> {
provider: P,
holders: Vec<Holder>,
health: Arc<PeerHealth>,
funder: F,
recovery: Arc<RecoveryGate>,
credentials: Option<CredentialSlot>,
stop: StopPolicy,
max_blob_bytes: u64,
}
async fn run_drive<P, F>(
state: Arc<StreamState>,
inputs: DriveInputs<P, F>,
read_ahead: u64,
lane_cap: usize,
cache: Arc<dyn BlobCache>,
hash: [u8; 32],
) -> anyhow::Result<()>
where
P: SourceProvider,
F: Funder,
{
let DriveInputs {
provider,
holders,
health,
funder,
recovery,
credentials,
stop,
max_blob_bytes,
} = inputs;
let pacer = WindowPacer::new(read_ahead.max(PULL_WINDOW_FLOOR));
let on_progress = {
let state = Arc::clone(&state);
move |_position: u64, _total: u64| state.progressed.notify_waiters()
};
let downstream = {
let state = Arc::clone(&state);
move || DownstreamFrontier {
served_paid: state.cursor.load(Ordering::SeqCst),
serve_demand: 0,
}
};
let wait = ConsumedWait {
state: Arc::clone(&state),
clock: Arc::clone(&stop.clock),
};
let pacing = ConsumptionPacing {
downstream: &downstream,
pacing_wait: &wait,
};
let mut sources = SourceSet::new(&provider, hash, health, holders);
let ranges = [(0, state.store.bound())];
let env = AcquireEnv {
pacer: &pacer,
funder: &funder,
recovery: &recovery,
credentials: credentials.as_ref(),
max_lanes: lane_cap.max(1),
stop: &stop,
on_progress: Some(&on_progress),
ledgers: None,
pacing: Some(&pacing),
max_blob_bytes,
run: None,
};
let result = acquire(
AcquireTarget {
store: &state.store,
hash,
total_bytes: state.hint,
ranges: &ranges,
},
&mut sources,
&env,
)
.await;
if result.is_ok() {
if cache.caches()
&& let Some(proven) = state.store.proven()
&& let Ok(whole) = state.store.read(0, proven).await
{
let _ = cache.put(hash, 0, whole).await;
}
}
result
}
pub struct Streamer<'a, P, F> {
provider: P,
holders: Vec<Holder>,
health: Arc<PeerHealth>,
funder: F,
recovery: Arc<RecoveryGate>,
credentials: Option<CredentialSlot>,
scratch: &'a Path,
max_blob_bytes: u64,
}
impl<P, F> std::fmt::Debug for Streamer<'_, P, F> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Streamer")
.field("holders", &self.holders.len())
.field("recovery", &self.recovery)
.field("scratch", &self.scratch)
.finish_non_exhaustive()
}
}
impl<'a, P, F> Streamer<'a, P, F> {
#[must_use]
pub fn new(provider: P, funder: F, scratch: &'a Path) -> Self {
Self {
provider,
holders: Vec::new(),
health: Arc::default(),
funder,
recovery: Arc::default(),
credentials: None,
scratch,
max_blob_bytes: 0,
}
}
#[must_use]
pub fn holders(mut self, holders: Vec<Holder>) -> Self {
self.holders = holders;
self
}
#[must_use]
pub fn recovery(mut self, recovery: Arc<RecoveryGate>) -> Self {
self.recovery = recovery;
self
}
#[must_use]
pub fn credentials(mut self, slot: CredentialSlot) -> Self {
self.credentials = Some(slot);
self
}
#[doc(hidden)]
#[must_use]
pub fn health(mut self, health: Arc<PeerHealth>) -> Self {
self.health = health;
self
}
#[must_use]
pub const fn max_blob_bytes(mut self, max_blob_bytes: u64) -> Self {
self.max_blob_bytes = max_blob_bytes;
self
}
}
impl<P, F> Streamer<'_, P, F>
where
P: SourceProvider,
F: Funder,
{
pub async fn open<'a>(
self,
hash: [u8; 32],
total_bytes: u64,
config: &PullConfig,
cache: Arc<dyn BlobCache>,
stop: StopPolicy,
) -> anyhow::Result<(VerifiedReader, StreamDrive<'a>)>
where
P: 'a,
F: 'a,
{
if let Some(cached) = cache.get(hash, 0, total_bytes).await?
&& u64::try_from(cached.len()).is_ok_and(|len| len == total_bytes)
{
return Ok((
VerifiedReader::Cached {
data: cached,
pos: 0,
},
StreamDrive {
fetch: None,
state: None,
error: None,
},
));
}
let stem = blake3::Hash::from_bytes(hash).to_hex();
let total_bytes = if self.max_blob_bytes > 0 {
total_bytes.min(self.max_blob_bytes)
} else {
total_bytes
};
let store = {
let (scratch, stem) = (self.scratch.to_path_buf(), stem.to_string());
tokio::task::spawn_blocking(move || {
ClientRangedStore::open_or_create(&scratch, &stem, hash, total_bytes)
.map_err(|e| anyhow::anyhow!("open stream store for {stem}: {e}"))
})
.await
.map_err(|e| anyhow::anyhow!("open stream store task: {e}"))??
};
let state = Arc::new(StreamState {
store,
cursor: AtomicU64::new(0),
hint: total_bytes,
consumed: Notify::new(),
progressed: Notify::new(),
outcome: Mutex::new(None),
});
let fetch = Box::pin(run_drive(
Arc::clone(&state),
DriveInputs {
provider: self.provider,
holders: self.holders,
health: self.health,
funder: self.funder,
recovery: self.recovery,
credentials: self.credentials,
stop,
max_blob_bytes: self.max_blob_bytes,
},
config.read_ahead_bytes,
config.streamer_lane_cap,
cache,
hash,
));
Ok((
VerifiedReader::Live(LiveReader {
state: Arc::clone(&state),
read: None,
buffered: Bytes::new(),
}),
StreamDrive {
fetch: Some(fetch),
state: Some(state),
error: None,
},
))
}
}
pub struct StreamDrive<'a> {
fetch: Option<Pin<Box<dyn Future<Output = anyhow::Result<()>> + 'a>>>,
state: Option<Arc<StreamState>>,
error: Option<anyhow::Error>,
}
impl std::fmt::Debug for StreamDrive<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StreamDrive")
.field("running", &self.fetch.is_some())
.finish_non_exhaustive()
}
}
impl StreamDrive<'_> {
pub async fn alongside<T>(&mut self, consume: impl Future<Output = T>) -> T {
tokio::pin!(consume);
tokio::select! {
biased;
() = &mut *self => consume.await,
out = &mut consume => out,
}
}
pub const fn take_error(&mut self) -> Option<anyhow::Error> {
self.error.take()
}
}
impl Future for StreamDrive<'_> {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = &mut *self;
let Some(fetch) = this.fetch.as_mut() else {
return Poll::Ready(());
};
match fetch.as_mut().poll(cx) {
Poll::Ready(result) => {
this.fetch = None;
let recorded = match result {
Ok(()) => Ok(()),
Err(err) => {
let message = anyhow::anyhow!("{err:#}");
this.error = Some(err);
Err(message)
}
};
if let Some(state) = &this.state {
state.finish(recorded);
}
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for StreamDrive<'_> {
fn drop(&mut self) {
if self.fetch.is_some()
&& let Some(state) = &self.state
{
state.finish(Err(anyhow::anyhow!(
"stream fetch stopped before the blob was complete"
)));
}
}
}
#[non_exhaustive]
pub enum VerifiedReader {
Cached {
data: Bytes,
pos: usize,
},
Live(LiveReader),
}
type VerifiedRead = Pin<Box<dyn Future<Output = io::Result<Option<Bytes>>> + Send>>;
pub struct LiveReader {
state: Arc<StreamState>,
read: Option<VerifiedRead>,
buffered: Bytes,
}
async fn present_frontier(store: &ClientRangedStore) -> anyhow::Result<u64> {
let present = store
.present_ranges()
.await
.context("read present ranges")?;
Ok(
match crate::driver::contiguous_byte_ranges(&present, store.bound()).first() {
Some(&(0, len)) => len,
_ => 0,
},
)
}
async fn next_verified(state: Arc<StreamState>, cursor: u64) -> io::Result<Option<Bytes>> {
loop {
if state.store.proven().is_some_and(|proven| cursor >= proven) {
return Ok(None);
}
let progressed = state.progressed.notified();
tokio::pin!(progressed);
progressed.as_mut().enable();
let grew = state.store.present_grew().notified();
tokio::pin!(grew);
grew.as_mut().enable();
let outcome = state.outcome();
let frontier = present_frontier(&state.store)
.await
.map_err(|e| io::Error::other(format!("present frontier: {e:#}")))?;
if frontier > cursor {
let len = (frontier - cursor).min(MAX_READ_CHUNK);
return state.store.read(cursor, len).await.map(Some).map_err(|e| {
io::Error::other(format!("read verified prefix: {:#}", anyhow::Error::new(e)))
});
}
match outcome {
Some(Err(e)) => {
return Err(io::Error::other(format!("stream fetch failed: {e:#}")));
}
Some(Ok(())) => {
return Err(io::Error::other(format!(
"stream fetch finished with {cursor} of {} bytes readable",
state.store.bound()
)));
}
None => {}
}
tokio::select! {
() = progressed => {}
() = grew => {}
}
}
}
impl std::fmt::Debug for LiveReader {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LiveReader")
.field("cursor", &self.state.cursor.load(Ordering::SeqCst))
.field("bound", &self.state.store.bound())
.field("buffered", &self.buffered.len())
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for VerifiedReader {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Cached { data, pos } => f
.debug_struct("VerifiedReader::Cached")
.field("len", &data.len())
.field("pos", pos)
.finish(),
Self::Live(_) => f
.debug_struct("VerifiedReader::Live")
.finish_non_exhaustive(),
}
}
}
impl LiveReader {
fn poll_read(&mut self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
loop {
if !self.buffered.is_empty() {
let n = self.buffered.len().min(buf.remaining());
if n == 0 {
return Poll::Ready(Ok(()));
}
let Ok(taken) = u64::try_from(n) else {
return Poll::Ready(Err(io::Error::other("read length does not fit u64")));
};
let chunk = self.buffered.split_to(n);
buf.put_slice(&chunk);
self.state.cursor.fetch_add(taken, Ordering::SeqCst);
self.state.consumed.notify_waiters();
return Poll::Ready(Ok(()));
}
let read = self.read.get_or_insert_with(|| {
let state = Arc::clone(&self.state);
let cursor = state.cursor.load(Ordering::SeqCst);
Box::pin(next_verified(state, cursor))
});
match read.as_mut().poll(cx) {
Poll::Ready(Ok(Some(bytes))) => {
self.read = None;
self.buffered = bytes;
}
Poll::Ready(Ok(None)) => {
self.read = None;
return Poll::Ready(Ok(()));
}
Poll::Ready(Err(e)) => {
self.read = None;
return Poll::Ready(Err(e));
}
Poll::Pending => return Poll::Pending,
}
}
}
}
impl AsyncRead for VerifiedReader {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
match self.get_mut() {
VerifiedReader::Cached { data, pos } => {
let remaining = data.len().saturating_sub(*pos);
if remaining == 0 {
return Poll::Ready(Ok(()));
}
let n = remaining.min(buf.remaining());
if let Some(slice) = data.get(*pos..pos.saturating_add(n)) {
buf.put_slice(slice);
*pos = pos.saturating_add(n);
}
Poll::Ready(Ok(()))
}
VerifiedReader::Live(live) => live.poll_read(cx, buf),
}
}
}
#[cfg(test)]
mod tests;