use {
super::{
Revision,
lease::{ManagedLease, ManagedLeaseFactory},
retry::retry_etcd_legacy,
},
crate::{
lease::{LeaseExpiredNotify, ManagedLeaseWeak},
retry::{retry_etcd, retry_etcd_txn},
watcher::WatchClientExt,
},
core::fmt,
etcd_client::{Compare, CompareOp, GetOptions, LockOptions, Txn, TxnOp, TxnResponse},
futures::{
FutureExt, StreamExt,
future::{BoxFuture, Shared, join_all},
},
retry::delay::Fixed,
std::{
future::Future,
pin::Pin,
task::{Context, Poll},
time::Duration,
},
thiserror::Error,
tokio::{
sync::mpsc,
task::{JoinError, JoinHandle},
},
tonic::Code,
tracing::{info, trace},
};
enum DeleteQueueCommand {
Delete(Vec<u8>),
}
#[derive(Clone)]
pub struct LockManager {
etcd: etcd_client::Client,
delete_queue_tx: mpsc::UnboundedSender<DeleteQueueCommand>,
manager_lease_factory: ManagedLeaseFactory,
try_locking_timeout: Duration,
#[allow(dead_code)]
lock_manager_handle_entangled_tx: mpsc::UnboundedSender<()>,
}
pub struct LockManagerHandle {
inner: JoinHandle<()>,
}
impl Future for LockManagerHandle {
type Output = Result<(), JoinError>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.inner.poll_unpin(cx)
}
}
pub struct ManagedLockRevokeNotify {
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.clone(),
lease_expired_notify: self.lease_expired_notify.clone(),
}
}
}
impl ManagedLockRevokeNotify {
pub async fn wait_for_revoke(self) {
let watch_lock_delete = self.watch_lock_delete;
tokio::select! {
_ = self.lease_expired_notify.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()
}
pub fn spawn_lock_manager(etcd: etcd_client::Client) -> (LockManagerHandle, LockManager) {
let (lease_factory, _) = ManagedLeaseFactory::spawn(etcd.clone());
spawn_lock_manager_with_lease_factory(etcd, lease_factory)
}
pub fn spawn_lock_manager_with_lease_factory(
etcd: etcd_client::Client,
managed_lease_factory: ManagedLeaseFactory,
) -> (LockManagerHandle, LockManager) {
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let etcd2 = etcd.clone();
let (entangled_tx, mut entangled_rx) = tokio::sync::mpsc::unbounded_channel();
let tx2 = tx.clone();
let handle = tokio::spawn(async move {
let _tx2 = tx2;
loop {
let cmd = tokio::select! {
cmd = rx.recv() => {
cmd.expect("command rx droped")
}
maybe = entangled_rx.recv() => {
match maybe {
Some(_) => unreachable!("entangled_rx should not have any message"),
None => {
break
},
}
}
};
match cmd {
DeleteQueueCommand::Delete(lock_id) => {
let kv_client = etcd2.kv_client();
let lock_id2 = lock_id.clone();
let result = retry_etcd_legacy(Fixed::from_millis(10), move || {
let lock_id = lock_id2.clone();
let mut kv_client = kv_client.clone();
async move { kv_client.delete(lock_id, None).await }
})
.await;
match result {
Ok(_) => {
let lock_id = String::from_utf8(lock_id).expect("lock id is not utf8");
info!("Deleted lock {lock_id}");
}
Err(e) => {
if !matches!(e, etcd_client::Error::GRpcStatus(ref status) if status.code() == Code::NotFound)
{
tracing::error!("Failed to revoke lock: {e}");
}
}
}
}
}
}
let mut futures = vec![];
while let Ok(cmd) = rx.try_recv() {
match cmd {
DeleteQueueCommand::Delete(lock_id) => {
let mut kv_client = etcd2.kv_client();
let fut = async move { kv_client.delete(lock_id, None).await };
futures.push(fut);
}
}
}
let _ = join_all(futures).await;
});
let handle = LockManagerHandle { inner: handle };
(
handle,
LockManager {
etcd,
delete_queue_tx: tx,
try_locking_timeout: Duration::from_secs(1),
manager_lease_factory: managed_lease_factory,
lock_manager_handle_entangled_tx: entangled_tx,
},
)
}
#[derive(Debug, thiserror::Error)]
pub enum LockingError {
#[error("Etcd error: {0:?}")]
EtcdError(etcd_client::Error),
}
impl LockManager {
pub async fn try_lock<S>(
&self,
name: S,
lease_duration: Duration,
) -> Result<ManagedLock, TryLockError>
where
S: AsRef<str>,
{
let name = name.as_ref();
if self.delete_queue_tx.is_closed() {
panic!("LockManager lifecycle thread is stopped.");
}
let gopts = GetOptions::new().with_prefix();
trace!("Trying to lock {name}...");
let get_response = retry_etcd(
self.etcd.clone(),
(name.to_string(), gopts),
move |etcd, (name, gopts)| async move { etcd.kv_client().get(name, Some(gopts)).await },
)
.await
.map_err(TryLockError::EtcdError)?;
if get_response.count() > 0 {
return Err(TryLockError::AlreadyTaken);
}
let managed_lease = self
.manager_lease_factory
.new_lease(lease_duration, None)
.await
.map_err(TryLockError::EtcdError)?;
let lease_id = managed_lease.lease_id;
let lock_fut = retry_etcd(
self.etcd.clone(),
(name.to_string(), LockOptions::new().with_lease(lease_id)),
|mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
);
let lease_expire_notify = managed_lease.get_lease_expire_notify();
let (revision, lock_key) = tokio::select! {
_ = tokio::time::sleep(self.try_locking_timeout) => {
return Err(TryLockError::LockingDeadlineExceeded)
}
result = lock_fut => {
let lock_response = match result {
Ok(lock_response) => lock_response,
Err(e) => {
match e {
etcd_client::Error::GRpcStatus(status) => {
if status.code() == Code::Unknown {
if status.message() == "etcdserver: requested lease not found" {
return Err(TryLockError::LeaseExpired)
} else {
return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
}
} else {
return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
}
}
_ => return Err(TryLockError::EtcdError(e))
}
}
};
(lock_response.header().expect("empty header for etcd lock").revision(), lock_response.key().to_vec())
}
_ = lease_expire_notify.recv() => {
return Err(TryLockError::LeaseExpired)
}
};
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,
})
}
pub async fn try_lock_with_lease<S>(
&self,
name: S,
managed_lease: ManagedLease,
) -> Result<ManagedLock, TryLockError>
where
S: AsRef<str>,
{
let name = name.as_ref();
if self.delete_queue_tx.is_closed() {
panic!("LockManager lifecycle thread is stopped.");
}
let gopts = GetOptions::new().with_prefix();
const TRY_LOCKING_DURATION: Duration = Duration::from_millis(1000);
trace!("Trying to lock {name}...");
let get_response = retry_etcd(
self.etcd.clone(),
(name.to_string(), gopts),
move |etcd, (name, gopts)| async move { etcd.kv_client().get(name, Some(gopts)).await },
)
.await
.map_err(TryLockError::EtcdError)?;
if get_response.count() > 0 {
return Err(TryLockError::AlreadyTaken);
}
let lease_id = managed_lease.lease_id;
let lock_fut = retry_etcd(
self.etcd.clone(),
(name.to_string(), LockOptions::new().with_lease(lease_id)),
|mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
);
let lease_expire_notify = managed_lease.get_lease_expire_notify();
let (revision, lock_key) = tokio::select! {
_ = tokio::time::sleep(TRY_LOCKING_DURATION) => {
return Err(TryLockError::LockingDeadlineExceeded)
}
result = lock_fut => {
let lock_response = match result {
Ok(lock_response) => lock_response,
Err(e) => {
match e {
etcd_client::Error::GRpcStatus(status) => {
if status.code() == Code::Unknown {
if status.message() == "etcdserver: requested lease not found" {
return Err(TryLockError::LeaseExpired)
} else {
return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
}
} else {
return Err(TryLockError::EtcdError(etcd_client::Error::GRpcStatus(status)))
}
}
_ => return Err(TryLockError::EtcdError(e))
}
}
};
(lock_response.header().expect("empty header for etcd lock").revision(), lock_response.key().to_vec())
}
_ = lease_expire_notify.recv() => {
return Err(TryLockError::LeaseExpired)
}
};
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,
})
}
pub async fn lock<S>(
&self,
name: S,
lease_duration: Duration,
) -> Result<ManagedLock, etcd_client::Error>
where
S: AsRef<str>,
{
if self.delete_queue_tx.is_closed() {
panic!("LockManager lifecycle thread is stopped.");
}
let managed_lease = self
.manager_lease_factory
.new_lease(lease_duration, None)
.await?;
self.lock_with_lease(name, managed_lease).await
}
pub async fn lock_with_lease<S>(
&self,
name: S,
managed_lease: ManagedLease,
) -> Result<ManagedLock, etcd_client::Error>
where
S: AsRef<str>,
{
if self.delete_queue_tx.is_closed() {
panic!("LockManager lifecycle thread is stopped.");
}
let name = name.as_ref();
let lease_id = managed_lease.lease_id;
let lock_fut = retry_etcd(
self.etcd.clone(),
(name.to_string(), LockOptions::new().with_lease(lease_id)),
|mut etcd, (name, opts)| async move { etcd.lock(name, Some(opts)).await },
);
let lock_response = tokio::select! {
result = lock_fut => {
result?
}
};
let (revision, lock_key) = (
lock_response
.header()
.expect("empty header for etcd lock")
.revision(),
lock_response.key().to_vec(),
);
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,
};
Ok(managed_lock)
}
}
pub struct ManagedLock {
pub(crate) lock_key: Vec<u8>,
managed_lease: ManagedLease,
pub created_at_revision: Revision,
pub(crate) etcd: etcd_client::Client,
delete_signal_tx: tokio::sync::mpsc::UnboundedSender<DeleteQueueCommand>,
revoke_callback: Shared<BoxFuture<'static, ()>>,
}
impl fmt::Debug for ManagedLock {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ManagedLock")
.field("lock_key", &String::from_utf8_lossy(&self.lock_key))
.field("lease_id", &self.managed_lease.lease_id)
.field("created_At_revision", &self.created_at_revision)
.finish()
}
}
impl Drop for ManagedLock {
fn drop(&mut self) {
info!(
"Destructor called for ManagedLock({})",
String::from_utf8_lossy(&self.lock_key)
);
let _ = self
.delete_signal_tx
.send(DeleteQueueCommand::Delete(self.lock_key.clone()));
}
}
#[derive(Debug)]
pub enum LockError {
LockRevoked,
}
impl fmt::Display for LockError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
LockError::LockRevoked => f.write_str("lock revoked"),
}
}
}
pub struct ManagedLockGuard<'a> {
pub(crate) managed_lock: &'a ManagedLock,
}
impl ManagedLockGuard<'_> {
pub(crate) fn get_key(&self) -> &[u8] {
self.managed_lock.lock_key.as_slice()
}
}
impl ManagedLock {
pub fn lease_id(&self) -> i64 {
self.managed_lease.lease_id
}
pub async fn txn(&self, operations: impl Into<Vec<TxnOp>>) -> TxnResponse {
let txn = Txn::new()
.when(vec![Compare::version(
self.lock_key.clone(),
CompareOp::Greater,
0,
)])
.and_then(operations);
retry_etcd_txn(self.etcd.clone(), txn)
.await
.expect("failed txn")
}
pub fn get_revoke_notify(&self) -> ManagedLockRevokeNotify {
ManagedLockRevokeNotify {
watch_lock_delete: self.revoke_callback.clone(),
lease_expired_notify: self.managed_lease.get_lease_expire_notify(),
}
}
pub async fn is_alive(&self) -> bool {
let get_response = self
.etcd
.kv_client()
.get(self.lock_key.as_slice(), None)
.await
.expect("failed to communicate with etcd");
get_response.count() == 1
}
pub fn get_key(&self) -> Vec<u8> {
self.lock_key.clone()
}
pub async fn scope<T, Fut>(&self, fut: Fut) -> Result<T, LockError>
where
T: Send + 'static,
Fut: Future<Output = T> + Send + 'static,
{
self.scope_with(move |_| fut).await
}
pub async fn scope_with<'a, T, F, Fut>(&'a self, func: F) -> Result<T, LockError>
where
T: Send + 'a,
F: FnOnce(ManagedLockGuard<'a>) -> Fut,
Fut: Future<Output = T> + Send + 'a,
{
let revoke_callback = self.revoke_callback.clone();
tokio::select! {
result = func(ManagedLockGuard { managed_lock: self }) => Ok(result),
_ = revoke_callback => Err(LockError::LockRevoked),
}
}
pub fn get_managed_lease_weak_ref(&self) -> ManagedLeaseWeak {
self.managed_lease.get_weak()
}
pub async fn into_revoked_fut(self) {
let _ = self.scope(futures::future::pending::<()>()).await;
}
}
#[derive(Debug, Error)]
pub enum TryLockError {
#[error("Already taken")]
AlreadyTaken,
#[error("Locking deadline exceeded")]
LockingDeadlineExceeded,
#[error("Lease expired before the lock")]
LeaseExpired,
#[error("Etcd error: {0:?}")]
EtcdError(etcd_client::Error),
}