use std::time::Duration;
use tonic::{Code, Status};
use crate::proto::udb::entity::v1::ErrorDetail;
pub const ERROR_DETAIL_TRAILER: &str = "udb-error-detail-bin";
#[derive(Debug, Clone)]
pub struct UdbError {
pub status: Status,
pub detail: Option<ErrorDetail>,
}
impl UdbError {
pub fn from_status(status: Status) -> Self {
let detail = status
.metadata()
.get_bin(ERROR_DETAIL_TRAILER)
.and_then(|v| v.to_bytes().ok())
.and_then(|bytes| <ErrorDetail as prost::Message>::decode(bytes).ok());
Self { status, detail }
}
pub fn code(&self) -> Code {
self.status.code()
}
pub fn message(&self) -> &str {
self.status.message()
}
pub fn is_retryable(&self) -> bool {
match &self.detail {
Some(d) => d.retryable,
None => matches!(
self.status.code(),
Code::Unavailable | Code::ResourceExhausted
),
}
}
pub fn retry_after(&self) -> Option<Duration> {
let ms = self.detail.as_ref()?.retry_after_ms;
(ms > 0).then(|| Duration::from_millis(ms as u64))
}
pub fn capability_required(&self) -> Option<&str> {
let c = self.detail.as_ref()?.capability_required.as_str();
(!c.is_empty()).then_some(c)
}
pub fn correlation_id(&self) -> Option<&str> {
let c = self.detail.as_ref()?.correlation_id.as_str();
(!c.is_empty()).then_some(c)
}
pub fn field_violations(&self) -> &[crate::proto::udb::entity::v1::ErrorFieldViolation] {
self.detail
.as_ref()
.map(|d| d.field_violations.as_slice())
.unwrap_or_default()
}
}
impl std::fmt::Display for UdbError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}: {}", self.status.code(), self.status.message())?;
if let Some(detail) = &self.detail {
if !detail.backend.is_empty() {
write!(f, " [backend={}", detail.backend)?;
if !detail.operation.is_empty() {
write!(f, " op={}", detail.operation)?;
}
write!(f, "]")?;
}
if !detail.capability_required.is_empty() {
write!(f, " (requires {})", detail.capability_required)?;
}
}
Ok(())
}
}
impl std::error::Error for UdbError {}
impl From<Status> for UdbError {
fn from(status: Status) -> Self {
Self::from_status(status)
}
}
#[derive(Debug, Clone, Copy)]
pub struct CallPolicy {
pub deadline: Option<Duration>,
pub max_attempts: u32,
pub base_backoff: Duration,
pub idempotent: bool,
}
impl Default for CallPolicy {
fn default() -> Self {
Self {
deadline: Some(Duration::from_secs(30)),
max_attempts: 3,
base_backoff: Duration::from_millis(100),
idempotent: false,
}
}
}
impl CallPolicy {
pub fn idempotent() -> Self {
Self {
idempotent: true,
..Default::default()
}
}
pub fn from_contract(path: &str) -> Self {
if crate::generated_rpcs::is_retry_safe(path) {
Self::idempotent()
} else {
Self::default()
}
}
pub fn once() -> Self {
Self {
max_attempts: 1,
..Default::default()
}
}
pub fn with_deadline(mut self, deadline: Duration) -> Self {
self.deadline = Some(deadline);
self
}
pub fn with_max_attempts(mut self, attempts: u32) -> Self {
self.max_attempts = attempts.max(1);
self
}
pub(crate) fn backoff_for(&self, attempt: u32, err: &UdbError) -> Duration {
if let Some(after) = err.retry_after() {
return after;
}
let exp = self
.base_backoff
.saturating_mul(1u32 << attempt.min(6).saturating_sub(1));
exp.min(Duration::from_secs(5))
}
pub(crate) fn should_retry(&self, attempt: u32, err: &UdbError) -> bool {
self.idempotent && attempt < self.max_attempts && err.is_retryable()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn status_without_detail(code: Code) -> UdbError {
UdbError::from_status(Status::new(code, "boom"))
}
#[test]
fn missing_trailer_falls_back_to_safe_codes() {
assert!(status_without_detail(Code::Unavailable).is_retryable());
assert!(status_without_detail(Code::ResourceExhausted).is_retryable());
assert!(
!status_without_detail(Code::InvalidArgument).is_retryable(),
"a rejected argument will be rejected again"
);
assert!(
!status_without_detail(Code::Internal).is_retryable(),
"INTERNAL may have applied a mutation; never assume repeatable"
);
}
#[test]
fn malformed_trailer_is_treated_as_absent() {
let mut status = Status::new(Code::Internal, "boom");
status.metadata_mut().insert_bin(
ERROR_DETAIL_TRAILER,
tonic::metadata::MetadataValue::from_bytes(b"not-a-protobuf-at-all"),
);
let err = UdbError::from_status(status);
assert_eq!(err.code(), Code::Internal);
assert_eq!(err.message(), "boom");
}
#[test]
fn broker_detail_overrides_the_code_guess() {
let detail = ErrorDetail {
retryable: true,
retry_after_ms: 250,
capability_required: "postgres_backend".into(),
correlation_id: "corr-9".into(),
..Default::default()
};
let mut status = Status::new(Code::Internal, "boom");
let mut buf = Vec::new();
prost::Message::encode(&detail, &mut buf).expect("encode");
status.metadata_mut().insert_bin(
ERROR_DETAIL_TRAILER,
tonic::metadata::MetadataValue::from_bytes(&buf),
);
let err = UdbError::from_status(status);
assert!(
err.is_retryable(),
"INTERNAL, but the broker said retryable"
);
assert_eq!(err.retry_after(), Some(Duration::from_millis(250)));
assert_eq!(err.capability_required(), Some("postgres_backend"));
assert_eq!(err.correlation_id(), Some("corr-9"));
}
#[test]
fn mutations_are_not_retried_by_default() {
let err = status_without_detail(Code::Unavailable);
assert!(
!CallPolicy::default().should_retry(1, &err),
"default policy must not repeat a possible mutation"
);
assert!(CallPolicy::idempotent().should_retry(1, &err));
}
#[test]
fn retry_stops_at_max_attempts() {
let err = status_without_detail(Code::Unavailable);
let policy = CallPolicy::idempotent().with_max_attempts(2);
assert!(policy.should_retry(1, &err));
assert!(!policy.should_retry(2, &err), "attempt 2 of 2 is the last");
}
#[test]
fn backoff_prefers_the_brokers_retry_after() {
let detail = ErrorDetail {
retryable: true,
retry_after_ms: 1234,
..Default::default()
};
let err = UdbError {
status: Status::new(Code::Unavailable, "boom"),
detail: Some(detail),
};
assert_eq!(
CallPolicy::idempotent().backoff_for(3, &err),
Duration::from_millis(1234)
);
}
#[test]
fn backoff_grows_and_is_capped() {
let err = status_without_detail(Code::Unavailable);
let p = CallPolicy::idempotent();
assert_eq!(p.backoff_for(1, &err), Duration::from_millis(100));
assert_eq!(p.backoff_for(2, &err), Duration::from_millis(200));
assert_eq!(p.backoff_for(3, &err), Duration::from_millis(400));
assert!(p.backoff_for(20, &err) <= Duration::from_secs(5));
}
}