use super::{
ClosureError, ClosureLimits, GenerationExportManifest, GenerationExportManifestError,
prove_generation_closure_async, validate_generation_export_manifest,
};
use crate::async_storage::AsyncObjectStore;
use crate::cancellation::{CancellationError, CancellationToken};
use crate::foundation::VolumeId;
use crate::model::VolumeConfig;
use crate::performance::{
MeasuredResult, OperationFailure, OperationReceipt, WorkBudget, WorkCounters, WorkError,
};
use crate::storage::{ObjectId, ObjectRead, ObjectReadRequest, ObjectStoreError};
use bytes::Bytes;
use std::mem::size_of;
use thiserror::Error;
#[derive(Clone, Copy, Debug, Default, Eq, Ord, PartialEq, PartialOrd)]
pub struct TransferCursor(u64);
impl TransferCursor {
pub const START: Self = Self(0);
#[must_use]
pub const fn new(next_object: u64) -> Self {
Self(next_object)
}
#[must_use]
pub const fn next_object(self) -> u64 {
self.0
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GenerationTransferBatch {
pub first_object: TransferCursor,
pub next: Option<TransferCursor>,
pub objects: Vec<ObjectRead>,
}
#[derive(Debug, Error)]
pub enum GenerationTransferError {
#[error(transparent)]
Cancelled(#[from] CancellationError),
#[error(transparent)]
Manifest(#[from] GenerationExportManifestError),
#[error(transparent)]
Closure(#[from] ClosureError),
#[error(transparent)]
Object(#[from] ObjectStoreError),
#[error("generation transfer manifest does not match its authenticated closure")]
ManifestMismatch,
#[error("generation transfer cursor is outside the manifest")]
InvalidCursor,
#[error("generation transfer batch bound must be positive")]
EmptyBatch,
#[error("generation transfer batch exceeds its object bound")]
TooManyObjects,
#[error("generation transfer request allocation failed")]
AllocationFailed,
#[error(transparent)]
Work(#[from] WorkError),
}
pub type GenerationTransferResult<T> = MeasuredResult<OperationReceipt<T>, GenerationTransferError>;
pub async fn build_generation_export_manifest_async<S: AsyncObjectStore>(
store: &S,
volume_id: VolumeId,
config: VolumeConfig,
generation_root: ObjectId,
closure_limits: ClosureLimits,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> GenerationTransferResult<GenerationExportManifest> {
cancellation
.check()
.map_err(|error| OperationFailure::before_work(error.into()))?;
config
.validate()
.map_err(GenerationExportManifestError::from)
.map_err(|error| OperationFailure::before_work(error.into()))?;
let proof = prove_generation_closure_async(
store,
generation_root,
closure_limits,
budget,
cancellation,
)
.await
.map_err(|failure| OperationFailure::new(failure.error.into(), *failure.work))?;
let manifest = GenerationExportManifest {
volume_id,
config,
generation_root,
generation_id: proof.generation_id,
objects: proof.objects,
file_count: proof.file_count,
};
validate_generation_export_manifest(&manifest, config.limits.maximum_objects_per_generation)
.map_err(|error| OperationFailure::new(error.into(), proof.work))?;
if proof.root.volume_id != volume_id {
return Err(OperationFailure::new(
GenerationTransferError::ManifestMismatch,
proof.work,
));
}
Ok(OperationReceipt {
value: manifest,
work: proof.work,
})
}
pub async fn authenticate_generation_export_manifest_async<S: AsyncObjectStore>(
store: &S,
manifest: &GenerationExportManifest,
closure_limits: ClosureLimits,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> GenerationTransferResult<()> {
cancellation
.check()
.map_err(|error| OperationFailure::before_work(error.into()))?;
validate_generation_export_manifest(
manifest,
manifest.config.limits.maximum_objects_per_generation,
)
.map_err(|error| OperationFailure::before_work(error.into()))?;
let proof = prove_generation_closure_async(
store,
manifest.generation_root,
closure_limits,
budget,
cancellation,
)
.await
.map_err(|failure| OperationFailure::new(failure.error.into(), *failure.work))?;
if proof.root.volume_id != manifest.volume_id
|| proof.generation_id != manifest.generation_id
|| proof.objects != manifest.objects
|| proof.file_count != manifest.file_count
{
return Err(OperationFailure::new(
GenerationTransferError::ManifestMismatch,
proof.work,
));
}
Ok(OperationReceipt {
value: (),
work: proof.work,
})
}
pub async fn export_generation_batch_async<S: AsyncObjectStore + ?Sized>(
store: &S,
manifest: &GenerationExportManifest,
cursor: TransferCursor,
maximum_objects: u32,
maximum_object_bytes: u64,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> GenerationTransferResult<GenerationTransferBatch> {
cancellation
.check()
.map_err(|error| OperationFailure::before_work(error.into()))?;
validate_generation_export_manifest(
manifest,
manifest.config.limits.maximum_objects_per_generation,
)
.map_err(|error| OperationFailure::before_work(error.into()))?;
if maximum_objects == 0 || maximum_object_bytes == 0 {
return Err(OperationFailure::before_work(
GenerationTransferError::EmptyBatch,
));
}
let object_count = u64::try_from(manifest.objects.len()).unwrap_or(u64::MAX);
if cursor.0 > object_count {
return Err(OperationFailure::before_work(
GenerationTransferError::InvalidCursor,
));
}
if cursor.0 == object_count {
return Ok(OperationReceipt {
value: GenerationTransferBatch {
first_object: cursor,
next: None,
objects: Vec::new(),
},
work: WorkCounters::default(),
});
}
let remaining_objects = object_count - cursor.0;
let count = remaining_objects.min(u64::from(maximum_objects));
let start = usize::try_from(cursor.0)
.map_err(|_| OperationFailure::before_work(GenerationTransferError::InvalidCursor))?;
let count_usize = usize::try_from(count)
.map_err(|_| OperationFailure::before_work(GenerationTransferError::TooManyObjects))?;
let request_bytes = count
.checked_mul(u64::try_from(size_of::<ObjectReadRequest>()).unwrap_or(u64::MAX))
.ok_or_else(|| {
OperationFailure::before_work(GenerationTransferError::Work(WorkError::Overflow))
})?;
let mut work = WorkCounters {
bytes_copied: request_bytes,
allocation_operations: 1,
peak_allocation_bytes: request_bytes,
..WorkCounters::default()
};
work.verify(budget)
.map_err(|error| OperationFailure::before_work(error.into()))?;
let mut requests = Vec::new();
requests
.try_reserve_exact(count_usize)
.map_err(|_| OperationFailure::new(GenerationTransferError::AllocationFailed, work))?;
let admitted_object_bytes =
maximum_object_bytes.min(manifest.config.limits.maximum_object_bytes);
#[allow(
clippy::indexing_slicing,
reason = "`cursor.0 > object_count` was already rejected above, and `cursor.0 == \
object_count` returns early, so `cursor.0 < object_count`; `count = \
remaining_objects.min(maximum_objects)` where `remaining_objects = \
object_count - cursor.0`, so `count <= remaining_objects`, giving `start + \
count_usize <= start + remaining_objects == object_count == \
manifest.objects.len()`"
)]
for object_id in &manifest.objects[start..start + count_usize] {
requests.push(ObjectReadRequest {
object_id: *object_id,
maximum_bytes: admitted_object_bytes,
});
}
let nested = store
.read_many(
&requests,
work.remaining(budget).map_err(|error| {
OperationFailure::new(GenerationTransferError::Work(error), work)
})?,
cancellation,
)
.await
.map_err(|failure| failure.map_with_prior_work(work, GenerationTransferError::from))?;
let nested_peak = request_bytes.saturating_add(nested.work.peak_allocation_bytes);
work = work
.checked_add(nested.work)
.map_err(|error| OperationFailure::new(error.into(), work))?;
work.peak_allocation_bytes = work.peak_allocation_bytes.max(nested_peak);
work.verify(budget)
.map_err(|error| OperationFailure::new(error.into(), work))?;
let next_index = cursor.0 + count;
Ok(OperationReceipt {
value: GenerationTransferBatch {
first_object: cursor,
next: (next_index < object_count).then_some(TransferCursor(next_index)),
objects: nested.value,
},
work,
})
}
pub async fn import_generation_batch_async<S: AsyncObjectStore + ?Sized>(
store: &S,
manifest: &GenerationExportManifest,
cursor: TransferCursor,
bodies: &[Bytes],
maximum_objects: u32,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> GenerationTransferResult<TransferCursor> {
cancellation
.check()
.map_err(|error| OperationFailure::before_work(error.into()))?;
validate_generation_export_manifest(
manifest,
manifest.config.limits.maximum_objects_per_generation,
)
.map_err(|error| OperationFailure::before_work(error.into()))?;
if maximum_objects == 0 {
return Err(OperationFailure::before_work(
GenerationTransferError::EmptyBatch,
));
}
let object_count = u64::try_from(manifest.objects.len()).unwrap_or(u64::MAX);
let body_count = u64::try_from(bodies.len()).unwrap_or(u64::MAX);
if cursor.0 > object_count {
return Err(OperationFailure::before_work(
GenerationTransferError::InvalidCursor,
));
}
if cursor.0 == object_count && bodies.is_empty() {
return Ok(OperationReceipt {
value: cursor,
work: WorkCounters::default(),
});
}
if bodies.is_empty() {
return Err(OperationFailure::before_work(
GenerationTransferError::EmptyBatch,
));
}
if body_count > u64::from(maximum_objects)
|| cursor
.0
.checked_add(body_count)
.is_none_or(|end| end > object_count)
{
return Err(OperationFailure::before_work(
GenerationTransferError::TooManyObjects,
));
}
let start = usize::try_from(cursor.0)
.map_err(|_| OperationFailure::before_work(GenerationTransferError::InvalidCursor))?;
let mut work = WorkCounters::default();
for (offset, body) in bodies.iter().enumerate() {
cancellation
.check()
.map_err(|error| OperationFailure::new(error.into(), work))?;
#[allow(
clippy::indexing_slicing,
reason = "the check above rejects `body_count > maximum_objects` or \
`cursor.0.checked_add(body_count)` overflowing or exceeding `object_count`, \
so `cursor.0 + body_count <= object_count` (== manifest.objects.len()); \
`start == cursor.0` as usize and `offset < bodies.len() == body_count`, so \
`start + offset < start + body_count <= manifest.objects.len()`"
)]
let object_id = manifest.objects[start + offset];
let receipt = store
.put(
object_id,
body.clone(),
work.remaining(budget).map_err(|error| {
OperationFailure::new(GenerationTransferError::Work(error), work)
})?,
cancellation,
)
.await
.map_err(|failure| failure.map_with_prior_work(work, GenerationTransferError::from))?;
work = work
.checked_add(receipt.work)
.map_err(|error| OperationFailure::new(error.into(), work))?;
}
Ok(OperationReceipt {
value: TransferCursor(cursor.0 + body_count),
work,
})
}
#[cfg(test)]
#[path = "tests/transfer.rs"]
mod tests;