"lib/vscode:/vscode.git/clone" did not exist on "bd91dca6141e05bcfbe9bd4dea54cc58b9e37d75"
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
...@@ -32,47 +32,6 @@ pub use path::*; ...@@ -32,47 +32,6 @@ pub use path::*;
use super::utils::build_in_runtime; 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 /// ETCD Client
#[derive(Clone)] #[derive(Clone)]
pub struct Client { pub struct Client {
...@@ -120,16 +79,14 @@ impl Client { ...@@ -120,16 +79,14 @@ impl Client {
})?; })?;
let lease_id = if config.attach_lease { let lease_id = if config.attach_lease {
let lease = create_lease(connector.clone(), 10, token) create_lease(connector.clone(), 10, token)
.await .await
.with_context(|| { .with_context(|| {
format!( format!(
"Unable to create lease. Check etcd server status at {}", "Unable to create lease. Check etcd server status at {}",
config.etcd_url.join(", ") config.etcd_url.join(", ")
) )
})?; })?
lease.id
} else { } else {
0 0
}; };
...@@ -159,30 +116,6 @@ impl Client { ...@@ -159,30 +116,6 @@ impl Client {
self.primary_lease 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<()> { 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 id = lease_id.unwrap_or(self.lease_id());
let put_options = PutOptions::new().with_lease(id as i64); let put_options = PutOptions::new().with_lease(id as i64);
...@@ -284,7 +217,7 @@ impl Client { ...@@ -284,7 +217,7 @@ impl Client {
) -> Result<PutResponse> { ) -> Result<PutResponse> {
let options = options let options = options
.unwrap_or_default() .unwrap_or_default()
.with_lease(self.primary_lease().id() as i64); .with_lease(self.lease_id() as i64);
self.connector self.connector
.get_client() .get_client()
.kv_client() .kv_client()
...@@ -702,10 +635,7 @@ mod tests { ...@@ -702,10 +635,7 @@ mod tests {
let value = b"test_value"; let value = b"test_value";
let client = drt.etcd_client().expect("etcd client should be available"); let client = drt.etcd_client().expect("etcd client should be available");
let lease_id = drt let lease_id = drt.connection_id();
.primary_lease()
.expect("primary lease should be available")
.id();
// Create the key // Create the key
let result = client.kv_create(key, value.to_vec(), Some(lease_id)).await; let result = client.kv_create(key, value.to_vec(), Some(lease_id)).await;
......
...@@ -2,24 +2,26 @@ ...@@ -2,24 +2,26 @@
// SPDX-License-Identifier: Apache-2.0 // SPDX-License-Identifier: Apache-2.0
use super::connector::Connector; use super::connector::Connector;
use super::*;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::{Duration, Instant};
use tokio::time::{sleep, timeout}; 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( pub async fn create_lease(
connector: Arc<Connector>, connector: Arc<Connector>,
ttl: u64, ttl: u64,
token: CancellationToken, token: CancellationToken,
) -> Result<Lease> { ) -> anyhow::Result<u64> {
let mut lease_client = connector.get_client().lease_client(); let mut lease_client = connector.get_client().lease_client();
let lease = lease_client.grant(ttl as i64, None).await?; let lease = lease_client.grant(ttl as i64, None).await?;
let id = lease.id() as u64; let id = lease.id() as u64;
let ttl = lease.ttl() as u64; let ttl = lease.ttl() as u64;
let child = token.child_token(); let child = token.child_token();
let clone = token.clone();
tokio::spawn(async move { tokio::spawn(async move {
match keep_alive(connector, id, ttl, child).await { match keep_alive(connector, id, ttl, child).await {
...@@ -34,22 +36,7 @@ pub async fn create_lease( ...@@ -34,22 +36,7 @@ pub async fn create_lease(
} }
}); });
Ok(Lease { Ok(id)
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())
}
}
} }
/// Task to keep leases alive with reconnection support. /// Task to keep leases alive with reconnection support.
...@@ -60,8 +47,8 @@ async fn keep_alive( ...@@ -60,8 +47,8 @@ async fn keep_alive(
lease_id: u64, lease_id: u64,
mut ttl: u64, mut ttl: u64,
token: CancellationToken, token: CancellationToken,
) -> Result<()> { ) -> anyhow::Result<()> {
let mut deadline = create_deadline(ttl)?; let mut deadline = Instant::now() + Duration::from_secs(ttl);
loop { loop {
// Try to establish or re-establish the keep-alive stream // Try to establish or re-establish the keep-alive stream
...@@ -99,9 +86,9 @@ async fn keep_alive( ...@@ -99,9 +86,9 @@ async fn keep_alive(
// Keep-alive loop with the established stream // Keep-alive loop with the established stream
loop { loop {
if deadline < std::time::Instant::now() { if deadline < std::time::Instant::now() {
return Err(error!( anyhow::bail!(
"Unable to refresh lease - deadline exceeded. Check etcd server status" "Unable to refresh lease - deadline exceeded. Check etcd server status"
)); );
} }
tokio::select! { tokio::select! {
...@@ -114,10 +101,10 @@ async fn keep_alive( ...@@ -114,10 +101,10 @@ async fn keep_alive(
// Update ttl and deadline from response // Update ttl and deadline from response
ttl = resp.ttl() as u64; ttl = resp.ttl() as u64;
deadline = create_deadline(ttl)?; deadline = Instant::now() + Duration::from_secs(ttl);
if resp.ttl() == 0 { 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) => { Ok(None) => {
...@@ -164,8 +151,3 @@ async fn keep_alive( ...@@ -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