use bytesize::ByteSize;
use serde::{Deserialize, Deserializer};
use std::time::Duration;
fn de_byte_size<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
ByteSize::deserialize(d).map(|b| b.as_u64())
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct BatchConfig {
pub max_rows: u64,
#[serde(deserialize_with = "de_byte_size")]
pub max_bytes: u64,
#[serde(with = "humantime_serde")]
pub linger: Duration,
}
impl Default for BatchConfig {
fn default() -> Self {
BatchConfig {
max_rows: 500_000,
max_bytes: 128 * 1024 * 1024,
linger: Duration::from_secs(1),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct InflightConfig {
pub max_per_shard: usize,
}
impl Default for InflightConfig {
fn default() -> Self {
InflightConfig { max_per_shard: 2 }
}
}
#[derive(Clone, Copy, Debug, PartialEq, Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct RetryConfig {
#[serde(with = "humantime_serde")]
pub initial: Duration,
#[serde(with = "humantime_serde")]
pub max: Duration,
pub multiplier: f64,
pub jitter: f64,
pub max_attempts: u32,
}
impl Default for RetryConfig {
fn default() -> Self {
RetryConfig {
initial: Duration::from_millis(100),
max: Duration::from_secs(10),
multiplier: 2.0,
jitter: 0.2,
max_attempts: 0,
}
}
}
#[derive(Clone, Debug, PartialEq, thiserror::Error)]
#[non_exhaustive]
pub enum RetryConfigError {
#[error("retry.multiplier must be a finite number in [1.0, 1e9] (got {0})")]
Multiplier(f64),
#[error("retry.jitter must be a finite fraction in [0.0, 1.0] (got {0})")]
Jitter(f64),
#[error("retry.initial and retry.max must be non-zero")]
ZeroDelay,
#[error("retry.initial ({initial:?}) must not exceed retry.max ({max:?})")]
InitialExceedsMax {
initial: Duration,
max: Duration,
},
}
impl RetryConfig {
pub fn validate(&self) -> Result<(), RetryConfigError> {
if !self.multiplier.is_finite() || !(1.0..=1e9).contains(&self.multiplier) {
return Err(RetryConfigError::Multiplier(self.multiplier));
}
if !self.jitter.is_finite() || !(0.0..=1.0).contains(&self.jitter) {
return Err(RetryConfigError::Jitter(self.jitter));
}
if self.initial.is_zero() || self.max.is_zero() {
return Err(RetryConfigError::ZeroDelay);
}
if self.initial > self.max {
return Err(RetryConfigError::InitialExceedsMax {
initial: self.initial,
max: self.max,
});
}
Ok(())
}
pub(crate) fn stalls_indefinitely(&self) -> bool {
const CEILING: Duration = Duration::from_secs(300);
self.max_attempts == 0 && self.max > CEILING
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct BreakerConfig {
pub failure_threshold: u32,
#[serde(with = "humantime_serde")]
pub open_for: Duration,
pub half_open_probes: u32,
}
impl Default for BreakerConfig {
fn default() -> Self {
BreakerConfig {
failure_threshold: 3,
open_for: Duration::from_secs(5),
half_open_probes: 1,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct SinkPoolConfig {
pub batch: BatchConfig,
pub inflight: InflightConfig,
pub retry: RetryConfig,
pub breaker: BreakerConfig,
}
#[cfg(test)]
mod tests {
use super::*;
fn retry(mutate: impl FnOnce(&mut RetryConfig)) -> RetryConfig {
let mut cfg = RetryConfig::default();
mutate(&mut cfg);
cfg
}
#[test]
fn validate_rejects_policies_that_misbehave_at_runtime() {
let cases = [
(retry(|c| c.multiplier = 0.5), "multiplier"),
(retry(|c| c.multiplier = -2.0), "multiplier"),
(retry(|c| c.multiplier = f64::NAN), "multiplier"),
(retry(|c| c.multiplier = f64::INFINITY), "multiplier"),
(retry(|c| c.multiplier = 1e9 + 1.0), "multiplier"),
(retry(|c| c.jitter = 1.5), "jitter"),
(retry(|c| c.jitter = -0.1), "jitter"),
(retry(|c| c.jitter = f64::NAN), "jitter"),
(retry(|c| c.initial = Duration::ZERO), "non-zero"),
(retry(|c| c.max = Duration::ZERO), "non-zero"),
(
retry(|c| {
c.initial = Duration::from_secs(10);
c.max = Duration::from_secs(1);
}),
"must not exceed",
),
];
for (cfg, needle) in cases {
let err = cfg
.validate()
.expect_err(&format!("{cfg:?} must fail"))
.to_string();
assert!(err.contains(needle), "expected `{needle}` in `{err}`");
}
}
#[test]
fn the_error_messages_survived_the_move_into_the_framework() {
assert_eq!(
RetryConfigError::Multiplier(0.5).to_string(),
"retry.multiplier must be a finite number in [1.0, 1e9] (got 0.5)"
);
assert_eq!(
RetryConfigError::Jitter(1.5).to_string(),
"retry.jitter must be a finite fraction in [0.0, 1.0] (got 1.5)"
);
assert_eq!(
RetryConfigError::ZeroDelay.to_string(),
"retry.initial and retry.max must be non-zero"
);
assert_eq!(
RetryConfigError::InitialExceedsMax {
initial: Duration::from_secs(10),
max: Duration::from_secs(1),
}
.to_string(),
"retry.initial (10s) must not exceed retry.max (1s)"
);
}
#[test]
fn only_unbounded_attempts_with_a_long_ceiling_count_as_a_stall() {
let long = Duration::from_secs(3600);
assert!(retry(|c| c.max = long).stalls_indefinitely());
assert!(
!retry(|c| {
c.max = long;
c.max_attempts = 5;
})
.stalls_indefinitely()
);
assert!(!RetryConfig::default().stalls_indefinitely());
assert!(!retry(|c| c.max = Duration::from_secs(300)).stalls_indefinitely());
assert!(retry(|c| c.max = Duration::from_secs(301)).stalls_indefinitely());
}
#[test]
fn validate_accepts_the_default_and_the_boundaries() {
assert!(RetryConfig::default().validate().is_ok());
let boundary = RetryConfig {
initial: Duration::from_nanos(1),
max: Duration::from_nanos(1),
multiplier: 1.0,
jitter: 0.0,
max_attempts: 0,
};
assert!(boundary.validate().is_ok(), "{boundary:?}");
let upper = RetryConfig {
multiplier: 1e9,
jitter: 1.0,
..RetryConfig::default()
};
assert!(upper.validate().is_ok(), "{upper:?}");
}
}