use keyhog_profile::RetryCause;
use std::fs::{File, Metadata};
use std::io;
use std::path::Path;
use std::time::Duration;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
}
impl RetryPolicy {
pub const DEFAULT: Self = Self {
max_attempts: 3,
initial_backoff: Duration::from_millis(5),
max_backoff: Duration::from_millis(40),
};
#[must_use]
pub fn backoff_for(&self, next_attempt: u32) -> Duration {
let doublings = next_attempt.saturating_sub(2).min(16);
let scaled = self
.initial_backoff
.saturating_mul(1u32 << doublings.min(16));
scaled.min(self.max_backoff)
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::DEFAULT
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PathOrigin {
OperatorSupplied,
Enumerated,
}
#[must_use]
pub fn classify_io(error: &io::Error, origin: PathOrigin) -> Option<RetryCause> {
match error.kind() {
io::ErrorKind::Interrupted => return Some(RetryCause::Interrupted),
io::ErrorKind::WouldBlock => return Some(RetryCause::WouldBlock),
io::ErrorKind::TimedOut
| io::ErrorKind::ConnectionReset
| io::ErrorKind::ConnectionAborted => return Some(RetryCause::Network),
io::ErrorKind::NotFound if matches!(origin, PathOrigin::Enumerated) => {
return Some(RetryCause::VanishedUnderWalk)
}
_ => {}
}
classify_raw_os(error.raw_os_error()?, origin)
}
fn classify_raw_os(errno: i32, origin: PathOrigin) -> Option<RetryCause> {
#[cfg(unix)]
{
const ESTALE: i32 = 116;
const EBUSY: i32 = 16;
const ETXTBSY: i32 = 26;
match errno {
ESTALE if matches!(origin, PathOrigin::Enumerated) => {
return Some(RetryCause::VanishedUnderWalk)
}
EBUSY | ETXTBSY => return Some(RetryCause::Locked),
_ => {}
}
}
#[cfg(windows)]
{
const ERROR_SHARING_VIOLATION: i32 = 32;
const ERROR_LOCK_VIOLATION: i32 = 33;
if matches!(errno, ERROR_SHARING_VIOLATION | ERROR_LOCK_VIOLATION) {
return Some(RetryCause::Locked);
}
}
let _ = (errno, origin); None
}
pub fn retry_classified<T, E, C, F>(policy: RetryPolicy, classify: C, mut op: F) -> Result<T, E>
where
C: Fn(&E) -> Option<RetryCause>,
F: FnMut() -> Result<T, E>,
{
let mut attempt = 1u32;
loop {
match op() {
Ok(value) => return Ok(value),
Err(error) => {
if attempt >= policy.max_attempts {
return Err(error);
}
let Some(cause) = classify(&error) else {
return Err(error);
};
attempt += 1;
keyhog_profile::record_retry(cause);
let backoff = policy.backoff_for(attempt);
if !backoff.is_zero() {
std::thread::sleep(backoff);
}
}
}
}
}
pub fn retry_io<T, F>(policy: RetryPolicy, origin: PathOrigin, op: F) -> io::Result<T>
where
F: FnMut() -> io::Result<T>,
{
retry_classified(policy, |error| classify_io(error, origin), op)
}
pub struct OpenedFile {
pub file: File,
pub metadata: Metadata,
}
pub fn open_enumerated(path: &Path) -> io::Result<OpenedFile> {
retry_io(RetryPolicy::DEFAULT, PathOrigin::Enumerated, || {
let file = File::open(path)?;
let metadata = file.metadata()?;
Ok(OpenedFile { file, metadata })
})
}
pub struct RetryingContentSource<'a> {
inner: &'a dyn crate::FileContentSource,
policy: RetryPolicy,
}
impl<'a> RetryingContentSource<'a> {
#[must_use]
pub fn new(inner: &'a dyn crate::FileContentSource) -> Self {
Self {
inner,
policy: RetryPolicy::DEFAULT,
}
}
}
impl crate::FileContentSource for RetryingContentSource<'_> {
fn read_prefix(
&self,
path: &str,
max_bytes: u64,
) -> Result<crate::FileContent, crate::ContentError> {
retry_classified(
self.policy,
|error| match error {
crate::ContentError::TransientRead => Some(RetryCause::VanishedUnderWalk),
crate::ContentError::PermanentRead | crate::ContentError::NotUtf8 => None,
},
|| self.inner.read_prefix(path, max_bytes),
)
}
}