tokio_etcd_lock/

lib.rs

mod error;
mod lease;
mod lock;

use std::sync::Arc;

pub use crate::{error::*, lease::*, lock::*};
use etcd_client::{
    Client, Compare, CompareOp, Event, PutOptions, ResponseHeader, Txn, TxnOp, WatchOptions,
    WatchStream,
};
use futures::prelude::*;
use log::*;
use rand::random;

pub enum TtlOrLeaseId {
    Ttl(i64),
    LeaseId(i64),
}

pub async fn try_race(
    client: &mut Client,
    key: &str,
    ttl_or_lease_id: TtlOrLeaseId,
) -> CrateResult<bool> {
    let value = random::<u128>().to_le_bytes();

    let lease_id = match ttl_or_lease_id {
        TtlOrLeaseId::Ttl(ttl) => client.lease_grant(ttl, None).await?.id(),
        TtlOrLeaseId::LeaseId(id) => id,
    };

    let locked = client
        .txn(
            Txn::new()
                .when(vec![Compare::create_revision(
                    key.as_bytes(),
                    CompareOp::Equal,
                    0,
                )])
                .and_then(vec![TxnOp::put(
                    key.as_bytes(),
                    &value[..],
                    Some(PutOptions::new().with_lease(lease_id)),
                )]),
        )
        .await?
        .succeeded();

    Ok(locked)
}

fn get_lock_txn(name: &str, lock_value: &[u8], lease_id: i64) -> Txn {
    Txn::new()
        .when([Compare::create_revision(name, CompareOp::Equal, 0)])
        .and_then([TxnOp::put(
            name,
            lock_value,
            Some(PutOptions::new().with_lease(lease_id)),
        )])
}

pub(crate) fn get_release_txn(name: &str, lock_value: &[u8], lease_id: i64) -> Txn {
    Txn::new()
        .when([
            Compare::value(name, CompareOp::Equal, lock_value),
            Compare::lease(name, CompareOp::Equal, lease_id),
        ])
        .and_then([TxnOp::delete(name, None)])
}

async fn watch_until_lock_available(
    etcd_client: &mut Client,
    watch_stream: &mut Option<WatchStream>,
    name: &str,
    last_locked_revision: Option<i64>,
) -> CrateResult<()> {
    loop {
        let mut etcd_client = etcd_client.clone();
        let mut inner_watch_stream = match watch_stream.take() {
            Some(watch_stream) => watch_stream,
            None => {
                let mut watch_options = WatchOptions::new();
                if let Some(last_locked_revision) = last_locked_revision {
                    watch_options = watch_options.with_start_revision(last_locked_revision + 1);
                }

                let (_, watch_stream) = etcd_client.watch(name, Some(watch_options)).await?;

                watch_stream
            }
        };

        let watch_response = inner_watch_stream
            .try_next()
            .await?
            .expect("Unexpected end of etcd watch stream");

        *watch_stream = Some(inner_watch_stream);

        let is_available = matches!(watch_response.events().last().and_then(Event::kv), Some(kv) if kv.create_revision() == 0);

        trace!(
            "Got update from watch stream for key: {} (avilable: {})",
            name,
            is_available
        );

        if is_available {
            return Ok(());
        }
    }
}

impl Lease {
    pub(crate) async fn stop_keep_alive(self) -> Result<(), etcd_client::Error> {
        if let Ok(LeaseInner {
            stop_tx,
            keep_alive_handle,
        }) = Arc::try_unwrap(self.inner)
        {
            let _ = stop_tx.send(());
            keep_alive_handle.await.unwrap()?;
        }

        Ok(())
    }

    /// Stop the keep-alive task, propagating any errors that it may have
    /// encountered, and ask the etcd cluster to revoke this lease. Any etcd
    /// keys associated with this lease, including keys used as locks by
    /// this library's [`LockGuard`], will be removed.
    ///
    /// See also: [etcd docs]
    ///
    /// [`LockGuard`]: ./struct.LockGuard.html
    /// [etcd docs]: https://etcd.io/docs/v3.3.12/dev-guide/interacting_v3/#revoke-leases
    pub async fn release(self, client: &mut Client) -> Result<(), etcd_client::Error> {
        let lease_id = self.lease_id;

        self.stop_keep_alive().await?;
        client.lease_revoke(lease_id).await?;

        debug!("Released lease: {:#x}", lease_id);

        Ok(())
    }

    async fn _lock<'l, 'n>(
        &'l mut self,
        client: &mut Client,
        name: &'n str,
        wait: bool,
    ) -> CrateResult<LockGuard<'l, 'n>> {
        let lock_value = random::<u128>().to_le_bytes();

        let mut watch_stream = None;
        let locked = loop {
            let txn_response = client
                .txn(get_lock_txn(name, &lock_value, self.lease_id))
                .await?;

            if txn_response.succeeded() {
                break true;
            } else if !wait {
                break false;
            }

            watch_until_lock_available(
                client,
                &mut watch_stream,
                name,
                txn_response.header().map(ResponseHeader::revision),
            )
            .await?;
        };

        if !locked {
            return Err(Error::Taken);
        } else {
            debug!("Acquired lock: {} (lease: {:#x})", name, self.lease_id);
        }

        Ok(LockGuard {
            lock_value,
            name,
            lease: &*self,
        })
    }

    pub async fn lock<'l, 'n>(
        &'l mut self,
        client: &mut Client,
        name: &'n str,
    ) -> CrateResult<LockGuard<'l, 'n>> {
        self._lock(client, name, true).await
    }

    /// Anologous to [`lock`], but if the lock is taken, return [`Error::Taken`]
    /// instead of waiting for the lock to become available.
    ///
    /// [`lock`]: ./struct.Lease.html#method.lock
    /// [`Error::taken`]: ./enum.Error.html#variant.Taken
    pub async fn try_lock<'l, 'n>(
        &'l mut self,
        client: &mut Client,
        name: &'n str,
    ) -> CrateResult<LockGuard<'l, 'n>> {
        self._lock(client, name, false).await
    }

    /// Run and await future-returning function `f`, but only in the context of
    /// an acquired lock.
    ///
    /// Attempt to acquire lock with `name`, waiting for it to become available
    /// if needed, and run `f` after it is successfuly acquired. After `f`
    /// is finished and the future it returned is awaited, the lock is released
    /// automatically.
    ///
    /// # Panic handling
    /// This function does not attempt to catch panics from `f`, and does not
    /// guarantee that the lock is released in such cases. If the program
    /// exits abruptly, the lock will be not be available to other clients until
    /// this lease's TTL runs out.
    ///
    /// # Example
    /// ```no_run
    /// # async fn test(mut lease: tokio_etcd_lock::Lease, mut client: etcd_client::Client) -> Result<(), tokio_etcd_lock::Error> {
    ///
    /// lease.with_lock(&mut client, "lock-key", || async {
    ///     println!("Holding mutually exclusive access to the lock!");
    /// }).await?;
    ///
    /// // at this point, the lock has been released and is available for other clients
    /// # Ok(())
    /// # }
    /// ```
    pub async fn with_lock<F, Fut, T>(
        &mut self,
        client: &mut Client,
        name: &str,
        f: F,
    ) -> CrateResult<T>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = T>,
    {
        let guard = self.lock(client, name).await?;
        let result = f().await;
        guard.release(client).await?;

        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::{random, random_range};
    use std::{
        ops::AddAssign,
        sync::{
            atomic::{AtomicUsize, Ordering},
            Arc,
        },
        time::Duration,
    };
    use tokio::{
        sync::{Barrier, Mutex},
        time::sleep,
    };

    async fn create_client() -> CrateResult<Client> {
        Ok(Client::connect(vec!["http://localhost:2379"], None).await?)
    }

    fn create_lock_name() -> String {
        format!("{:x}", random::<u128>())
    }

    #[tokio::test]
    async fn test_mutex() -> anyhow::Result<()> {
        const WORKERS: usize = 10;

        type Result<T> = anyhow::Result<T, anyhow::Error>;

        let lock_name = create_lock_name();

        let client = create_client().await?;
        let locked_count = Arc::new(Mutex::new(0_usize));

        stream::iter(0..WORKERS)
            .map(Result::Ok)
            .try_for_each_concurrent(None, |_| {
                let mut client = client.clone();
                let lock_name = lock_name.clone();
                let locked_count = locked_count.clone();

                async move {
                    let mut lease = acquire_lease(&mut client, 10).await?;

                    lease
                        .with_lock(&mut client, &lock_name, || async move {
                            {
                                let mut locked_count = locked_count.try_lock()?;
                                sleep(Duration::from_millis(random_range(0..2_000))).await;
                                locked_count.add_assign(1);
                            }

                            Result::Ok(())
                        })
                        .await??;

                    Result::Ok(())
                }
            })
            .await?;

        assert_eq!(*locked_count.try_lock().unwrap(), WORKERS);

        Ok(())
    }

    #[tokio::test]
    async fn test_lease_expire() -> CrateResult<()> {
        let lock_name = create_lock_name();

        let mut client = create_client().await?;

        let lease1 = acquire_lease(&mut client, 10).await?;
        let mut lease2 = acquire_lease(&mut client, 10).await?;

        let barrier1 = Arc::new(Barrier::new(2));
        let barrier2 = Arc::new(Barrier::new(2));

        future::try_join(
            {
                let mut client = client.clone();
                let lock_name = lock_name.clone();
                let barrier1 = barrier1.clone();
                let barrier2 = barrier2.clone();
                async move {
                    client
                        .txn(get_lock_txn(&lock_name, b"", lease1.lease_id))
                        .await?;
                    barrier1.wait().await;
                    barrier2.wait().await;
                    lease1.stop_keep_alive().await?;
                    Ok(())
                }
            },
            {
                let mut client = client.clone();

                async move {
                    barrier1.wait().await;

                    // make sure the lock is still taken at this point
                    assert!(matches!(
                        lease2.try_lock(&mut client, &lock_name).await.unwrap_err(),
                        Error::Taken
                    ));

                    // allow first worker to stop the lease keep-alive
                    barrier2.wait().await;

                    lease2
                        .with_lock(&mut client, &lock_name, || async {
                            info!("Inside lock.");
                        })
                        .await
                }
            },
        )
        .await?;

        Ok(())
    }

    #[tokio::test]
    async fn test_lock_lost() -> CrateResult<()> {
        let lock_name = create_lock_name();

        let mut client = create_client().await?;
        let mut lease = acquire_lease(&mut client, 10).await?;

        let lease_id = lease.lease_id;

        let res = lease
            .with_lock(&mut client.clone(), &lock_name, || async {
                client.lease_revoke(lease_id).await.unwrap();
            })
            .await;

        assert!(matches!(res, Err(Error::Lost)));

        Ok(())
    }

    #[tokio::test]
    async fn test_try_race() -> CrateResult<()> {
        const TTL: u64 = 5;

        let lock_name = create_lock_name();

        let mut client = create_client().await?;
        let locked_count = Arc::new(AtomicUsize::new(0));

        let wait_for_ttl = sleep(Duration::from_secs(TTL + 1));

        stream::iter(0..20)
            .map(|_| {
                let mut client = client.clone();
                let lock_name = lock_name.clone();
                let locked_count = locked_count.clone();

                async move {
                    sleep(Duration::from_secs(random_range(0..TTL - 1))).await;

                    if try_race(&mut client, &lock_name, TtlOrLeaseId::Ttl(TTL as i64)).await? {
                        locked_count.fetch_add(1, Ordering::SeqCst);
                    }

                    CrateResult::Ok(())
                }
            })
            .buffer_unordered(20)
            .try_collect::<Vec<_>>()
            .await?;

        assert_eq!(locked_count.load(Ordering::SeqCst), 1);

        wait_for_ttl.await;
        assert!(try_race(&mut client, &lock_name, TtlOrLeaseId::Ttl(10)).await?);

        Ok(())
    }
}