Skip to main content

rust_etcd_utils/
log.rs

1use etcd_client::{Compare, CompareOp, Txn, TxnOp, WatchOptions};
2use futures::StreamExt;
3use serde::{Serialize, de::DeserializeOwned};
4
5use crate::{
6    lock::ManagedLockGuard,
7    retry::retry_etcd_txn,
8    watcher::{EtcdJsonPutWatchStream, WatchClientExt},
9};
10
11pub struct LogWatcher<T> {
12    stream: EtcdJsonPutWatchStream<T>,
13}
14
15pub struct ExclusiveLogUpdater<'a, T> {
16    etcd: etcd_client::Client,
17    managed_lock_scope: ManagedLockGuard<'a>,
18    last_log_version: i64,
19    log_name: String,
20    _phantom: std::marker::PhantomData<T>,
21}
22
23#[derive(Debug, thiserror::Error)]
24pub enum WriteError {
25    #[error("serialization error: {0}")]
26    SerializationError(String),
27    #[error("Failed to append to log due to lost exclusivity: lock or lease lost")]
28    LostExclusivity,
29    #[error("etcd error: {0}")]
30    EtcdError(#[from] etcd_client::Error),
31}
32
33impl<'a, T> ExclusiveLogUpdater<'a, T>
34where
35    T: Serialize,
36{
37    pub async fn with_scope(
38        log_name: impl AsRef<str>,
39        lock_scope: ManagedLockGuard<'a>,
40    ) -> Result<Self, etcd_client::Error> {
41        let mut etcd = lock_scope.managed_lock.etcd.clone();
42        let mut resp = etcd.get(log_name.as_ref(), None).await?;
43
44        let last_log_version = resp
45            .take_kvs()
46            .into_iter()
47            .max_by_key(|kv| kv.mod_revision())
48            .map(|kv| kv.version())
49            .unwrap_or(0);
50
51        tracing::trace!("last log version: {}", last_log_version);
52
53        Ok(Self {
54            etcd,
55            last_log_version,
56            managed_lock_scope: lock_scope,
57            log_name: log_name.as_ref().to_string(),
58            _phantom: std::marker::PhantomData,
59        })
60    }
61
62    pub async fn write(&mut self, value: T) -> Result<i64, WriteError> {
63        let ser_value = serde_json::to_string(&value)
64            .map_err(|e| WriteError::SerializationError(e.to_string()))?;
65
66        tracing::trace!(
67            "Writing to log `{}`, last version: {}",
68            self.log_name,
69            self.last_log_version
70        );
71        let txn = Txn::new()
72            .when(vec![
73                Compare::version(self.managed_lock_scope.get_key(), CompareOp::Greater, 0),
74                Compare::version(
75                    self.log_name.as_str(),
76                    CompareOp::Equal,
77                    self.last_log_version,
78                ),
79            ])
80            .and_then(vec![TxnOp::put(self.log_name.as_str(), ser_value, None)]);
81
82        let txn_resp = retry_etcd_txn(self.etcd.clone(), txn)
83            .await
84            .map_err(WriteError::EtcdError)?;
85
86        if !txn_resp.succeeded() {
87            tracing::warn!("Failed to append to protected log: {:?}", txn_resp);
88            return Err(WriteError::LostExclusivity);
89        }
90
91        self.last_log_version += 1;
92        Ok(self.last_log_version)
93    }
94}
95
96#[derive(Debug, thiserror::Error)]
97#[error("channel error: {0}")]
98pub struct ReadError(String);
99
100impl<T> LogWatcher<T>
101where
102    T: DeserializeOwned + Send + 'static,
103{
104    ///
105    /// Creates a new log watcher for the given log name.
106    ///
107    pub async fn new(
108        mut etcd: etcd_client::Client,
109        log_name: impl AsRef<str>,
110    ) -> Result<Self, etcd_client::Error> {
111        let mut get_resp = etcd.get(log_name.as_ref(), None).await?;
112
113        // If the key exists, resume from its latest mod revision so callers can
114        // observe current state and any subsequent updates in order.
115        // If the key does not exist yet, start from header revision + 1 to avoid
116        // missing writes between the read and watch establishment.
117        let maybe_watch_opts = get_resp
118            .take_kvs()
119            .into_iter()
120            .map(|kv| kv.mod_revision())
121            .max()
122            .map(|max_mod_rev| WatchOptions::new().with_start_revision(max_mod_rev))
123            .or_else(|| {
124                get_resp
125                    .header()
126                    .map(|h| WatchOptions::new().with_start_revision(h.revision() + 1))
127            });
128
129        let stream = etcd
130            .watch_client()
131            .json_put_watch_stream::<T>(log_name.as_ref(), maybe_watch_opts);
132
133        Ok(Self { stream })
134    }
135
136    ///
137    /// Observes the log for new entries.
138    ///
139    pub async fn observe(&mut self) -> Option<T> {
140        self.stream.next().await.map(|(_revision, val)| val)
141    }
142}