use crate::{Error, Result, RetryHint};
use asupersync::{time, Cx};
use std::future::Future;
use std::time::Duration;
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Idempotency {
Idempotent,
NonIdempotent,
}
impl Idempotency {
pub fn classify_sql(sql: &str) -> Idempotency {
let head = sql.trim_start();
let keyword = head
.split(|ch: char| !ch.is_ascii_alphabetic())
.next()
.unwrap_or("");
if keyword.eq_ignore_ascii_case("select") {
Idempotency::Idempotent
} else {
Idempotency::NonIdempotent
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RetryAction {
Surface,
RetrySameConnection { backoff: Duration },
ReconnectThenRetry { backoff: Duration },
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetryPolicy {
pub max_retries: u32,
pub base_backoff: Duration,
pub max_backoff: Duration,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
base_backoff: Duration::from_millis(50),
max_backoff: Duration::from_secs(2),
}
}
}
impl RetryPolicy {
pub const fn none() -> Self {
Self {
max_retries: 0,
base_backoff: Duration::ZERO,
max_backoff: Duration::ZERO,
}
}
fn backoff_for(&self, attempts_made: u32) -> Duration {
if self.base_backoff.is_zero() {
return Duration::ZERO;
}
let shift = attempts_made.saturating_sub(1).min(31);
let scaled = self
.base_backoff
.checked_mul(1u32 << shift)
.unwrap_or(self.max_backoff);
scaled.min(self.max_backoff)
}
pub fn decide(&self, err: &Error, idempotency: Idempotency, attempts_made: u32) -> RetryAction {
if idempotency == Idempotency::NonIdempotent {
return RetryAction::Surface;
}
if attempts_made > self.max_retries {
return RetryAction::Surface;
}
let backoff = self.backoff_for(attempts_made);
match err.retry_hint() {
RetryHint::Never => RetryAction::Surface,
RetryHint::RetrySameConnectionIfIdempotent => {
RetryAction::RetrySameConnection { backoff }
}
RetryHint::ReconnectThenRetryIfIdempotent => {
RetryAction::ReconnectThenRetry { backoff }
}
}
}
}
async fn backoff_sleep(cx: &Cx, backoff: Duration) {
if backoff.is_zero() || cx.cancel_reason().is_some() {
return;
}
time::sleep(time::wall_now(), backoff).await;
}
pub async fn run_with_retry<T, MkFut, Fut>(
cx: &Cx,
policy: &RetryPolicy,
idempotency: Idempotency,
op: MkFut,
) -> Result<T>
where
MkFut: FnMut() -> Fut,
Fut: Future<Output = Result<T>>,
{
run_inner(cx, policy, idempotency, op, NoReconnect).await
}
pub async fn run_with_retry_reconnecting<T, MkFut, Fut, Recon, ReconFut>(
cx: &Cx,
policy: &RetryPolicy,
idempotency: Idempotency,
op: MkFut,
reconnect: Recon,
) -> Result<T>
where
MkFut: FnMut() -> Fut,
Fut: Future<Output = Result<T>>,
Recon: FnMut() -> ReconFut,
ReconFut: Future<Output = Result<()>>,
{
run_inner(cx, policy, idempotency, op, Reconnect(reconnect)).await
}
trait ReconnectHook {
fn reconnect(&mut self) -> impl Future<Output = Result<bool>>;
}
struct NoReconnect;
impl ReconnectHook for NoReconnect {
async fn reconnect(&mut self) -> Result<bool> {
Ok(false)
}
}
struct Reconnect<R>(R);
impl<R, Fut> ReconnectHook for Reconnect<R>
where
R: FnMut() -> Fut,
Fut: Future<Output = Result<()>>,
{
async fn reconnect(&mut self) -> Result<bool> {
(self.0)().await.map(|()| true)
}
}
async fn run_inner<T, MkFut, Fut, H>(
cx: &Cx,
policy: &RetryPolicy,
idempotency: Idempotency,
mut op: MkFut,
mut hook: H,
) -> Result<T>
where
MkFut: FnMut() -> Fut,
Fut: Future<Output = Result<T>>,
H: ReconnectHook,
{
let mut attempts_made: u32 = 0;
loop {
match op().await {
Ok(value) => return Ok(value),
Err(err) => {
attempts_made += 1;
match policy.decide(&err, idempotency, attempts_made) {
RetryAction::Surface => return Err(err),
RetryAction::RetrySameConnection { backoff } => {
backoff_sleep(cx, backoff).await;
}
RetryAction::ReconnectThenRetry { backoff } => {
backoff_sleep(cx, backoff).await;
match hook.reconnect().await {
Ok(true) => {}
Ok(false) => return Err(err),
Err(reconnect_err) => return Err(reconnect_err),
}
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use asupersync::runtime::{reactor, RuntimeBuilder};
use std::cell::Cell;
fn structured_error(code: u32) -> Error {
Error::Protocol(oracledb_protocol::ProtocolError::ServerErrorInfo(Box::new(
oracledb_protocol::ServerErrorDetails {
message: format!("ORA-{code:05}: synthetic"),
code,
..Default::default()
},
)))
}
fn transient() -> Error {
structured_error(54)
}
fn connection_lost() -> Error {
structured_error(3113)
}
fn permanent() -> Error {
structured_error(942)
}
fn instant_policy(max_retries: u32) -> RetryPolicy {
RetryPolicy {
max_retries,
base_backoff: Duration::ZERO,
max_backoff: Duration::ZERO,
}
}
#[test]
fn transient_idempotent_retries_within_budget() {
let p = instant_policy(3);
assert_eq!(
p.decide(&transient(), Idempotency::Idempotent, 1),
RetryAction::RetrySameConnection {
backoff: Duration::ZERO
}
);
}
#[test]
fn transient_non_idempotent_always_surfaces() {
let p = instant_policy(3);
assert_eq!(
p.decide(&transient(), Idempotency::NonIdempotent, 1),
RetryAction::Surface
);
}
#[test]
fn connection_lost_idempotent_asks_for_reconnect() {
let p = instant_policy(3);
assert_eq!(
p.decide(&connection_lost(), Idempotency::Idempotent, 1),
RetryAction::ReconnectThenRetry {
backoff: Duration::ZERO
}
);
}
#[test]
fn connection_lost_non_idempotent_surfaces() {
let p = instant_policy(3);
assert_eq!(
p.decide(&connection_lost(), Idempotency::NonIdempotent, 1),
RetryAction::Surface
);
}
#[test]
fn permanent_never_retries_even_when_idempotent() {
let p = instant_policy(3);
assert_eq!(
p.decide(&permanent(), Idempotency::Idempotent, 1),
RetryAction::Surface
);
}
#[test]
fn budget_exhaustion_surfaces() {
let p = instant_policy(2);
assert!(matches!(
p.decide(&transient(), Idempotency::Idempotent, 2),
RetryAction::RetrySameConnection { .. }
));
assert_eq!(
p.decide(&transient(), Idempotency::Idempotent, 3),
RetryAction::Surface
);
}
#[test]
fn zero_retry_policy_surfaces_first_failure() {
let p = instant_policy(0);
assert_eq!(
p.decide(&transient(), Idempotency::Idempotent, 1),
RetryAction::Surface
);
}
#[test]
fn backoff_is_exponential_and_capped() {
let p = RetryPolicy {
max_retries: 10,
base_backoff: Duration::from_millis(10),
max_backoff: Duration::from_millis(50),
};
assert_eq!(p.backoff_for(1), Duration::from_millis(10));
assert_eq!(p.backoff_for(2), Duration::from_millis(20));
assert_eq!(p.backoff_for(3), Duration::from_millis(40));
assert_eq!(p.backoff_for(4), Duration::from_millis(50));
assert_eq!(p.backoff_for(1000), Duration::from_millis(50));
}
#[test]
fn classify_sql_only_select_is_idempotent() {
assert_eq!(
Idempotency::classify_sql("SELECT * FROM dual"),
Idempotency::Idempotent
);
assert_eq!(
Idempotency::classify_sql(" select 1 from dual"),
Idempotency::Idempotent
);
assert_eq!(
Idempotency::classify_sql("INSERT INTO t VALUES (1)"),
Idempotency::NonIdempotent
);
assert_eq!(
Idempotency::classify_sql("UPDATE t SET x = 1"),
Idempotency::NonIdempotent
);
assert_eq!(
Idempotency::classify_sql("BEGIN proc(); END;"),
Idempotency::NonIdempotent
);
assert_eq!(
Idempotency::classify_sql("MERGE INTO t USING s ON (t.k=s.k)"),
Idempotency::NonIdempotent
);
assert_eq!(
Idempotency::classify_sql("WITH q AS (SELECT 1) SELECT * FROM q"),
Idempotency::NonIdempotent
);
assert_eq!(Idempotency::classify_sql(""), Idempotency::NonIdempotent);
}
struct ScriptedOp {
script: Vec<std::result::Result<u64, Error>>,
next: Cell<usize>,
calls: Cell<usize>,
}
impl ScriptedOp {
fn new(script: Vec<std::result::Result<u64, Error>>) -> Self {
Self {
script,
next: Cell::new(0),
calls: Cell::new(0),
}
}
fn run_once(&self) -> Result<u64> {
let i = self.next.get();
self.calls.set(self.calls.get() + 1);
self.next.set(i + 1);
match self.script.get(i) {
Some(Ok(v)) => Ok(*v),
Some(Err(e)) => Err(clone_error(e)),
None => panic!("scripted op called more times than scripted"),
}
}
fn calls(&self) -> usize {
self.calls.get()
}
}
fn clone_error(e: &Error) -> Error {
match e.ora_code() {
Some(code) => structured_error(code as u32),
None => Error::Runtime("scripted".into()),
}
}
fn block_on<F: Future<Output = ()>>(fut: F) {
let reactor = reactor::create_reactor().expect("reactor");
let runtime = RuntimeBuilder::current_thread()
.with_reactor(reactor)
.build()
.expect("runtime");
runtime.block_on(fut);
}
#[test]
fn idempotent_transient_then_success_retries() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![Err(transient()), Ok(42)]);
let out = run_with_retry(&cx, &instant_policy(3), Idempotency::Idempotent, || async {
op.run_once()
})
.await;
assert_eq!(out.unwrap(), 42);
assert_eq!(op.calls(), 2, "one retry after the transient failure");
});
}
#[test]
fn non_idempotent_transient_is_surfaced_not_rerun() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![Err(transient()), Ok(99)]);
let out = run_with_retry(
&cx,
&instant_policy(3),
Idempotency::NonIdempotent,
|| async { op.run_once() },
)
.await;
assert!(out.is_err(), "non-idempotent op surfaces the failure");
assert_eq!(out.unwrap_err().ora_code(), Some(54));
assert_eq!(op.calls(), 1, "non-idempotent op must NOT be re-run");
});
}
#[test]
fn idempotent_exhausts_budget_then_surfaces() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![
Err(transient()),
Err(transient()),
Err(transient()),
Err(transient()),
]);
let out = run_with_retry(&cx, &instant_policy(2), Idempotency::Idempotent, || async {
op.run_once()
})
.await;
assert!(out.is_err());
assert_eq!(op.calls(), 3);
});
}
#[test]
fn connection_lost_without_hook_surfaces_on_same_conn_path() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![Err(connection_lost()), Ok(7)]);
let out = run_with_retry(&cx, &instant_policy(3), Idempotency::Idempotent, || async {
op.run_once()
})
.await;
assert!(out.is_err(), "same-conn path cannot reconnect");
assert_eq!(op.calls(), 1);
});
}
#[test]
fn connection_lost_reconnects_then_retries() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![Err(connection_lost()), Ok(7)]);
let reconnects = Cell::new(0usize);
let out = run_with_retry_reconnecting(
&cx,
&instant_policy(3),
Idempotency::Idempotent,
|| async { op.run_once() },
|| async {
reconnects.set(reconnects.get() + 1);
Ok(())
},
)
.await;
assert_eq!(out.unwrap(), 7);
assert_eq!(op.calls(), 2);
assert_eq!(reconnects.get(), 1, "reconnected once before the retry");
});
}
#[test]
fn reconnect_failure_is_surfaced() {
block_on(async {
let cx = Cx::current().expect("cx");
let op = ScriptedOp::new(vec![Err(connection_lost()), Ok(7)]);
let out = run_with_retry_reconnecting(
&cx,
&instant_policy(3),
Idempotency::Idempotent,
|| async { op.run_once() },
|| async { Err(Error::Runtime("reconnect refused".into())) },
)
.await;
let err = out.unwrap_err();
assert!(matches!(err, Error::Runtime(msg) if msg.contains("reconnect refused")));
assert_eq!(op.calls(), 1);
});
}
}