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

Signed-off-by: Graham King <grahamk@nvidia.com>

lib/runtime/src/transports/etcd.rs

@@ -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

@@ -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))
-}