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

rust_etcd_utils/
lock.rs
