use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use async_trait::async_trait;
use tokio::sync::Notify;
use crate::InfraClientError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OperationProgress {
Pending,
Succeeded,
Failed,
Canceled,
}
#[derive(Debug, Clone)]
pub struct OperationObservation {
pub progress: OperationProgress,
pub next_poll_after: Option<Duration>,
pub error_code: Option<String>,
}
#[async_trait]
pub trait OperationPoller<T>: Send + Sync + fmt::Debug {
async fn poll(&self, operation_id: &str) -> Result<T, InfraClientError>;
async fn cancel(&self, operation_id: &str) -> Result<(), InfraClientError>;
fn observe(&self, operation: &T) -> OperationObservation;
}
#[derive(Debug, Clone)]
pub struct WaitOptions {
pub deadline: Duration,
pub minimum_poll_interval: Duration,
pub maximum_poll_interval: Duration,
pub cancellation: CancellationToken,
}
impl WaitOptions {
pub fn deadline(deadline: Duration) -> Self {
Self {
deadline,
..Self::default()
}
}
pub fn cancellation(mut self, cancellation: CancellationToken) -> Self {
self.cancellation = cancellation;
self
}
}
impl Default for WaitOptions {
fn default() -> Self {
Self {
deadline: Duration::from_secs(15 * 60),
minimum_poll_interval: Duration::from_millis(100),
maximum_poll_interval: Duration::from_secs(5),
cancellation: CancellationToken::new(),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct CancellationToken {
inner: Arc<CancellationState>,
}
#[derive(Debug, Default)]
struct CancellationState {
canceled: AtomicBool,
notify: Notify,
}
impl CancellationToken {
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
if !self.inner.canceled.swap(true, Ordering::AcqRel) {
self.inner.notify.notify_waiters();
}
}
pub fn is_canceled(&self) -> bool {
self.inner.canceled.load(Ordering::Acquire)
}
async fn canceled(&self) {
if self.is_canceled() {
return;
}
self.inner.notify.notified().await;
}
}
#[derive(Clone)]
pub struct OperationHandle<T> {
operation_id: Arc<str>,
initial: Option<T>,
poller: Arc<dyn OperationPoller<T>>,
}
impl<T> fmt::Debug for OperationHandle<T> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("OperationHandle")
.field("operation_id", &self.operation_id)
.finish_non_exhaustive()
}
}
impl<T> OperationHandle<T>
where
T: Clone + Send + Sync + 'static,
{
pub fn new(
operation_id: impl Into<Arc<str>>,
initial: Option<T>,
poller: Arc<dyn OperationPoller<T>>,
) -> Self {
Self {
operation_id: operation_id.into(),
initial,
poller,
}
}
pub fn id(&self) -> &str {
&self.operation_id
}
pub async fn refresh(&self) -> Result<T, InfraClientError> {
self.poller.poll(&self.operation_id).await
}
pub async fn cancel(&self) -> Result<(), InfraClientError> {
self.poller.cancel(&self.operation_id).await
}
pub async fn wait(mut self, options: WaitOptions) -> Result<T, InfraClientError> {
if options.deadline.is_zero()
|| options.deadline > Duration::from_secs(7 * 24 * 60 * 60)
|| options.minimum_poll_interval.is_zero()
|| options.maximum_poll_interval < options.minimum_poll_interval
|| options.maximum_poll_interval > Duration::from_secs(60)
{
return Err(InfraClientError::InvalidOptions {
message: "operation wait requires a deadline within 1ns..=7d and poll intervals within 1ns..=60s with minimum <= maximum".into(),
});
}
let started = Instant::now();
let mut current = match self.initial.take() {
Some(initial) => initial,
None => match tokio::time::timeout(options.deadline, self.refresh()).await {
Ok(result) => result?,
Err(_) => {
return Err(InfraClientError::DeadlineExceeded {
service: "operation",
});
}
},
};
let mut fallback_delay = options.minimum_poll_interval;
loop {
let observation = self.poller.observe(¤t);
match observation.progress {
OperationProgress::Succeeded => return Ok(current),
OperationProgress::Failed => {
return Err(InfraClientError::OperationTerminal {
operation_id: self.operation_id.to_string(),
state: "failed",
code: observation.error_code,
});
}
OperationProgress::Canceled => {
return Err(InfraClientError::OperationTerminal {
operation_id: self.operation_id.to_string(),
state: "canceled",
code: observation.error_code,
});
}
OperationProgress::Pending => {}
}
if started.elapsed() >= options.deadline {
return Err(InfraClientError::DeadlineExceeded {
service: "operation",
});
}
if options.cancellation.is_canceled() {
self.cancel_bounded(Duration::from_secs(5)).await?;
return Err(InfraClientError::Canceled {
operation_id: self.operation_id.to_string(),
});
}
let remaining = options.deadline.saturating_sub(started.elapsed());
let delay = observation
.next_poll_after
.unwrap_or(fallback_delay)
.clamp(options.minimum_poll_interval, options.maximum_poll_interval)
.min(remaining);
tokio::select! {
() = tokio::time::sleep(delay) => {},
() = options.cancellation.canceled() => {
self.cancel_bounded(Duration::from_secs(5).min(remaining)).await?;
return Err(InfraClientError::Canceled {
operation_id: self.operation_id.to_string(),
});
}
}
fallback_delay = fallback_delay
.saturating_mul(2)
.min(options.maximum_poll_interval);
let remaining = options.deadline.saturating_sub(started.elapsed());
current = tokio::select! {
result = tokio::time::timeout(remaining, self.refresh()) => match result {
Ok(result) => result?,
Err(_) => return Err(InfraClientError::DeadlineExceeded { service: "operation" }),
},
() = options.cancellation.canceled() => {
self.cancel_bounded(remaining.min(Duration::from_secs(5))).await?;
return Err(InfraClientError::Canceled {
operation_id: self.operation_id.to_string(),
});
}
};
}
}
async fn cancel_bounded(&self, budget: Duration) -> Result<(), InfraClientError> {
match tokio::time::timeout(budget.max(Duration::from_millis(1)), self.cancel()).await {
Ok(result) => result,
Err(_) => Err(InfraClientError::DeadlineExceeded {
service: "operation-cancel",
}),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
#[derive(Debug)]
struct Poller {
polls: AtomicUsize,
canceled: AtomicBool,
}
#[async_trait]
impl OperationPoller<usize> for Poller {
async fn poll(&self, _operation_id: &str) -> Result<usize, InfraClientError> {
Ok(self.polls.fetch_add(1, Ordering::Relaxed) + 1)
}
async fn cancel(&self, _operation_id: &str) -> Result<(), InfraClientError> {
self.canceled.store(true, Ordering::Relaxed);
Ok(())
}
fn observe(&self, operation: &usize) -> OperationObservation {
OperationObservation {
progress: if *operation >= 2 {
OperationProgress::Succeeded
} else {
OperationProgress::Pending
},
next_poll_after: Some(Duration::from_millis(1)),
error_code: None,
}
}
}
#[tokio::test]
async fn wait_polls_to_terminal_without_collecting_unbounded_state() {
let poller = Arc::new(Poller {
polls: AtomicUsize::new(0),
canceled: AtomicBool::new(false),
});
let result = OperationHandle::new("op", Some(0), poller)
.wait(WaitOptions {
deadline: Duration::from_secs(1),
minimum_poll_interval: Duration::from_millis(1),
maximum_poll_interval: Duration::from_millis(2),
cancellation: CancellationToken::new(),
})
.await
.unwrap();
assert_eq!(result, 2);
}
#[tokio::test]
async fn cancellation_propagates_to_the_remote_operation() {
let poller = Arc::new(Poller {
polls: AtomicUsize::new(0),
canceled: AtomicBool::new(false),
});
let cancellation = CancellationToken::new();
cancellation.cancel();
let error = OperationHandle::new("op-cancel", Some(0), poller.clone())
.wait(WaitOptions {
deadline: Duration::from_secs(1),
minimum_poll_interval: Duration::from_millis(1),
maximum_poll_interval: Duration::from_millis(2),
cancellation,
})
.await
.unwrap_err();
assert!(matches!(error, InfraClientError::Canceled { .. }));
assert!(poller.canceled.load(Ordering::Relaxed));
}
#[derive(Debug)]
struct SlowPoller;
#[async_trait]
impl OperationPoller<usize> for SlowPoller {
async fn poll(&self, _operation_id: &str) -> Result<usize, InfraClientError> {
tokio::time::sleep(Duration::from_millis(100)).await;
Ok(0)
}
async fn cancel(&self, _operation_id: &str) -> Result<(), InfraClientError> {
Ok(())
}
fn observe(&self, _operation: &usize) -> OperationObservation {
OperationObservation {
progress: OperationProgress::Pending,
next_poll_after: None,
error_code: None,
}
}
}
#[tokio::test]
async fn deadline_bounds_the_initial_poll() {
let error = OperationHandle::new("op-timeout", None, Arc::new(SlowPoller))
.wait(WaitOptions::deadline(Duration::from_millis(5)))
.await
.unwrap_err();
assert!(matches!(
error,
InfraClientError::DeadlineExceeded {
service: "operation"
}
));
}
#[tokio::test]
async fn invalid_poll_intervals_return_error_instead_of_panicking_or_spinning() {
let error = OperationHandle::new("op-invalid", Some(0), Arc::new(SlowPoller))
.wait(WaitOptions {
deadline: Duration::from_secs(1),
minimum_poll_interval: Duration::from_secs(2),
maximum_poll_interval: Duration::from_secs(1),
cancellation: CancellationToken::new(),
})
.await
.unwrap_err();
assert!(matches!(error, InfraClientError::InvalidOptions { .. }));
}
#[derive(Debug)]
struct FailedPoller;
#[async_trait]
impl OperationPoller<usize> for FailedPoller {
async fn poll(&self, _operation_id: &str) -> Result<usize, InfraClientError> {
Ok(0)
}
async fn cancel(&self, _operation_id: &str) -> Result<(), InfraClientError> {
Ok(())
}
fn observe(&self, _operation: &usize) -> OperationObservation {
OperationObservation {
progress: OperationProgress::Failed,
next_poll_after: None,
error_code: Some("PROVIDER_UNAVAILABLE".into()),
}
}
}
#[tokio::test]
async fn terminal_operation_preserves_the_stable_remote_error_code() {
let error = OperationHandle::new("op-failed", Some(0), Arc::new(FailedPoller))
.wait(WaitOptions::default())
.await
.unwrap_err();
assert_eq!(error.code(), "PROVIDER_UNAVAILABLE");
assert!(error.to_string().contains("PROVIDER_UNAVAILABLE"));
}
}