rust_etcd_utils/
log.rs

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