455 lines
14 KiB
Rust
455 lines
14 KiB
Rust
|
use std::future::Future;
|
||
|
|
use std::time::{Duration, Instant};
|
||
|
|
|
||
|
|
use tokio::time::timeout;
|
||
|
|
|
||
|
|
use crate::cancel::CancelToken;
|
||
|
|
use crate::error::{StreamError, StreamSide};
|
||
|
|
use crate::failure::{StructuredFailure, classify_stream_error};
|
||
|
|
use crate::integrity::{IntegrityState, Sha256Accountant};
|
||
|
|
use crate::ports::{SourceBody, TargetBody, TargetFinish};
|
||
|
|
use crate::progress::{TransferCheckpoint, TransferProgress};
|
||
|
|
|
||
|
|
/// Race an I/O future against cancellation and the no-progress bound.
|
||
|
|
async fn bounded_io<F, T>(
|
||
|
|
cancel: &CancelToken,
|
||
|
|
bound: Duration,
|
||
|
|
side: StreamSide,
|
||
|
|
fut: F,
|
||
|
|
) -> Result<T, StreamError>
|
||
|
|
where
|
||
|
|
F: Future<Output = Result<T, StreamError>>,
|
||
|
|
{
|
||
|
|
let started = Instant::now();
|
||
|
|
tokio::select! {
|
||
|
|
biased;
|
||
|
|
() = cancel.cancelled() => Err(StreamError::Cancelled),
|
||
|
|
result = timeout(bound, fut) => match result {
|
||
|
|
Ok(inner) => inner,
|
||
|
|
Err(_) => {
|
||
|
|
let elapsed_ms =
|
||
|
|
u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||
|
|
Err(StreamError::NoProgress { side, elapsed_ms })
|
||
|
|
}
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Bounded-buffer copy configuration.
|
||
|
|
///
|
||
|
|
/// Memory used by the kernel for payload data is one buffer of `buffer_size`
|
||
|
|
/// bytes plus fixed integrity state, independent of total transfer size.
|
||
|
|
///
|
||
|
|
/// `initial_completed_bytes` is the absolute prefix already retained before this
|
||
|
|
/// invocation. When it is non-zero, callers must seed integrity with
|
||
|
|
/// [`copy_stream_seeded`] so the accountant's `bytes_hashed` matches that
|
||
|
|
/// prefix; otherwise the copy fails closed.
|
||
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
||
|
|
pub struct CopyConfig {
|
||
|
|
pub buffer_size: usize,
|
||
|
|
pub progress_interval_bytes: u64,
|
||
|
|
pub checkpoint_interval_bytes: u64,
|
||
|
|
pub no_progress_timeout: Duration,
|
||
|
|
pub initial_completed_bytes: u64,
|
||
|
|
pub total_bytes: Option<u64>,
|
||
|
|
pub completed_members: u64,
|
||
|
|
pub total_members: Option<u64>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl CopyConfig {
|
||
|
|
pub const DEFAULT_BUFFER_SIZE: usize = 64 * 1024;
|
||
|
|
pub const MIN_BUFFER_SIZE: usize = 4 * 1024;
|
||
|
|
pub const MAX_BUFFER_SIZE: usize = 8 * 1024 * 1024;
|
||
|
|
|
||
|
|
#[must_use]
|
||
|
|
pub fn sanitized(self) -> Self {
|
||
|
|
let buffer_size = self
|
||
|
|
.buffer_size
|
||
|
|
.clamp(Self::MIN_BUFFER_SIZE, Self::MAX_BUFFER_SIZE);
|
||
|
|
let progress_interval_bytes = if self.progress_interval_bytes == 0 {
|
||
|
|
buffer_size as u64
|
||
|
|
} else {
|
||
|
|
self.progress_interval_bytes
|
||
|
|
};
|
||
|
|
let checkpoint_interval_bytes = if self.checkpoint_interval_bytes == 0 {
|
||
|
|
progress_interval_bytes.saturating_mul(16)
|
||
|
|
} else {
|
||
|
|
self.checkpoint_interval_bytes
|
||
|
|
};
|
||
|
|
let no_progress_timeout = if self.no_progress_timeout.is_zero() {
|
||
|
|
Duration::from_secs(60)
|
||
|
|
} else {
|
||
|
|
self.no_progress_timeout
|
||
|
|
};
|
||
|
|
Self {
|
||
|
|
buffer_size,
|
||
|
|
progress_interval_bytes,
|
||
|
|
checkpoint_interval_bytes,
|
||
|
|
no_progress_timeout,
|
||
|
|
initial_completed_bytes: self.initial_completed_bytes,
|
||
|
|
total_bytes: self.total_bytes,
|
||
|
|
completed_members: self.completed_members,
|
||
|
|
total_members: self.total_members,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Default for CopyConfig {
|
||
|
|
fn default() -> Self {
|
||
|
|
Self {
|
||
|
|
buffer_size: Self::DEFAULT_BUFFER_SIZE,
|
||
|
|
progress_interval_bytes: Self::DEFAULT_BUFFER_SIZE as u64,
|
||
|
|
checkpoint_interval_bytes: (Self::DEFAULT_BUFFER_SIZE as u64).saturating_mul(16),
|
||
|
|
no_progress_timeout: Duration::from_secs(60),
|
||
|
|
initial_completed_bytes: 0,
|
||
|
|
total_bytes: None,
|
||
|
|
completed_members: 0,
|
||
|
|
total_members: None,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Hooks invoked when progress or checkpoint thresholds advance.
|
||
|
|
pub trait CopyHooks: Send {
|
||
|
|
fn on_progress(
|
||
|
|
&mut self,
|
||
|
|
progress: &TransferProgress,
|
||
|
|
) -> impl Future<Output = Result<(), StreamError>> + Send;
|
||
|
|
|
||
|
|
fn on_checkpoint(
|
||
|
|
&mut self,
|
||
|
|
checkpoint: &TransferCheckpoint,
|
||
|
|
) -> impl Future<Output = Result<(), StreamError>> + Send;
|
||
|
|
}
|
||
|
|
|
||
|
|
impl<H: CopyHooks + ?Sized> CopyHooks for &mut H {
|
||
|
|
fn on_progress(
|
||
|
|
&mut self,
|
||
|
|
progress: &TransferProgress,
|
||
|
|
) -> impl Future<Output = Result<(), StreamError>> + Send {
|
||
|
|
(**self).on_progress(progress)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn on_checkpoint(
|
||
|
|
&mut self,
|
||
|
|
checkpoint: &TransferCheckpoint,
|
||
|
|
) -> impl Future<Output = Result<(), StreamError>> + Send {
|
||
|
|
(**self).on_checkpoint(checkpoint)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// No-op hooks for callers that record progress elsewhere.
|
||
|
|
pub struct NoopHooks;
|
||
|
|
|
||
|
|
impl CopyHooks for NoopHooks {
|
||
|
|
async fn on_progress(&mut self, _progress: &TransferProgress) -> Result<(), StreamError> {
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn on_checkpoint(&mut self, _checkpoint: &TransferCheckpoint) -> Result<(), StreamError> {
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Successful completion of one member copy.
|
||
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
||
|
|
pub struct CopyResult {
|
||
|
|
pub progress: TransferProgress,
|
||
|
|
pub integrity: IntegrityState,
|
||
|
|
pub target: TargetFinish,
|
||
|
|
/// Configured payload buffer length used for this copy.
|
||
|
|
pub buffer_len: usize,
|
||
|
|
/// Reserved capacity of the payload buffer (never grown during the copy).
|
||
|
|
pub buffer_capacity: usize,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Terminal result of a bounded copy.
|
||
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
||
|
|
pub enum CopyOutcome {
|
||
|
|
Completed(CopyResult),
|
||
|
|
Cancelled {
|
||
|
|
progress: TransferProgress,
|
||
|
|
integrity: IntegrityState,
|
||
|
|
buffer_len: usize,
|
||
|
|
buffer_capacity: usize,
|
||
|
|
},
|
||
|
|
Failed {
|
||
|
|
failure: StructuredFailure,
|
||
|
|
progress: TransferProgress,
|
||
|
|
integrity: IntegrityState,
|
||
|
|
buffer_len: usize,
|
||
|
|
buffer_capacity: usize,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Copy from `source` into `target` with a single bounded buffer and a fresh
|
||
|
|
/// integrity accountant (session starts at zero hashed bytes).
|
||
|
|
pub async fn copy_stream<S, T, H>(
|
||
|
|
source: S,
|
||
|
|
target: T,
|
||
|
|
config: CopyConfig,
|
||
|
|
cancel: &CancelToken,
|
||
|
|
hooks: H,
|
||
|
|
) -> CopyOutcome
|
||
|
|
where
|
||
|
|
S: SourceBody,
|
||
|
|
T: TargetBody,
|
||
|
|
H: CopyHooks,
|
||
|
|
{
|
||
|
|
copy_stream_seeded(
|
||
|
|
source,
|
||
|
|
target,
|
||
|
|
config,
|
||
|
|
cancel,
|
||
|
|
hooks,
|
||
|
|
Sha256Accountant::new(),
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Copy with a pre-seeded integrity accountant for resume.
|
||
|
|
///
|
||
|
|
/// `accountant.bytes_hashed()` must equal `config.initial_completed_bytes`. A
|
||
|
|
/// mismatch fails closed so absolute progress and integrity stay consistent.
|
||
|
|
pub async fn copy_stream_seeded<S, T, H>(
|
||
|
|
mut source: S,
|
||
|
|
mut target: T,
|
||
|
|
config: CopyConfig,
|
||
|
|
cancel: &CancelToken,
|
||
|
|
mut hooks: H,
|
||
|
|
mut accountant: Sha256Accountant,
|
||
|
|
) -> CopyOutcome
|
||
|
|
where
|
||
|
|
S: SourceBody,
|
||
|
|
T: TargetBody,
|
||
|
|
H: CopyHooks,
|
||
|
|
{
|
||
|
|
let config = config.sanitized();
|
||
|
|
let mut buffer = vec![0u8; config.buffer_size];
|
||
|
|
let buffer_len = buffer.len();
|
||
|
|
let buffer_capacity = buffer.capacity();
|
||
|
|
let initial = config.initial_completed_bytes;
|
||
|
|
let mut completed_bytes = initial;
|
||
|
|
let mut session_bytes = 0u64;
|
||
|
|
let mut last_progress_at = initial;
|
||
|
|
let mut last_checkpoint_at = initial;
|
||
|
|
|
||
|
|
let progress_snapshot = |completed_bytes: u64, session_bytes: u64| {
|
||
|
|
TransferProgress::new(
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
config.total_bytes,
|
||
|
|
config.completed_members,
|
||
|
|
config.total_members,
|
||
|
|
)
|
||
|
|
};
|
||
|
|
|
||
|
|
let fail = |error: StreamError,
|
||
|
|
completed_bytes: u64,
|
||
|
|
session_bytes: u64,
|
||
|
|
accountant: Sha256Accountant| {
|
||
|
|
let progress = progress_snapshot(completed_bytes, session_bytes);
|
||
|
|
let integrity = accountant.finalize();
|
||
|
|
if error.is_cancelled() {
|
||
|
|
CopyOutcome::Cancelled {
|
||
|
|
progress,
|
||
|
|
integrity,
|
||
|
|
buffer_len,
|
||
|
|
buffer_capacity,
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
CopyOutcome::Failed {
|
||
|
|
failure: classify_stream_error(&error),
|
||
|
|
progress,
|
||
|
|
integrity,
|
||
|
|
buffer_len,
|
||
|
|
buffer_capacity,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
if accountant.bytes_hashed() != initial {
|
||
|
|
return fail(
|
||
|
|
StreamError::internal(format!(
|
||
|
|
"integrity seed covers {} bytes but initial_completed_bytes is {initial}",
|
||
|
|
accountant.bytes_hashed()
|
||
|
|
)),
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
loop {
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
let n = match bounded_io(
|
||
|
|
cancel,
|
||
|
|
config.no_progress_timeout,
|
||
|
|
StreamSide::Source,
|
||
|
|
source.read(&mut buffer),
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
Ok(0) => break,
|
||
|
|
Ok(n) => n,
|
||
|
|
Err(error) => {
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
if n > buffer.len() {
|
||
|
|
return fail(
|
||
|
|
StreamError::internal(format!(
|
||
|
|
"source read returned {n} bytes into a {}-byte buffer",
|
||
|
|
buffer.len()
|
||
|
|
)),
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
let chunk = &buffer[..n];
|
||
|
|
if let Err(error) = bounded_io(
|
||
|
|
cancel,
|
||
|
|
config.no_progress_timeout,
|
||
|
|
StreamSide::Target,
|
||
|
|
target.write_all(chunk),
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
|
||
|
|
accountant.update(chunk);
|
||
|
|
session_bytes = session_bytes.saturating_add(n as u64);
|
||
|
|
completed_bytes = initial.saturating_add(session_bytes);
|
||
|
|
|
||
|
|
if completed_bytes.saturating_sub(last_progress_at) >= config.progress_interval_bytes {
|
||
|
|
let progress = progress_snapshot(completed_bytes, session_bytes);
|
||
|
|
if let Err(error) = hooks.on_progress(&progress).await {
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
last_progress_at = completed_bytes;
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if completed_bytes.saturating_sub(last_checkpoint_at) >= config.checkpoint_interval_bytes {
|
||
|
|
let checkpoint =
|
||
|
|
TransferCheckpoint::new(completed_bytes, session_bytes, accountant.snapshot());
|
||
|
|
if let Err(error) = hooks.on_checkpoint(&checkpoint).await {
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
last_checkpoint_at = completed_bytes;
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Final progress/checkpoint at EOF when bytes advanced past the last report.
|
||
|
|
if completed_bytes > last_progress_at || session_bytes == 0 {
|
||
|
|
let progress = progress_snapshot(completed_bytes, session_bytes);
|
||
|
|
if let Err(error) = hooks.on_progress(&progress).await {
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if completed_bytes > last_checkpoint_at && session_bytes > 0 {
|
||
|
|
let checkpoint =
|
||
|
|
TransferCheckpoint::new(completed_bytes, session_bytes, accountant.snapshot());
|
||
|
|
if let Err(error) = hooks.on_checkpoint(&checkpoint).await {
|
||
|
|
return fail(error, completed_bytes, session_bytes, accountant);
|
||
|
|
}
|
||
|
|
if cancel.is_cancelled() {
|
||
|
|
return fail(
|
||
|
|
StreamError::Cancelled,
|
||
|
|
completed_bytes,
|
||
|
|
session_bytes,
|
||
|
|
accountant,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let integrity = accountant.finalize();
|
||
|
|
let progress = progress_snapshot(completed_bytes, session_bytes);
|
||
|
|
match bounded_io(
|
||
|
|
cancel,
|
||
|
|
config.no_progress_timeout,
|
||
|
|
StreamSide::Target,
|
||
|
|
target.finish(),
|
||
|
|
)
|
||
|
|
.await
|
||
|
|
{
|
||
|
|
Ok(target_finish) => CopyOutcome::Completed(CopyResult {
|
||
|
|
progress,
|
||
|
|
integrity,
|
||
|
|
target: target_finish,
|
||
|
|
buffer_len,
|
||
|
|
buffer_capacity,
|
||
|
|
}),
|
||
|
|
Err(error) => {
|
||
|
|
if error.is_cancelled() {
|
||
|
|
CopyOutcome::Cancelled {
|
||
|
|
progress,
|
||
|
|
integrity,
|
||
|
|
buffer_len,
|
||
|
|
buffer_capacity,
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
CopyOutcome::Failed {
|
||
|
|
failure: classify_stream_error(&error),
|
||
|
|
progress,
|
||
|
|
integrity,
|
||
|
|
buffer_len,
|
||
|
|
buffer_capacity,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|