rust_etcd_utils/
lock.rs

1///
2/// This module provides a lock manager to create "managed" locks.
3///
4/// Managed lock's lifecycle is managed by the lock manager background thread, that includes:
5///    - Automatic lease refresh
6///    - Lock revocation when the lock is dropped
7///
8/// You can clone [`LockManager`] to share it across threads, it is really cheap to do so.
9///
10/// See [`spawn_lock_manager`] to create a new lock manager.
11///
12/// # Examples
13///
14/// ```no_run
15///
16/// use etcd_client::Client;
17/// use rust_etcd_utils::{lease::ManagedLeaseFactory, lock::spawn_lock_manager, ManagedLock};
18///
19/// let etcd = Client::connect(["http://localhost:2379"], None).await.expect("failed to connect to etcd");
20///
21/// let managed_lease_factory = ManagedLeaseFactory::new(etcd.clone());
22///
23/// let (lock_man_handle, lock_man) = spawn_lock_manager(etcd.clone(), managed_lease_factory.clone());
24///
25/// // Do something with the lock manager
26///
27/// let managed_lock: ManagedLock = lock_man.try_lock("test").await.expect("failed to lock");
28///
29///
30/// drop(lock_man);
31///
32/// // Wait for the lock manager background thread to stop
33/// lock_man_handle.await.expect("failed to release lock manager handle");
34/// ```
35use {
36    super::{
37        lease::{ManagedLease, ManagedLeaseFactory},
38        retry::retry_etcd_legacy,
39        Revision,
40    },
41    crate::{
42        lease::{LeaseExpiredNotify, ManagedLeaseWeak},
43        retry::{retry_etcd, retry_etcd_txn},
44        watcher::WatchClientExt,
45    },
46    core::fmt,
47    etcd_client::{Compare, CompareOp, GetOptions, LockOptions, Txn, TxnOp, TxnResponse},
48    futures::{future::join_all, FutureExt},
49    retry::delay::Fixed,
50    std::{
51        future::Future,
52        pin::Pin,
53        task::{Context, Poll},
54        time::Duration,
55    },
56    thiserror::Error,
57    tokio::{
58        sync::{broadcast, mpsc},
59        task::{JoinError, JoinHandle},
60    },
61    tonic::Code,
62    tracing::{info, trace},
63};
64
65enum DeleteQueueCommand {
66    Delete(Vec<u8>),
67}
68
69///
70/// A lock manager to create "managed" locks.
71///
72/// Managed lock's lifecycle is managed by the lock manager background thread, that includes:
73///     - Automatic lease refresh
74///     - Lock revocation when the lock is dropped
75///
76/// You can clone [`LockManager`] to share it across threads, it is really cheap to do so.
77///
78/// See [`spawn_lock_manager`] to create a new lock manager.
79///
80#[derive(Clone)]
81pub struct LockManager {
82    etcd: etcd_client::Client,
83    delete_queue_tx: mpsc::UnboundedSender<DeleteQueueCommand>,
84    manager_lease_factory: ManagedLeaseFactory,
85    try_locking_timeout: Duration,
86    #[allow(dead_code)]
87    // When all Sender to this channel is dropped, the Lock manager background thread will stop.
88    lock_manager_handle_entangled_tx: mpsc::UnboundedSender<()>,
89}
90
91///
92/// Handle to the lock manager background thread.
93///
94/// See [`spawn_lock_manager`] to create a new lock manager.
95pub struct LockManagerHandle {
96    inner: JoinHandle<()>,
97}
98
99impl Future for LockManagerHandle {
100    type Output = Result<(), JoinError>;
101    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
102        self.inner.poll_unpin(cx)
103    }
104}
105
106///
107/// Used to notify when a lock is revoked.
108///
109/// Examples
110///
111/// ```no_run
112///
113/// use etcd_client::Client;
114/// use rust_etcd_utils::{lease::ManagedLeaseFactory, lock::spawn_lock_manager, ManagedLock};
115///
116/// let etcd = Client::connect(["http://localhost:2379"], None).await.expect("failed to connect to etcd");
117///
118/// let managed_lease_factory = ManagedLeaseFactory::new(etcd.clone());
119///
120/// let (lock_man_handle, lock_man) = spawn_lock_manager(etcd.clone(), managed_lease_factory.clone());
121///
122/// // Do something with the lock manager
123///
124/// let managed_lock: ManagedLock = lock_man.try_lock("test").await.expect("failed to lock");
125///
126/// let revoke_notify = managed_lock.get_revoke_notify();
127///
128/// // Can be cloned
129/// let revoke_notify2 = revoke_notify.clone();
130///
131/// // Can create multiple instances..
132///
133/// let revoke_notify3 = managed_lock.get_revoke_notify();
134///
135/// etcd.delete("test", None).await.expect("failed to delete");
136///
137/// revoke_notify.wait_for_revoke().await;
138/// revoke_notify2.wait_for_revoke().await;
139/// revoke_notify3.wait_for_revoke().await;
140///
141/// println!("All revoke notify received");
142/// ```
143///
144pub struct ManagedLockRevokeNotify {
145    watch_lock_delete: broadcast::Receiver<Revision>,
146    lease_expired_notify: LeaseExpiredNotify,
147}
148
149impl Clone for ManagedLockRevokeNotify {
150    fn clone(&self) -> Self {
151        Self {
152            watch_lock_delete: self.watch_lock_delete.resubscribe(),
153            lease_expired_notify: self.lease_expired_notify.clone(),
154        }
155    }
156}
157
158impl ManagedLockRevokeNotify {
159    ///
160    /// Wait for the lock to be revoked.
161    ///
162    pub async fn wait_for_revoke(mut self) {
163        tokio::select! {
164            _ = self.lease_expired_notify.recv() => {}
165            _ = self.watch_lock_delete.recv() => {}
166        }
167    }
168}
169
170///
171/// Creates a lock manager to create "managed" locks.
172///
173/// Managed lock's lifecycle is managed by the lock manager background thread, that includes:
174///     - Automatic lease refresh
175///     - Lock revocation when the lock is dropped
176///
177/// You can clone [[`LockManager`]] to share it across threads, it is really cheap to do so.
178///
179/// The lock manager background thread will stop when all the [[`LockManager`]] is dropped.
180/// You can await on the [[`LockManagerHandle`]] to wait for the lock manager background thread to stop.
181///
182/// Dropping the [[`LockManagerHandle`]] will not stop the lock manager background thread, but it is not recommended to do so
183/// as you are suppoed to `await` the handle to gracefully shutdown.
184///
185/// Examples
186///
187/// ```no_run
188/// use etcd_client::Client;
189/// use rust_etcd_utils::{lock::spawn_lock_manager, ManagedLock};
190///
191/// let etcd = Client::connect(["http://localhost:2379"], None).await.expect("failed to connect to etcd");
192///
193/// let (lock_man_handle, lock_man) = spawn_lock_manager(etcd.clone());
194///
195/// // Do something with the lock manager
196///
197/// let managed_lock: ManagedLock = lock_man.try_lock("test").await.expect("failed to lock");
198///
199///
200/// drop(lock_man);
201///
202/// // Wait for the lock manager background thread to stop
203/// lock_man_handle.await.expect("failed to release lock manager handle");
204/// ```
205///
206/// Cloning the lock manager is cheap and can be shared across threads.
207///
208/// ```no_run
209///
210/// use etcd_client::Client;
211/// use rust_etcd_utils::{lock::spawn_lock_manager, lock::ManagedLock, lock::TryLockError};
212///
213/// let etcd = Client::connect(["http://localhost:2379"], None).await.expect("failed to connect to etcd");
214/// let (_, lock_man) = spawn_lock_manager(etcd.clone());
215///     
216/// let lock_man2 = lock_man.clone();
217/// let task1 = tokio::spawn(async move {
218///     let managed_lock: ManagedLock = lock_man2.try_lock("test").await.expect("failed to lock");
219///
220///     tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
221/// });
222///
223///
224/// let err = lock_man.try_lock("test").await;
225///
226/// assert!(matches!(err, Err(TryLockError::AlreadyTaken)));
227/// ```
228///
229pub fn spawn_lock_manager(etcd: etcd_client::Client) -> (LockManagerHandle, LockManager) {
230    let (lease_factory, _) = ManagedLeaseFactory::spawn(etcd.clone());
231    spawn_lock_manager_with_lease_factory(etcd, lease_factory)
232}
233
234///
235/// Creates a lock manager to create "managed" locks.
236///
237/// See [`spawn_lock_manager`] for more details.
238///
239pub fn spawn_lock_manager_with_lease_factory(
240    etcd: etcd_client::Client,
241    managed_lease_factory: ManagedLeaseFactory,
242) -> (LockManagerHandle, LockManager) {
243    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
244    let etcd2 = etcd.clone();
245
246    let (entangled_tx, mut entangled_rx) = tokio::sync::mpsc::unbounded_channel();
247    let tx2 = tx.clone();
248    let handle = tokio::spawn(async move {
249        let _tx2 = tx2;
250        loop {
251            let cmd = tokio::select! {
252                cmd = rx.recv() => {
253                    cmd.expect("command rx droped")
254                }
255                maybe = entangled_rx.recv() => {
256                    match maybe {
257                        Some(_) => unreachable!("entangled_rx should not have any message"),
258                        None => {
259                            // If entangled_rx is closed, we should stop the loop because it means there is no LockManager instance alive.
260                            break
261                        },
262                    }
263                }
264            };
265            match cmd {
266                DeleteQueueCommand::Delete(lock_id) => {
267                    let kv_client = etcd2.kv_client();
268                    let lock_id2 = lock_id.clone();
269                    let result = retry_etcd_legacy(Fixed::from_millis(10), move || {
270                        let lock_id = lock_id2.clone();
271                        let mut kv_client = kv_client.clone();
272                        async move { kv_client.delete(lock_id, None).await }
273                    })
274                    .await;
275                    match result {
276                        Ok(_) => {
277                            let lock_id = String::from_utf8(lock_id).expect("lock id is not utf8");
278                            info!("Deleted lock {lock_id}");
279                        }
280                        Err(e) => {
281                            if !matches!(e, etcd_client::Error::GRpcStatus(ref status) if status.code() == Code::NotFound)
282                            {
283                                tracing::error!("Failed to revoke lock: {e}");
284                                // panic!("Failed to delete lock: {e}");
285                            }
286                        }
287                    }
288                }
289            }
290        }
291        let mut futures = vec![];
292        // Drain any remaining delete commands
293        while let Ok(cmd) = rx.try_recv() {
294            match cmd {
295                DeleteQueueCommand::Delete(lock_id) => {
296                    let mut kv_client = etcd2.kv_client();
297                    let fut = async move { kv_client.delete(lock_id, None).await };
298                    futures.push(fut);
299                }
300            }
301        }
302        // Since we are closing the channel, we can ignore the result of the futures
303        let _ = join_all(futures).await;
304    });
305    let handle = LockManagerHandle { inner: handle };
306    (
307        handle,
308        LockManager {
309            etcd,
310            delete_queue_tx: tx,
311            try_locking_timeout: Duration::from_secs(1),
312            manager_lease_factory: managed_lease_factory,
313            lock_manager_handle_entangled_tx: entangled_tx,
314        },
315    )
316}
317
318///
319/// Error that can occur when trying to lock a key.
320///
321#[derive(Debug, thiserror::Error)]
322pub enum LockingError {
323    #[error("Etcd error: {0:?}")]
324    EtcdError(etcd_client::Error),
325}
326
327impl LockManager {
328    ///
329    /// Tries to lock a key with automatic lease refresh and lock revocation when dropped.
330    ///
331    /// If the key is already held by another lock, it will return an error immediately.
332    ///
333    pub async fn try_lock<S>(
334        &self,
335        name: S,
336        lease_duration: Duration,
337    ) -> Result<ManagedLock, TryLockError>
338    where
339        S: AsRef<str>,
340    {
341        let name = name.as_ref();
342        if self.delete_queue_tx.is_closed() {
343            panic!("LockManager lifecycle thread is stopped.");
344        }
345        let gopts = GetOptions::new().with_prefix();
346        trace!("Trying to lock {name}...");
347        let get_response = retry_etcd(
348            self.etcd.clone(),
349            (name.to_string(), gopts),
350            move |etcd, (name, gopts)| async move { etcd.kv_client().get(name, Some(gopts)).await },
351        )
352        .await
353        .map_err(TryLockError::EtcdError)?;
354
355        if get_response.count() > 0 {
356            return Err(TryLockError::AlreadyTaken);
357        }
358
359        let managed_lease = self
360            .manager_lease_factory
361            .new_lease(lease_duration, None)
362            .await
363            .map_err(TryLockError::EtcdError)?;
364        let lease_id = managed_lease.lease_id;
365
366        let lock_fut = retry_etcd(
367            self.etcd.clone(),
368            (name.to_string(), LockOptions::new().with_lease(lease_id)),
369            |mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
370        );
371
372        let lease_expire_notify = managed_lease.get_lease_expire_notify();
373
374        let (revision, lock_key) = tokio::select! {
375            _ = tokio::time::sleep(self.try_locking_timeout) => {
376                return Err(TryLockError::LockingDeadlineExceeded)
377            }
378            result = lock_fut => {
379                let lock_response = match result {
380                    Ok(lock_response) => lock_response,
381                    Err(e) => {
382                        match e {
383                            etcd_client::Error::GRpcStatus(status) => {
384                                if status.code() == Code::Unknown {
385                                    if status.message() == "etcdserver: requested lease not found" {
386                                        return Err(TryLockError::LeaseExpired)
387                                    } else {
388                                        return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
389                                    }
390                                } else {
391                                    return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
392                                }
393                            }
394                            _ => return Err(TryLockError::EtcdError(e))
395                        }
396                    }
397                };
398                (lock_response.header().expect("empty header for etcd lock").revision(), lock_response.key().to_vec())
399            }
400            _ = lease_expire_notify.recv() => {
401                return Err(TryLockError::LeaseExpired)
402            }
403        };
404
405        let watch_lock_delete = self
406            .etcd
407            .watch_client()
408            .watch_lock_key_change(lock_key.clone(), revision);
409        Ok(ManagedLock {
410            lock_key,
411            managed_lease,
412            etcd: self.etcd.clone(),
413            created_at_revision: revision,
414            delete_signal_tx: self.delete_queue_tx.clone(),
415            revoke_callback_rx: watch_lock_delete.subscribe(),
416        })
417    }
418
419    ///
420    /// Similar to [`LockManager::try_lock`] but with a custom lease.
421    ///
422    pub async fn try_lock_with_lease<S>(
423        &self,
424        name: S,
425        managed_lease: ManagedLease,
426    ) -> Result<ManagedLock, TryLockError>
427    where
428        S: AsRef<str>,
429    {
430        let name = name.as_ref();
431        if self.delete_queue_tx.is_closed() {
432            panic!("LockManager lifecycle thread is stopped.");
433        }
434        let gopts = GetOptions::new().with_prefix();
435        const TRY_LOCKING_DURATION: Duration = Duration::from_millis(1000);
436        trace!("Trying to lock {name}...");
437        let get_response = retry_etcd(
438            self.etcd.clone(),
439            (name.to_string(), gopts),
440            move |etcd, (name, gopts)| async move { etcd.kv_client().get(name, Some(gopts)).await },
441        )
442        .await
443        .map_err(TryLockError::EtcdError)?;
444
445        if get_response.count() > 0 {
446            return Err(TryLockError::AlreadyTaken);
447        }
448
449        let lease_id = managed_lease.lease_id;
450
451        let lock_fut = retry_etcd(
452            self.etcd.clone(),
453            (name.to_string(), LockOptions::new().with_lease(lease_id)),
454            |mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
455        );
456
457        let lease_expire_notify = managed_lease.get_lease_expire_notify();
458        let (revision, lock_key) = tokio::select! {
459            _ = tokio::time::sleep(TRY_LOCKING_DURATION) => {
460                return Err(TryLockError::LockingDeadlineExceeded)
461            }
462            result = lock_fut => {
463                let lock_response = match result {
464                    Ok(lock_response) => lock_response,
465                    Err(e) => {
466                        match e {
467                            etcd_client::Error::GRpcStatus(status) => {
468                                if status.code() == Code::Unknown {
469                                    if status.message() == "etcdserver: requested lease not found" {
470                                        return Err(TryLockError::LeaseExpired)
471                                    } else {
472                                        return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
473                                    }
474                                } else {
475                                    return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
476                                }
477                            }
478                            _ => return Err(TryLockError::EtcdError(e))
479                        }
480                    }
481                };
482
483                (lock_response.header().expect("empty header for etcd lock").revision(), lock_response.key().to_vec())
484            }
485            _ = lease_expire_notify.recv() => {
486                return Err(TryLockError::LeaseExpired)
487            }
488        };
489
490        let watch_lock_delete = self
491            .etcd
492            .watch_client()
493            .watch_lock_key_change(lock_key.clone(), revision);
494        Ok(ManagedLock {
495            lock_key,
496            managed_lease,
497            etcd: self.etcd.clone(),
498            created_at_revision: revision,
499            delete_signal_tx: self.delete_queue_tx.clone(),
500            revoke_callback_rx: watch_lock_delete.subscribe(),
501        })
502    }
503
504    ///
505    /// Locks a key with automatic lease refresh and lock revocation when dropped.
506    /// If the key is already held by another lock, it will wait until the lock is released.
507    ///
508    /// Be aware this method can await indefinitely if the key is never released.
509    pub async fn lock<S>(
510        &self,
511        name: S,
512        lease_duration: Duration,
513    ) -> Result<ManagedLock, etcd_client::Error>
514    where
515        S: AsRef<str>,
516    {
517        if self.delete_queue_tx.is_closed() {
518            panic!("LockManager lifecycle thread is stopped.");
519        }
520        let managed_lease = self
521            .manager_lease_factory
522            .new_lease(lease_duration, None)
523            .await?;
524        self.lock_with_lease(name, managed_lease).await
525    }
526
527    ///
528    /// Similar to [`LockManager::lock`] but with a custom lease.
529    ///
530    pub async fn lock_with_lease<S>(
531        &self,
532        name: S,
533        managed_lease: ManagedLease,
534    ) -> Result<ManagedLock, etcd_client::Error>
535    where
536        S: AsRef<str>,
537    {
538        if self.delete_queue_tx.is_closed() {
539            panic!("LockManager lifecycle thread is stopped.");
540        }
541        let name = name.as_ref();
542
543        let lease_id = managed_lease.lease_id;
544
545        let lock_fut = retry_etcd(
546            self.etcd.clone(),
547            (name.to_string(), LockOptions::new().with_lease(lease_id)),
548            |mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
549        );
550
551        let lock_response = tokio::select! {
552            result = lock_fut => {
553                result?
554            }
555        };
556
557        let (revision, lock_key) = (
558            lock_response
559                .header()
560                .expect("empty header for etcd lock")
561                .revision(),
562            lock_response.key().to_vec(),
563        );
564
565        let watch_lock_delete = self
566            .etcd
567            .watch_client()
568            .watch_lock_key_change(lock_key.clone(), revision);
569
570        let managed_lock = ManagedLock {
571            lock_key,
572            managed_lease,
573            etcd: self.etcd.clone(),
574            created_at_revision: revision,
575            delete_signal_tx: self.delete_queue_tx.clone(),
576            revoke_callback_rx: watch_lock_delete.subscribe(),
577        };
578
579        Ok(managed_lock)
580    }
581}
582
583///
584/// A Lock instance with automatic lease refresh and lock revocation when dropped.
585///
586pub struct ManagedLock {
587    pub(crate) lock_key: Vec<u8>,
588    managed_lease: ManagedLease,
589    pub created_at_revision: Revision,
590    pub(crate) etcd: etcd_client::Client,
591    delete_signal_tx: tokio::sync::mpsc::UnboundedSender<DeleteQueueCommand>,
592    revoke_callback_rx: broadcast::Receiver<Revision>,
593}
594
595impl fmt::Debug for ManagedLock {
596    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
597        f.debug_struct("ManagedLock")
598            .field("lock_key", &String::from_utf8_lossy(&self.lock_key))
599            .field("lease_id", &self.managed_lease.lease_id)
600            .field("created_At_revision", &self.created_at_revision)
601            .finish()
602    }
603}
604
605impl Drop for ManagedLock {
606    fn drop(&mut self) {
607        info!(
608            "Destructor called for ManagedLock({})",
609            String::from_utf8_lossy(&self.lock_key)
610        );
611        let _ = self
612            .delete_signal_tx
613            .send(DeleteQueueCommand::Delete(self.lock_key.clone()));
614    }
615}
616
617///
618/// Error that can occur when using a managed lock.
619///
620#[derive(Debug)]
621pub enum LockError {
622    LockRevoked,
623}
624
625impl fmt::Display for LockError {
626    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
627        match self {
628            LockError::LockRevoked => f.write_str("lock revoked"),
629        }
630    }
631}
632
633///
634/// Acquired only via [`ManagedLock::scope`] or [`ManagedLock::scope_with`].
635///
636/// This guard represents the scope of the managed lock, and can be used to express restriction on code execution so
637/// that it is only executed within the lock lifetime.
638///
639/// This guard can also be used with other modules in the library such as the `log` module.
640///
641pub struct ManagedLockGuard<'a> {
642    pub(crate) managed_lock: &'a ManagedLock,
643}
644
645impl ManagedLockGuard<'_> {
646    pub(crate) fn get_key(&self) -> &[u8] {
647        self.managed_lock.lock_key.as_slice()
648    }
649}
650
651impl ManagedLock {
652    pub fn lease_id(&self) -> i64 {
653        self.managed_lease.lease_id
654    }
655
656    ///
657    /// Execute an etcd transaction if and only if the lock is still alive.
658    ///
659    pub async fn txn(&self, operations: impl Into<Vec<TxnOp>>) -> TxnResponse {
660        let txn = Txn::new()
661            .when(vec![Compare::version(
662                self.lock_key.clone(),
663                CompareOp::Greater,
664                0,
665            )])
666            .and_then(operations);
667
668        retry_etcd_txn(self.etcd.clone(), txn)
669            .await
670            .expect("failed txn")
671    }
672
673    ///
674    /// Get a revoke notify handle to be notified when the lock is revoked.
675    ///
676    pub fn get_revoke_notify(&self) -> ManagedLockRevokeNotify {
677        ManagedLockRevokeNotify {
678            watch_lock_delete: self.revoke_callback_rx.resubscribe(),
679            lease_expired_notify: self.managed_lease.get_lease_expire_notify(),
680        }
681    }
682
683    ///
684    /// Check if the lock is still alive.
685    ///
686    pub async fn is_alive(&self) -> bool {
687        let get_response = self
688            .etcd
689            .kv_client()
690            .get(self.lock_key.as_slice(), None)
691            .await
692            .expect("failed to communicate with etcd");
693        get_response.count() == 1
694    }
695
696    ///
697    /// Get the underlying unique lock key.
698    ///
699    pub fn get_key(&self) -> Vec<u8> {
700        self.lock_key.clone()
701    }
702
703    ///
704    /// This function make sure the future is executed within a valid managed lock lifetime.
705    ///
706    /// If the lock is revoked, it will cancel the future and return a LockError::LockRevoked.
707    ///
708    /// Make sure the future returned by the closure is cancel safe.
709    ///
710    /// Examples
711    ///
712    /// ```no_run
713    /// use etcd_client::Client;
714    /// use rust_etcd_utils::{lease::ManagedLeaseFactory, lock::spawn_lock_manager, ManagedLock};
715    ///
716    /// let etcd = Client::connect(["http://localhost:2379"], None).await.expect("failed to connect to etcd");
717    ///
718    /// let (lock_man_handle, lock_man) = spawn_lock_manager(etcd.clone());
719    ///
720    /// // Do something with the lock manager
721    ///
722    /// let managed_lock: ManagedLock = lock_man.try_lock("test").await.expect("failed to lock");
723    ///
724    /// managed_lock.scope(async move {
725    ///    // execute only if my lock is valid
726    ///    access_protected_ressource().await;
727    /// });
728    ///
729    /// ```
730    pub async fn scope<T, Fut>(&self, fut: Fut) -> Result<T, LockError>
731    where
732        T: Send + 'static,
733        Fut: Future<Output = T> + Send + 'static,
734    {
735        self.scope_with(move |_| fut).await
736    }
737
738    ///
739    /// Similar to [`ManagedLock::scope`] but accept a closure to compute the future to execute against the lock.
740    ///
741    pub async fn scope_with<'a, T, F, Fut>(&'a self, func: F) -> Result<T, LockError>
742    where
743        T: Send + 'a,
744        F: FnOnce(ManagedLockGuard<'a>) -> Fut,
745        Fut: Future<Output = T> + Send + 'a,
746    {
747        let mut rx = self.revoke_callback_rx.resubscribe();
748
749        match rx.try_recv() {
750            Ok(_) => {
751                tracing::trace!("Lock revoked");
752                return Err(LockError::LockRevoked);
753            }
754            Err(broadcast::error::TryRecvError::Closed) => {
755                tracing::trace!("Lock revoked");
756                return Err(LockError::LockRevoked);
757            }
758            _ => {}
759        }
760        tokio::select! {
761            result = func(ManagedLockGuard { managed_lock: self }) => Ok(result),
762            _ = rx.recv() => Err(LockError::LockRevoked),
763        }
764    }
765
766    ///
767    /// Get a weak reference to the managed lease.
768    ///
769    pub fn get_managed_lease_weak_ref(&self) -> ManagedLeaseWeak {
770        self.managed_lease.get_weak()
771    }
772
773    ///
774    /// Convert the managed lock into a signal that will be resolved when the lock is revoked.
775    ///
776    pub async fn into_revoked_fut(self) {
777        let _ = self.scope(futures::future::pending::<()>()).await;
778    }
779}
780
781///
782/// Error that can occur when trying to lock a key.
783///
784#[derive(Debug, Error)]
785pub enum TryLockError {
786    #[error("Already taken")]
787    AlreadyTaken,
788    #[error("Locking deadline exceeded")]
789    LockingDeadlineExceeded,
790    #[error("Lease expired before the lock")]
791    LeaseExpired,
792    #[error("Etcd error: {0:?}")]
793    EtcdError(etcd_client::Error),
794}
rust_etcd_utils/lock.rs

rust_etcd_utils/
lock.rs
