rust_zero_core/
etcd.rs

1//! Etcd-backed configuration and service discovery.
2//!
3//! ```no_run
4//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
5//! use rust_zero_core::{EtcdClient, EtcdConfig};
6//! use std::time::Duration;
7//! let client = EtcdClient::connect(EtcdConfig::new(["http://127.0.0.1:2379"])).await?;
8//! let _lease = client
9//!     .publish("users", "users-1", "http://127.0.0.1:8080", Duration::from_secs(10))
10//!     .await?;
11//! let subscription = client.subscribe("users").await?;
12//! assert!(!subscription.endpoints().is_empty());
13//! # Ok(())
14//! # }
15//! ```
16
17use crate::{
18    ConfigFormat, DiscoveryReconnectBackoff, DynamicConfig, EndpointChangeFuture,
19    EndpointSubscription,
20};
21use etcd_client::{Client, ConnectOptions, EventType, GetOptions, PutOptions, WatchOptions};
22use serde::de::DeserializeOwned;
23use std::{
24    collections::{hash_map::DefaultHasher, BTreeMap},
25    error::Error,
26    fmt,
27    hash::{Hash, Hasher},
28    sync::Arc,
29    time::{Duration, SystemTime, UNIX_EPOCH},
30};
31use tokio::{
32    sync::{oneshot, watch},
33    task::JoinHandle,
34};
35
36/// Connection and namespace settings for the etcd configuration and discovery adapter.
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct EtcdConfig {
39    pub endpoints: Vec<String>,
40    pub namespace: String,
41    pub username: Option<String>,
42    pub password: Option<String>,
43    pub timeout: Duration,
44    pub reconnect_backoff: DiscoveryReconnectBackoff,
45}
46
47impl EtcdConfig {
48    pub fn new(endpoints: impl IntoIterator<Item = impl Into<String>>) -> Self {
49        Self {
50            endpoints: endpoints.into_iter().map(Into::into).collect(),
51            namespace: "/rust-zero".to_owned(),
52            username: None,
53            password: None,
54            timeout: Duration::from_secs(10),
55            reconnect_backoff: DiscoveryReconnectBackoff::default(),
56        }
57    }
58
59    pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
60        self.namespace = namespace.into();
61        self
62    }
63
64    pub fn with_credentials(
65        mut self,
66        username: impl Into<String>,
67        password: impl Into<String>,
68    ) -> Self {
69        self.username = Some(username.into());
70        self.password = Some(password.into());
71        self
72    }
73
74    pub fn with_timeout(mut self, timeout: Duration) -> Self {
75        assert!(!timeout.is_zero(), "etcd timeout must be positive");
76        self.timeout = timeout;
77        self
78    }
79
80    pub fn with_reconnect_backoff(mut self, backoff: DiscoveryReconnectBackoff) -> Self {
81        self.reconnect_backoff = backoff;
82        self
83    }
84}
85
86#[derive(Debug)]
87pub enum EtcdError {
88    EmptyEndpoints,
89    EmptyName(&'static str),
90    InvalidLeaseTtl,
91    MissingConfig(String),
92    InvalidConfig(crate::ConfigCenterError),
93    Client(etcd_client::Error),
94    Task(String),
95}
96
97impl fmt::Display for EtcdError {
98    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
99        match self {
100            Self::EmptyEndpoints => formatter.write_str("at least one etcd endpoint is required"),
101            Self::EmptyName(kind) => write!(formatter, "etcd {kind} cannot be empty"),
102            Self::InvalidLeaseTtl => formatter.write_str("etcd lease TTL must be positive"),
103            Self::MissingConfig(key) => {
104                write!(formatter, "etcd configuration key {key:?} is missing")
105            }
106            Self::InvalidConfig(error) => error.fmt(formatter),
107            Self::Client(error) => write!(formatter, "etcd operation failed: {error}"),
108            Self::Task(error) => write!(formatter, "etcd background task failed: {error}"),
109        }
110    }
111}
112
113impl Error for EtcdError {
114    fn source(&self) -> Option<&(dyn Error + 'static)> {
115        match self {
116            Self::InvalidConfig(error) => Some(error),
117            Self::Client(error) => Some(error),
118            _ => None,
119        }
120    }
121}
122
123impl From<etcd_client::Error> for EtcdError {
124    fn from(error: etcd_client::Error) -> Self {
125        Self::Client(error)
126    }
127}
128
129impl From<crate::ConfigCenterError> for EtcdError {
130    fn from(error: crate::ConfigCenterError) -> Self {
131        Self::InvalidConfig(error)
132    }
133}
134
135/// A cloneable etcd client scoped to a namespace.
136#[derive(Clone)]
137pub struct EtcdClient {
138    client: Client,
139    namespace: Arc<str>,
140    reconnect_backoff: DiscoveryReconnectBackoff,
141}
142
143impl EtcdClient {
144    pub async fn connect(config: EtcdConfig) -> Result<Self, EtcdError> {
145        if config.endpoints.is_empty() {
146            return Err(EtcdError::EmptyEndpoints);
147        }
148        let namespace = normalize_namespace(&config.namespace)?;
149        let reconnect_backoff = config.reconnect_backoff;
150        let mut options = ConnectOptions::new().with_timeout(config.timeout);
151        if let Some(username) = config.username {
152            options = options.with_user(username, config.password.unwrap_or_default());
153        }
154        let client = Client::connect(&config.endpoints, Some(options)).await?;
155        Ok(Self {
156            client,
157            namespace: Arc::from(namespace),
158            reconnect_backoff,
159        })
160    }
161
162    pub fn from_client(client: Client, namespace: impl AsRef<str>) -> Result<Self, EtcdError> {
163        Ok(Self {
164            client,
165            namespace: Arc::from(normalize_namespace(namespace.as_ref())?),
166            reconnect_backoff: DiscoveryReconnectBackoff::default(),
167        })
168    }
169
170    pub fn with_reconnect_backoff(mut self, backoff: DiscoveryReconnectBackoff) -> Self {
171        self.reconnect_backoff = backoff;
172        self
173    }
174
175    /// Loads a typed configuration value and watches future valid updates.
176    ///
177    /// Invalid updates leave the last known-good snapshot installed. The watcher remains active.
178    pub async fn watch_config<T>(
179        &self,
180        name: impl AsRef<str>,
181        format: ConfigFormat,
182    ) -> Result<EtcdConfigWatcher<T>, EtcdError>
183    where
184        T: DeserializeOwned + Send + Sync + 'static,
185    {
186        let key = self.config_key(name.as_ref())?;
187        let mut client = self.client.clone();
188        let response = client.get(key.clone(), None).await?;
189        let revision = response.header().map_or(0, |header| header.revision());
190        let value = response
191            .kvs()
192            .first()
193            .ok_or_else(|| EtcdError::MissingConfig(key.clone()))?
194            .value_str()?;
195        let config = DynamicConfig::new(value, format)?;
196        let watched = config.clone();
197        let task = tokio::spawn(async move {
198            let options = WatchOptions::new().with_start_revision(revision + 1);
199            let mut stream = client.watch(key, Some(options)).await?;
200            while let Some(response) = stream.message().await? {
201                for event in response.events() {
202                    if event.event_type() == EventType::Put {
203                        if let Some(value) = event.kv() {
204                            // Invalid data is deliberately rejected by DynamicConfig while the
205                            // watch continues to preserve the last known-good configuration.
206                            let _ = watched.update(value.value_str()?);
207                        }
208                    }
209                }
210            }
211            Err(EtcdError::Task("configuration watch closed".to_owned()))
212        });
213        Ok(EtcdConfigWatcher { config, task })
214    }
215
216    /// Publishes an endpoint under a renewable etcd lease.
217    pub async fn publish(
218        &self,
219        service: impl AsRef<str>,
220        instance: impl AsRef<str>,
221        endpoint: impl Into<String>,
222        ttl: Duration,
223    ) -> Result<EtcdServiceLease, EtcdError> {
224        if ttl.is_zero() {
225            return Err(EtcdError::InvalidLeaseTtl);
226        }
227        let key = self.service_key(service.as_ref(), instance.as_ref())?;
228        let endpoint = endpoint.into();
229        if endpoint.trim().is_empty() {
230            return Err(EtcdError::EmptyName("endpoint"));
231        }
232        let ttl_seconds = i64::try_from(ttl.as_secs().max(1)).unwrap_or(i64::MAX);
233        let mut client = self.client.clone();
234        let lease_id = client.lease_grant(ttl_seconds, None).await?.id();
235        client
236            .put(key, endpoint, Some(PutOptions::new().with_lease(lease_id)))
237            .await?;
238        let (shutdown, mut shutdown_receiver) = oneshot::channel();
239        let task = tokio::spawn(async move {
240            let (mut keeper, mut responses) = client.lease_keep_alive(lease_id).await?;
241            let mut interval = tokio::time::interval(Duration::from_secs(
242                u64::try_from((ttl_seconds / 3).max(1)).unwrap_or(1),
243            ));
244            loop {
245                tokio::select! {
246                    _ = &mut shutdown_receiver => {
247                        client.lease_revoke(lease_id).await?;
248                        return Ok(());
249                    }
250                    _ = interval.tick() => {
251                        keeper.keep_alive().await?;
252                        if responses.message().await?.is_none() {
253                            return Err(EtcdError::Task("lease keep-alive stream closed".to_owned()));
254                        }
255                    }
256                }
257            }
258        });
259        Ok(EtcdServiceLease {
260            shutdown: Some(shutdown),
261            task,
262        })
263    }
264
265    /// Subscribes to the complete, sorted endpoint set for a service.
266    pub async fn subscribe(
267        &self,
268        service: impl AsRef<str>,
269    ) -> Result<EtcdServiceSubscription, EtcdError> {
270        let prefix = self.service_prefix(service.as_ref())?;
271        let mut client = self.client.clone();
272        let response = client
273            .get(prefix.clone(), Some(GetOptions::new().with_prefix()))
274            .await?;
275        let revision = response.header().map_or(0, |header| header.revision());
276        let mut entries = BTreeMap::new();
277        for value in response.kvs() {
278            entries.insert(value.key().to_vec(), value.value_str()?.to_owned());
279        }
280        let initial = sorted_endpoints(&entries);
281        let (updates, receiver) = watch::channel(initial);
282        let backoff = self.reconnect_backoff;
283        let task = tokio::spawn(async move {
284            let seed = reconnect_seed(&prefix);
285            let mut revision = revision;
286            let mut attempt = 0_u32;
287            loop {
288                let options = WatchOptions::new()
289                    .with_prefix()
290                    .with_start_revision(revision.saturating_add(1));
291                let watch_result = client.watch(prefix.clone(), Some(options)).await;
292                if let Ok(mut stream) = watch_result {
293                    while let Ok(Some(response)) = stream.message().await {
294                        if let Some(header) = response.header() {
295                            revision = revision.max(header.revision());
296                        }
297                        let mut changed = false;
298                        for event in response.events() {
299                            let Some(value) = event.kv() else {
300                                continue;
301                            };
302                            match event.event_type() {
303                                EventType::Put => {
304                                    entries.insert(
305                                        value.key().to_vec(),
306                                        value.value_str()?.to_owned(),
307                                    );
308                                }
309                                EventType::Delete => {
310                                    entries.remove(value.key());
311                                }
312                            }
313                            changed = true;
314                        }
315                        if changed {
316                            updates.send_replace(sorted_endpoints(&entries));
317                        }
318                        attempt = 0;
319                    }
320                }
321
322                tokio::time::sleep(backoff.delay(attempt, seed.wrapping_add(u64::from(attempt))))
323                    .await;
324                attempt = attempt.saturating_add(1);
325
326                // Always relist after a broken stream. This repairs missed events and recovers
327                // transparently when etcd has compacted the previous watch revision.
328                match client
329                    .get(prefix.clone(), Some(GetOptions::new().with_prefix()))
330                    .await
331                {
332                    Ok(response) => {
333                        revision = response
334                            .header()
335                            .map_or(revision, |header| header.revision());
336                        entries.clear();
337                        for value in response.kvs() {
338                            if let Ok(endpoint) = value.value_str() {
339                                entries.insert(value.key().to_vec(), endpoint.to_owned());
340                            }
341                        }
342                        updates.send_replace(sorted_endpoints(&entries));
343                    }
344                    Err(_) => continue,
345                }
346            }
347            #[allow(unreachable_code)]
348            Ok(())
349        });
350        Ok(EtcdServiceSubscription { receiver, task })
351    }
352
353    fn config_key(&self, name: &str) -> Result<String, EtcdError> {
354        validate_name("configuration name", name)?;
355        Ok(format!("{}/config/{name}", self.namespace))
356    }
357
358    fn service_prefix(&self, service: &str) -> Result<String, EtcdError> {
359        validate_name("service name", service)?;
360        Ok(format!("{}/discovery/{service}/", self.namespace))
361    }
362
363    fn service_key(&self, service: &str, instance: &str) -> Result<String, EtcdError> {
364        validate_name("instance name", instance)?;
365        Ok(format!("{}{instance}", self.service_prefix(service)?))
366    }
367}
368
369pub struct EtcdConfigWatcher<T> {
370    config: DynamicConfig<T>,
371    task: JoinHandle<Result<(), EtcdError>>,
372}
373
374impl<T> EtcdConfigWatcher<T> {
375    pub fn config(&self) -> &DynamicConfig<T> {
376        &self.config
377    }
378
379    pub async fn wait(mut self) -> Result<(), EtcdError> {
380        (&mut self.task)
381            .await
382            .map_err(|error| EtcdError::Task(error.to_string()))?
383    }
384}
385
386impl<T> Drop for EtcdConfigWatcher<T> {
387    fn drop(&mut self) {
388        self.task.abort();
389    }
390}
391
392pub struct EtcdServiceLease {
393    shutdown: Option<oneshot::Sender<()>>,
394    task: JoinHandle<Result<(), EtcdError>>,
395}
396
397impl EtcdServiceLease {
398    pub async fn revoke(mut self) -> Result<(), EtcdError> {
399        if let Some(shutdown) = self.shutdown.take() {
400            let _ = shutdown.send(());
401        }
402        (&mut self.task)
403            .await
404            .map_err(|error| EtcdError::Task(error.to_string()))?
405    }
406}
407
408impl Drop for EtcdServiceLease {
409    fn drop(&mut self) {
410        if let Some(shutdown) = self.shutdown.take() {
411            let _ = shutdown.send(());
412        }
413    }
414}
415
416pub struct EtcdServiceSubscription {
417    receiver: watch::Receiver<Vec<String>>,
418    task: JoinHandle<Result<(), EtcdError>>,
419}
420
421impl EtcdServiceSubscription {
422    pub fn endpoints(&self) -> Vec<String> {
423        self.receiver.borrow().clone()
424    }
425
426    pub async fn changed(&mut self) -> Result<Vec<String>, EtcdError> {
427        self.receiver
428            .changed()
429            .await
430            .map_err(|_| EtcdError::Task("service watch closed".to_owned()))?;
431        Ok(self.endpoints())
432    }
433}
434
435impl EndpointSubscription for EtcdServiceSubscription {
436    type Error = EtcdError;
437
438    fn endpoints(&self) -> Vec<String> {
439        EtcdServiceSubscription::endpoints(self)
440    }
441
442    fn changed(&mut self) -> EndpointChangeFuture<'_, Self::Error> {
443        Box::pin(EtcdServiceSubscription::changed(self))
444    }
445}
446
447impl Drop for EtcdServiceSubscription {
448    fn drop(&mut self) {
449        self.task.abort();
450    }
451}
452
453fn normalize_namespace(namespace: &str) -> Result<String, EtcdError> {
454    let namespace = namespace.trim_matches('/');
455    validate_name("namespace", namespace)?;
456    Ok(format!("/{namespace}"))
457}
458
459fn validate_name(kind: &'static str, value: &str) -> Result<(), EtcdError> {
460    if value.trim().is_empty() || value.contains('/') {
461        Err(EtcdError::EmptyName(kind))
462    } else {
463        Ok(())
464    }
465}
466
467fn sorted_endpoints(entries: &BTreeMap<Vec<u8>, String>) -> Vec<String> {
468    let mut endpoints: Vec<_> = entries.values().cloned().collect();
469    endpoints.sort();
470    endpoints.dedup();
471    endpoints
472}
473
474fn reconnect_seed(scope: &str) -> u64 {
475    let mut hasher = DefaultHasher::new();
476    scope.hash(&mut hasher);
477    SystemTime::now()
478        .duration_since(UNIX_EPOCH)
479        .unwrap_or_default()
480        .as_nanos()
481        .hash(&mut hasher);
482    hasher.finish()
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488    use serde::Deserialize;
489
490    #[derive(Debug, Deserialize)]
491    struct Limits {
492        requests: u64,
493    }
494
495    #[test]
496    fn validates_and_normalizes_namespaces_and_names() {
497        assert_eq!(normalize_namespace("services").unwrap(), "/services");
498        assert_eq!(normalize_namespace("/services/").unwrap(), "/services");
499        assert!(normalize_namespace("/").is_err());
500        assert!(validate_name("service", "bad/name").is_err());
501    }
502
503    #[test]
504    fn endpoint_snapshots_are_sorted_and_deduplicated() {
505        let entries = BTreeMap::from([
506            (b"instance-b".to_vec(), "http://b:8080".to_owned()),
507            (b"instance-a".to_vec(), "http://a:8080".to_owned()),
508            (b"instance-c".to_vec(), "http://a:8080".to_owned()),
509        ]);
510        assert_eq!(
511            sorted_endpoints(&entries),
512            vec!["http://a:8080", "http://b:8080"]
513        );
514    }
515
516    #[tokio::test]
517    async fn integration_covers_config_discovery_and_lease_withdrawal() {
518        let Ok(endpoint) = std::env::var("RUST_ZERO_ETCD_ENDPOINT") else {
519            return;
520        };
521        let namespace = format!("rust-zero-{}", std::process::id());
522        let mut raw = Client::connect([&endpoint], None).await.unwrap();
523        let adapter = EtcdClient::from_client(raw.clone(), &namespace).unwrap();
524        let config_key = format!("/{namespace}/config/limits");
525        raw.put(config_key.clone(), "requests = 10", None)
526            .await
527            .unwrap();
528
529        let watcher = adapter
530            .watch_config::<Limits>("limits", ConfigFormat::Toml)
531            .await
532            .unwrap();
533        let mut changes = watcher.config().subscribe();
534        raw.put(config_key, "requests = 20", None).await.unwrap();
535        changes.changed().await.unwrap();
536        assert_eq!(changes.borrow().value().requests, 20);
537
538        let mut services = adapter.subscribe("users").await.unwrap();
539        let lease = adapter
540            .publish(
541                "users",
542                "instance-a",
543                "http://127.0.0.1:8080",
544                Duration::from_secs(3),
545            )
546            .await
547            .unwrap();
548        assert_eq!(
549            services.changed().await.unwrap(),
550            vec!["http://127.0.0.1:8080"]
551        );
552        lease.revoke().await.unwrap();
553        assert!(services.changed().await.unwrap().is_empty());
554        raw.delete(
555            format!("/{namespace}/"),
556            Some(etcd_client::DeleteOptions::new().with_prefix()),
557        )
558        .await
559        .unwrap();
560    }
561}
rust_zero_core/etcd.rs

rust_zero_core/
etcd.rs
