Skip to content

GitLab

Unverified Commit c78b5901 authored by Graham King's avatar Graham King Committed by GitHub
Browse files

chore(runtime): Do not expose etcd lease ID (#3915) 


Signed-off-by: default avatarGraham King <grahamk@nvidia.com>
parent a79122c6
Hide whitespace changes
Inline Side-by-side
Showing with 19 additions and 107 deletions
lib/runtime/src/transports/etcd.rs
View file @ c78b5901
......@@ -32,47 +32,6 @@ pub use path::*;
use super::utils::build_in_runtime;
#[derive(Debug, Clone)]
pub struct Lease {
/// ETCD lease ID
/// Delivered as i64 by etcd because of documented gRPC limitations.
id: u64,
/// [`CancellationToken`] associated with the lease
cancel_token: CancellationToken,
}
impl Lease {
/// Get the lease ID
pub fn id(&self) -> u64 {
self.id
}
/// Get the primary [`CancellationToken`] associated with the lease.
/// This token will revoke the lease if canceled.
pub fn primary_token(&self) -> CancellationToken {
self.cancel_token.clone()
}
/// Get a child [`CancellationToken`] from the lease's [`CancellationToken`].
/// This child token will be triggered if the lease is revoked, but will not revoke the lease if canceled.
pub fn child_token(&self) -> CancellationToken {
self.cancel_token.child_token()
}
/// Revoke the lease triggering the [`CancellationToken`].
pub fn revoke(&self) {
self.cancel_token.cancel();
}
/// Check if the lease is still valid (not revoked)
pub async fn is_valid(&self) -> Result<bool> {
// A lease is valid if its cancellation token has not been triggered
// We can use try_cancelled which returns immediately with a boolean
Ok(!self.cancel_token.is_cancelled())
}
}
/// ETCD Client
#[derive(Clone)]
pub struct Client {
......@@ -120,16 +79,14 @@ impl Client {
})?;
let lease_id = if config.attach_lease {
let lease = create_lease(connector.clone(), 10, token)
create_lease(connector.clone(), 10, token)
.await
.with_context(|| {
format!(
"Unable to create lease. Check etcd server status at {}",
config.etcd_url.join(", ")
)
})?;
lease.id
})?
} else {
0
};
......@@ -159,30 +116,6 @@ impl Client {
self.primary_lease
}
/// Primary [`Lease`]
pub fn primary_lease(&self) -> Lease {
Lease {
id: self.primary_lease,
cancel_token: self.runtime.primary_token(),
}
}
/// Create a [`Lease`] with a given time-to-live (TTL).
/// This [`Lease`] will be tied to the [`Runtime`], specifically a child [`CancellationToken`].
pub async fn create_lease(&self, ttl: u64) -> Result<Lease> {
let token = self.runtime.child_token();
self.rt
.spawn(create_lease(self.connector.clone(), ttl, token))
.await?
}
// Revoke an etcd lease given its lease id. A wrapper over etcd_client::LeaseClient::revoke
pub async fn revoke_lease(&self, lease_id: u64) -> Result<()> {
self.rt
.spawn(revoke_lease(self.connector.clone(), lease_id))
.await?
}
pub async fn kv_create(&self, key: &str, value: Vec<u8>, lease_id: Option<u64>) -> Result<()> {
let id = lease_id.unwrap_or(self.lease_id());
let put_options = PutOptions::new().with_lease(id as i64);
......@@ -284,7 +217,7 @@ impl Client {
) -> Result<PutResponse> {
let options = options
.unwrap_or_default()
.with_lease(self.primary_lease().id() as i64);
.with_lease(self.lease_id() as i64);
self.connector
.get_client()
.kv_client()
......@@ -702,10 +635,7 @@ mod tests {
let value = b"test_value";
let client = drt.etcd_client().expect("etcd client should be available");
let lease_id = drt
.primary_lease()
.expect("primary lease should be available")
.id();
let lease_id = drt.connection_id();
// Create the key
let result = client.kv_create(key, value.to_vec(), Some(lease_id)).await;
......
lib/runtime/src/transports/etcd/lease.rs
View file @ c78b5901
......@@ -2,24 +2,26 @@
// SPDX-License-Identifier: Apache-2.0
use super::connector::Connector;
use super::*;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::{sleep, timeout};
use std::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;
/// Create a [`Lease`] with a given time-to-live (TTL) attached to the [`CancellationToken`].
/// Create an etcd lease with the given TTL, attach it to the provided cancellation token,
/// spawn a keep-alive task, and return the lease id (u64).
///
/// Note: this function spawns a background task that maintains the lease until the token is
/// cancelled or an unrecoverable error occurs.
pub async fn create_lease(
connector: Arc<Connector>,
ttl: u64,
token: CancellationToken,
) -> Result<Lease> {
) -> anyhow::Result<u64> {
let mut lease_client = connector.get_client().lease_client();
let lease = lease_client.grant(ttl as i64, None).await?;
let id = lease.id() as u64;
let ttl = lease.ttl() as u64;
let child = token.child_token();
let clone = token.clone();
tokio::spawn(async move {
match keep_alive(connector, id, ttl, child).await {
......@@ -34,22 +36,7 @@ pub async fn create_lease(
}
});
Ok(Lease {
id,
cancel_token: clone,
})
}
/// Revoke a lease given its lease id. A wrapper over etcd_client::LeaseClient::revoke
pub async fn revoke_lease(connector: Arc<Connector>, lease_id: u64) -> Result<()> {
let mut lease_client = connector.get_client().lease_client();
match lease_client.revoke(lease_id as i64).await {
Ok(_) => Ok(()),
Err(e) => {
tracing::warn!("failed to revoke lease: {:?}", e);
Err(e.into())
}
}
Ok(id)
}
/// Task to keep leases alive with reconnection support.
......@@ -60,8 +47,8 @@ async fn keep_alive(
lease_id: u64,
mut ttl: u64,
token: CancellationToken,
) -> Result<()> {
let mut deadline = create_deadline(ttl)?;
) -> anyhow::Result<()> {
let mut deadline = Instant::now() + Duration::from_secs(ttl);
loop {
// Try to establish or re-establish the keep-alive stream
......@@ -99,9 +86,9 @@ async fn keep_alive(
// Keep-alive loop with the established stream
loop {
if deadline < std::time::Instant::now() {
return Err(error!(
anyhow::bail!(
"Unable to refresh lease - deadline exceeded. Check etcd server status"
));
);
}
tokio::select! {
......@@ -114,10 +101,10 @@ async fn keep_alive(
// Update ttl and deadline from response
ttl = resp.ttl() as u64;
deadline = create_deadline(ttl)?;
deadline = Instant::now() + Duration::from_secs(ttl);
if resp.ttl() == 0 {
return Err(error!("Unable to maintain lease - expired or revoked. Check etcd server status"));
anyhow::bail!("Unable to maintain lease - expired or revoked. Check etcd server status");
}
}
Ok(None) => {
......@@ -164,8 +151,3 @@ async fn keep_alive(
}
}
}
/// Create a deadline for a given time-to-live (TTL).
fn create_deadline(ttl: u64) -> Result<std::time::Instant> {
Ok(std::time::Instant::now() + std::time::Duration::from_secs(ttl))
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment