use std::future::Future;
use std::time::Duration;
use ferrin_spec::Headers;
use ferrin_spec::ProviderOptions;
use tokio_util::sync::CancellationToken;
use crate::USER_AGENT;
use crate::cancel::CallCancellation;
use crate::error::Error;
use crate::retry::RetryPolicy;
use crate::telemetry::TelemetryOptions;
use crate::timeout::TimeoutScope;
#[derive(Debug, Clone)]
pub(crate) struct ModalityOptions {
pub(crate) headers: Headers,
pub(crate) provider_options: ProviderOptions,
pub(crate) retry_policy: RetryPolicy,
pub(crate) cancellation: CancellationToken,
pub(crate) timeout: Option<Duration>,
pub(crate) telemetry: TelemetryOptions,
}
impl Default for ModalityOptions {
fn default() -> Self {
Self {
headers: Headers::new(),
provider_options: ProviderOptions::new(),
retry_policy: RetryPolicy::default(),
cancellation: CancellationToken::new(),
timeout: None,
telemetry: TelemetryOptions::default(),
}
}
}
impl ModalityOptions {
pub(crate) fn request_headers(&self) -> Headers {
self.headers.clone().with_user_agent_suffix([USER_AGENT])
}
pub(crate) async fn run<T, F, Fut>(self, operation: F) -> Result<T, Error>
where
F: FnOnce(Self, CancellationToken) -> Fut,
Fut: Future<Output = Result<T, Error>>,
{
let cancellation = CallCancellation::new(&self.cancellation);
let timeout = self.timeout;
let token = cancellation.token().clone();
cancellation
.with_timeout(TimeoutScope::Total, timeout, operation(self, token))
.await
.map_err(|error| cancellation.map_error(error))
}
}
macro_rules! impl_modality_builder {
($ty:ident $(< $($generic:ident),+ >)?) => {
$crate::modality::impl_modality_builder!(@common $ty $(< $($generic),+ >)?);
impl$(<$($generic),+>)? $ty$(<$($generic),+>)? {
#[must_use]
pub fn retry(mut self, policy: $crate::retry::RetryPolicy) -> Self {
self.base.retry_policy = policy;
self
}
#[must_use]
pub fn max_retries(mut self, max_retries: u32) -> Self {
self.base.retry_policy.max_retries = max_retries;
self
}
}
};
(@no_retry $ty:ident $(< $($generic:ident),+ >)?) => {
$crate::modality::impl_modality_builder!(@common $ty $(< $($generic),+ >)?);
};
(@common $ty:ident $(< $($generic:ident),+ >)?) => {
impl$(<$($generic),+>)? $ty$(<$($generic),+>)? {
#[must_use]
pub fn headers(mut self, headers: ::ferrin_spec::Headers) -> Self {
self.base.headers.merge(&headers);
self
}
#[must_use]
pub fn header(mut self, name: &str, value: &str) -> Self {
let headers = ::std::mem::take(&mut self.base.headers);
self.base.headers = headers.with(name, value);
self
}
#[must_use]
pub fn provider_options(mut self, options: ::ferrin_spec::ProviderOptions) -> Self {
self.base.provider_options = options;
self
}
#[must_use]
pub fn provider_option(
mut self,
provider: impl Into<String>,
options: ::ferrin_spec::JsonObject,
) -> Self {
self.base.provider_options.insert(provider.into(), options);
self
}
#[must_use]
pub fn cancellation(mut self, token: ::tokio_util::sync::CancellationToken) -> Self {
self.base.cancellation = token;
self
}
#[must_use]
pub fn timeout(mut self, timeout: ::std::time::Duration) -> Self {
self.base.timeout = Some(timeout);
self
}
#[must_use]
pub fn telemetry(mut self, options: $crate::telemetry::TelemetryOptions) -> Self {
self.base.telemetry = options;
self
}
}
};
}
pub(crate) use impl_modality_builder;
pub(crate) fn add_optional(a: Option<u64>, b: Option<u64>) -> Option<u64> {
match (a, b) {
(Some(a), Some(b)) => Some(a.saturating_add(b)),
(Some(value), None) | (None, Some(value)) => Some(value),
(None, None) => None,
}
}