#![forbid(unsafe_code)]
#![deny(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
use std::future::Future;
use std::pin::Pin;
use dynamic_config::{AsyncRemoteSource, Error, Fetched, Format};
use etcd_client::EventType;
pub use etcd_client::{Client, ConnectOptions};
#[cfg(feature = "tls")]
#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
pub use etcd_client::{Certificate, Identity, TlsOptions};
const INVALID_TOKEN: &str = "invalid auth token";
use tokio::sync::Mutex;
pub struct Etcd {
client: Mutex<Client>,
key: String,
format: Option<Format>,
endpoints: String,
}
impl Etcd {
pub async fn new<E, S>(endpoints: E, key: impl Into<String>) -> Result<Self, Error>
where
E: IntoIterator<Item = S>,
S: Into<String>,
{
Self::with_options(endpoints, key, ConnectOptions::new()).await
}
pub async fn with_options<E, S>(
endpoints: E,
key: impl Into<String>,
options: ConnectOptions,
) -> Result<Self, Error>
where
E: IntoIterator<Item = S>,
S: Into<String>,
{
let endpoints: Vec<String> = endpoints.into_iter().map(Into::into).collect();
let described = endpoints.join(", ");
let client = connect(&endpoints, &options, &described).await?;
Ok(Self {
client: Mutex::new(client),
key: key.into(),
format: None,
endpoints: described,
}
.with_format_from_key())
}
#[must_use]
pub fn from_client(client: Client, key: impl Into<String>) -> Self {
Self {
client: Mutex::new(client),
key: key.into(),
format: None,
endpoints: "<an existing client>".to_owned(),
}
.with_format_from_key()
}
fn with_format_from_key(mut self) -> Self {
self.format = Format::from_key(&self.key);
self
}
#[must_use]
pub fn with_format(mut self, format: Format) -> Self {
self.format = Some(format);
self
}
pub async fn watch<F>(&self, mut on_change: F) -> Result<(), Error>
where
F: FnMut(Fetched) -> Result<(), Error> + Send,
{
let format = self.format.ok_or_else(|| {
Error::remote(format!(
"{}: the key names no format; call `with_format`",
self.describe()
))
})?;
let mut stream = match self.watch_once(None).await {
Err(error) if is_expired_token(&error) => {
self.refresh_token().await?;
self.watch_once(None).await?
}
outcome => outcome?,
};
const MOST_TOKEN_RECOVERIES: u32 = 3;
let mut token_recoveries = 0_u32;
let mut resume_from: Option<i64> = None;
loop {
let response = match stream.message().await {
Ok(Some(response)) => {
token_recoveries = 0;
if let Some(header) = response.header() {
resume_from = Some(header.revision() + 1);
}
response
}
Ok(None) => break,
Err(error) => {
let wrapped =
Error::remote(format!("{}: the watch failed: {error}", self.describe()));
if is_expired_token(&wrapped) {
token_recoveries += 1;
if token_recoveries > MOST_TOKEN_RECOVERIES {
return Err(wrapped);
}
self.refresh_token().await?;
stream = self.watch_once(resume_from).await?;
continue;
}
return Err(wrapped);
}
};
if response.canceled() {
return Err(Error::remote(format!(
"{}: the store cancelled the watch: {}",
self.describe(),
response.cancel_reason()
)));
}
for event in response.events() {
if event.event_type() != EventType::Put {
continue;
}
let Some(value) = event.kv() else { continue };
let text = value.value_str().map_err(|error| {
Error::remote(format!(
"{}: the value is not UTF-8: {error}",
self.describe()
))
})?;
guarded(&mut on_change, Fetched::new(text, format), &self.describe())?;
}
}
Err(Error::remote(format!(
"{}: the watch ended; the connection was closed",
self.describe()
)))
}
async fn refresh_token(&self) -> Result<(), Error> {
self.client
.lock()
.await
.refresh_token()
.await
.map_err(|error| {
Error::remote(format!(
"{}: the auth token expired and could not be replaced: {error}",
self.describe()
))
})
}
}
impl Etcd {
async fn watch_once(
&self,
from_revision: Option<i64>,
) -> Result<etcd_client::WatchStream, Error> {
let options = from_revision
.map(|revision| etcd_client::WatchOptions::new().with_start_revision(revision));
self.client
.lock()
.await
.watch(self.key.as_str(), options)
.await
.map_err(|error| Error::remote(format!("{}: cannot watch: {error}", self.describe())))
}
async fn get_once(&self) -> Result<etcd_client::GetResponse, Error> {
self.client
.lock()
.await
.get(self.key.as_str(), None)
.await
.map_err(|error| Error::remote(format!("{}: {error}", self.describe())))
}
}
fn is_expired_token(error: &Error) -> bool {
error.to_string().contains(INVALID_TOKEN)
}
async fn connect(
endpoints: &[String],
options: &ConnectOptions,
described: &str,
) -> Result<Client, Error> {
Client::connect(endpoints, Some(options.clone()))
.await
.map_err(|error| Error::remote(format!("etcd {described}: {error}")))
}
impl AsyncRemoteSource for Etcd {
fn fetch(&self) -> Pin<Box<dyn Future<Output = Result<Fetched, Error>> + Send + '_>> {
Box::pin(async move {
let format = self.format.ok_or_else(|| {
Error::remote(format!(
"{}: the key names no format; call `with_format`",
self.describe()
))
})?;
let response = match self.get_once().await {
Err(error) if is_expired_token(&error) => {
self.refresh_token().await?;
self.get_once().await?
}
outcome => outcome?,
};
let value = response.kvs().first().ok_or_else(|| {
Error::remote(format!("{}: the key holds no value", self.describe()))
})?;
let text = value.value_str().map_err(|error| {
Error::remote(format!(
"{}: the value is not UTF-8: {error}",
self.describe()
))
})?;
Ok(Fetched::new(text, format))
})
}
fn describe(&self) -> String {
format!("etcd {} key {}", self.endpoints, self.key)
}
}
impl std::fmt::Debug for Etcd {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Etcd")
.field("endpoints", &self.endpoints)
.field("key", &self.key)
.field("format", &self.format)
.finish_non_exhaustive()
}
}
fn guarded<F>(on_change: &mut F, document: Fetched, described: &str) -> Result<(), Error>
where
F: FnMut(Fetched) -> Result<(), Error>,
{
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| on_change(document))).unwrap_or_else(
|_| {
Err(Error::remote(format!(
"{described}: the watch callback panicked; the watch is stopped"
)))
},
)
}