Add bounded streaming execution kernel
This commit is contained in:
parent
ac2289bd0a
commit
f569707f97
1 changed files with 2384 additions and 1 deletions
455
crates/remora-stream/src/copy.rs
Normal file
455
crates/remora-stream/src/copy.rs
Normal file
|
|
@ -0,0 +1,455 @@
|
|||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue