wip

Author: lvboudre

Cargo.lock

@@ -17,6 +17,7 @@ repository = "https://github.com/rpcpool/rust-etcd-utils"
 async-trait = "^0.1.83"
 etcd-client = "^0.17.0"
 futures = "^0.3.31"
+pin-project = "1.1.11"
 rand = "^0.8.5"
 retry = "2"
 serde = { version = "1", features = ["derive"] }
@@ -28,4 +29,4 @@ tonic = "^0.14.2"
 tracing = "^0.1.40"
 
 [dev-dependencies]
-tracing-subscriber = { version = "^0.3.1", features = ["ansi", "env-filter"] }
+tracing-subscriber = { version = "^0.3.1", features = ["ansi", "env-filter"] }

Cargo.toml

@@ -2,10 +2,25 @@ version: '3'
 
 services:
   etcd:
-    image: bitnamilegacy/etcd:3.5
+    # image: bitnamilegacy/etcd:3.5
+    # environment:
+    #   - ALLOW_NONE_AUTHENTICATION=yes
+    #   - ETCD_ADVERTISE_CLIENT_URLS=http://etcd:2379
+    # ports:
+    #   - "2379:2379"
+    #   - "2380:2380"
+    image: quay.io/coreos/etcd:v3.5.0
+    container_name: etcd
     environment:
-      - ALLOW_NONE_AUTHENTICATION=yes
+      - ETCD_NAME=etcd1
+      - ETCD_DATA_DIR=/etcd-data
+      - ETCD_LISTEN_CLIENT_URLS=http://0.0.0.0:2379
       - ETCD_ADVERTISE_CLIENT_URLS=http://etcd:2379
+      - ETCD_LISTEN_PEER_URLS=http://0.0.0.0:2380
+      - ETCD_INITIAL_ADVERTISE_PEER_URLS=http://etcd:2380
+      - ETCD_INITIAL_CLUSTER=etcd1=http://etcd:2380
+      - ETCD_INITIAL_CLUSTER_TOKEN=etcd-single
+      - ETCD_INITIAL_CLUSTER_STATE=new
     ports:
       - "2379:2379"
       - "2380:2380"

compose.yaml

@@ -34,9 +34,9 @@
 /// ```
 use {
     super::{
+        Revision,
         lease::{ManagedLease, ManagedLeaseFactory},
         retry::retry_etcd_legacy,
-        Revision,
     },
     crate::{
         lease::{LeaseExpiredNotify, ManagedLeaseWeak},
@@ -45,7 +45,10 @@ use {
     },
     core::fmt,
     etcd_client::{Compare, CompareOp, GetOptions, LockOptions, Txn, TxnOp, TxnResponse},
-    futures::{future::join_all, FutureExt},
+    futures::{
+        FutureExt, StreamExt,
+        future::{BoxFuture, Shared, join_all},
+    },
     retry::delay::Fixed,
     std::{
         future::Future,
@@ -55,7 +58,7 @@ use {
     },
     thiserror::Error,
     tokio::{
-        sync::{broadcast, mpsc},
+        sync::mpsc,
         task::{JoinError, JoinHandle},
     },
     tonic::Code,
@@ -142,14 +145,14 @@ impl Future for LockManagerHandle {
 /// ```
 ///
 pub struct ManagedLockRevokeNotify {
-    watch_lock_delete: broadcast::Receiver<Revision>,
+    watch_lock_delete: Shared<BoxFuture<'static, ()>>,
     lease_expired_notify: LeaseExpiredNotify,
 }
 
 impl Clone for ManagedLockRevokeNotify {
     fn clone(&self) -> Self {
         Self {
-            watch_lock_delete: self.watch_lock_delete.resubscribe(),
+            watch_lock_delete: self.watch_lock_delete.clone(),
             lease_expired_notify: self.lease_expired_notify.clone(),
         }
     }
@@ -159,14 +162,30 @@ impl ManagedLockRevokeNotify {
     ///
     /// Wait for the lock to be revoked.
     ///
-    pub async fn wait_for_revoke(mut self) {
+    pub async fn wait_for_revoke(self) {
+        let watch_lock_delete = self.watch_lock_delete;
         tokio::select! {
             _ = self.lease_expired_notify.recv() => {}
-            _ = self.watch_lock_delete.recv() => {}
+            _ = watch_lock_delete => {}
         }
     }
 }
 
+fn make_revoke_callback(
+    etcd: etcd_client::Client,
+    lock_key: Vec<u8>,
+    revision: Revision,
+) -> Shared<BoxFuture<'static, ()>> {
+    let mut watch_stream = etcd
+        .watch_client()
+        .watch_lock_key_change_stream(lock_key, revision);
+    async move {
+        let _ = watch_stream.next().await;
+    }
+    .boxed()
+    .shared()
+}
+
 ///
 /// Creates a lock manager to create "managed" locks.
 ///
@@ -402,17 +421,14 @@ impl LockManager {
             }
         };
 
-        let watch_lock_delete = self
-            .etcd
-            .watch_client()
-            .watch_lock_key_change(lock_key.clone(), revision);
+        let revoke_callback = make_revoke_callback(self.etcd.clone(), lock_key.clone(), revision);
         Ok(ManagedLock {
             lock_key,
             managed_lease,
             etcd: self.etcd.clone(),
             created_at_revision: revision,
             delete_signal_tx: self.delete_queue_tx.clone(),
-            revoke_callback_rx: watch_lock_delete.subscribe(),
+            revoke_callback,
         })
     }
 
@@ -487,17 +503,14 @@ impl LockManager {
             }
         };
 
-        let watch_lock_delete = self
-            .etcd
-            .watch_client()
-            .watch_lock_key_change(lock_key.clone(), revision);
+        let revoke_callback = make_revoke_callback(self.etcd.clone(), lock_key.clone(), revision);
         Ok(ManagedLock {
             lock_key,
             managed_lease,
             etcd: self.etcd.clone(),
             created_at_revision: revision,
             delete_signal_tx: self.delete_queue_tx.clone(),
-            revoke_callback_rx: watch_lock_delete.subscribe(),
+            revoke_callback,
         })
     }
 
@@ -562,18 +575,15 @@ impl LockManager {
             lock_response.key().to_vec(),
         );
 
-        let watch_lock_delete = self
-            .etcd
-            .watch_client()
-            .watch_lock_key_change(lock_key.clone(), revision);
+        let revoke_callback = make_revoke_callback(self.etcd.clone(), lock_key.clone(), revision);
 
         let managed_lock = ManagedLock {
             lock_key,
             managed_lease,
             etcd: self.etcd.clone(),
             created_at_revision: revision,
             delete_signal_tx: self.delete_queue_tx.clone(),
-            revoke_callback_rx: watch_lock_delete.subscribe(),
+            revoke_callback,
         };
 
         Ok(managed_lock)
@@ -589,7 +599,7 @@ pub struct ManagedLock {
     pub created_at_revision: Revision,
     pub(crate) etcd: etcd_client::Client,
     delete_signal_tx: tokio::sync::mpsc::UnboundedSender<DeleteQueueCommand>,
-    revoke_callback_rx: broadcast::Receiver<Revision>,
+    revoke_callback: Shared<BoxFuture<'static, ()>>,
 }
 
 impl fmt::Debug for ManagedLock {
@@ -675,7 +685,7 @@ impl ManagedLock {
     ///
     pub fn get_revoke_notify(&self) -> ManagedLockRevokeNotify {
         ManagedLockRevokeNotify {
-            watch_lock_delete: self.revoke_callback_rx.resubscribe(),
+            watch_lock_delete: self.revoke_callback.clone(),
             lease_expired_notify: self.managed_lease.get_lease_expire_notify(),
         }
     }
@@ -744,22 +754,10 @@ impl ManagedLock {
         F: FnOnce(ManagedLockGuard<'a>) -> Fut,
         Fut: Future<Output = T> + Send + 'a,
     {
-        let mut rx = self.revoke_callback_rx.resubscribe();
-
-        match rx.try_recv() {
-            Ok(_) => {
-                tracing::trace!("Lock revoked");
-                return Err(LockError::LockRevoked);
-            }
-            Err(broadcast::error::TryRecvError::Closed) => {
-                tracing::trace!("Lock revoked");
-                return Err(LockError::LockRevoked);
-            }
-            _ => {}
-        }
+        let revoke_callback = self.revoke_callback.clone();
         tokio::select! {
             result = func(ManagedLockGuard { managed_lock: self }) => Ok(result),
-            _ = rx.recv() => Err(LockError::LockRevoked),
+            _ = revoke_callback => Err(LockError::LockRevoked),
         }
     }
 

src/lock.rs

@@ -1,10 +1,15 @@
 use etcd_client::{Compare, CompareOp, Txn, TxnOp, WatchOptions};
-use serde::{de::DeserializeOwned, Serialize};
+use futures::StreamExt;
+use serde::{Serialize, de::DeserializeOwned};
 
-use crate::{lock::ManagedLockGuard, retry::retry_etcd_txn, sync::watch, watcher::WatchClientExt};
+use crate::{
+    lock::ManagedLockGuard,
+    retry::retry_etcd_txn,
+    watcher::{EtcdJsonPutWatchStream, WatchClientExt},
+};
 
 pub struct LogWatcher<T> {
-    rx: watch::Receiver<T>,
+    stream: EtcdJsonPutWatchStream<T>,
 }
 
 pub struct ExclusiveLogUpdater<'a, T> {
@@ -112,26 +117,17 @@ where
             .max()
             .map(|max_mod_rev| WatchOptions::new().with_start_revision(max_mod_rev));
 
-        let mut rx = etcd
+        let stream = etcd
             .watch_client()
-            .json_put_watch_channel::<T>(log_name.as_ref(), maybe_watch_opts);
+            .json_put_watch_stream::<T>(log_name.as_ref(), maybe_watch_opts);
 
-        let (mut wtx, wrx) = watch::watch::<T>();
-
-        let _channel_handle = tokio::spawn(async move {
-            loop {
-                let (_revision, val) = rx.recv().await.expect("watch channel closed");
-                let _ = wtx.update(val).await;
-            }
-        });
-
-        Ok(Self { rx: wrx })
+        Ok(Self { stream })
     }
 
     ///
     /// Observes the log for new entries.
     ///
     pub async fn observe(&mut self) -> Option<T> {
-        self.rx.recv().await
+        self.stream.next().await.map(|(_revision, val)| val)
     }
 }

src/log.rs

@@ -1,4 +1,4 @@
-use std::sync::{atomic::AtomicBool, Arc};
+use std::sync::{Arc, atomic::AtomicBool};
 use tokio::sync::{Mutex, Notify};
 
 struct Inner<T> {