use crate::background::{CheckpointBytes, CheckpointWorkContext, CheckpointWorkError};
use std::io::{self, Write};
pub(crate) struct AdmittedCompressionOutput {
bytes: CheckpointBytes,
used: usize,
}
impl AdmittedCompressionOutput {
pub(crate) fn new(
input_len: usize,
work: &CheckpointWorkContext,
) -> Result<Self, CheckpointWorkError> {
let capacity = zstd::zstd_safe::compress_bound(input_len)
.checked_add(input_len / (64 * 1024) * 32)
.and_then(|bytes| bytes.checked_add(1024))
.ok_or_else(|| {
work.record_failure(CheckpointWorkError::Allocation {
bytes: u64::MAX,
reason: "compressed output capacity overflows usize".into(),
})
})?;
Ok(Self {
bytes: CheckpointBytes::zeroed(capacity, work)?,
used: 0,
})
}
pub(crate) fn into_bytes(mut self) -> CheckpointBytes {
self.bytes.truncate(self.used);
self.bytes
}
}
impl Write for AdmittedCompressionOutput {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
let end = self
.used
.checked_add(bytes.len())
.ok_or_else(|| io::Error::other("compressed output length overflows usize"))?;
let destination = self
.bytes
.as_mut_slice()
.get_mut(self.used..end)
.ok_or_else(|| {
io::Error::other("compressed stream exceeds admitted output capacity")
})?;
destination.copy_from_slice(bytes);
self.used = end;
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}