fix(autoscaler): release distributed locks only for their owner

migrations/cassandra/keyspaces/nvcf_autoscaler/05_add_lock_token.up.sql

@@ -0,0 +1,2 @@
+-- Unique per-acquisition ownership token for conditional lock refresh/release.
+ALTER TABLE nvcf_autoscaler.locks ADD lock_token text;

src/control-plane-services/function-autoscaler/crates/server/src/cassandra/cassandra_service.rs

@@ -23,7 +23,7 @@ use crate::secrets::secrets_config::CassandraSslCertificates;
 use anyhow::{Context, Result};
 use async_trait::async_trait;
 use base64::Engine;
-use chrono::Utc;
+use chrono::{DateTime, Utc};
 use futures::TryStreamExt as _;
 use openssl::ssl::{SslContextBuilder, SslMethod, SslVerifyMode};
 use scylla::client::session::Session;
@@ -524,6 +524,7 @@ impl CassandraServiceManager {
                         &lock.lock_name,
                         &lock.node_id,
                         lock.acquired_at,
+                        &lock.lock_token,
                         ttl_seconds,
                     ),
                 )
@@ -582,6 +583,15 @@ impl CassandraServiceManager {
         node_id: &str,
         ttl_seconds: i32,
     ) -> Result<bool> {
+        let Some(lock) = self.get_lock(lock_name).await? else {
+            return Ok(false);
+        };
+        if lock.node_id != node_id {
+            return Ok(false);
+        }
+        // Keep the acquisition timestamp stable. It distinguishes a later
+        // acquisition by the same node from the lease observed above.
+        let acquired_at = lock.acquired_at;
         let session = self.get_session().await?;
         with_cassandra_timing("refresh_lock_ttl", || async {
             let mut prepared = session
@@ -590,17 +600,27 @@ impl CassandraServiceManager {
             prepared.set_consistency(Consistency::LocalQuorum);
             prepared.set_serial_consistency(Some(SerialConsistency::Serial));
             prepared.set_is_idempotent(false);
-            let refreshed_at = Utc::now();
             let result = session
                 .execute_unpaged(
                     &prepared,
-                    (ttl_seconds, node_id, refreshed_at, lock_name, node_id),
+                    (
+                        ttl_seconds,
+                        node_id,
+                        acquired_at.clone(),
+                        &lock.lock_token,
+                        lock_name,
+                        node_id,
+                        acquired_at,
+                        &lock.lock_token,
+                    ),
                 )
                 .await?;
             let rows = result.into_rows_result()?;
-            // On a CAS miss, Scylla returns [applied]=false + the IF-clause columns (node_id).
-            // Check for that shape first; if node_id is present we know applied=false.
-            if rows.column_specs().get_by_name("node_id").is_some() {
+            // On a CAS miss, Scylla returns [applied]=false plus the IF columns.
+            // Check for that shape first to avoid decoding the returned row as (bool,).
+            if rows.column_specs().get_by_name("node_id").is_some()
+                || rows.column_specs().get_by_name("acquired_at").is_some()
+            {
                 return Ok(false);
             }
             if let Some(row) = rows.rows::<(bool,)>()?.next() {
@@ -636,6 +656,41 @@ impl CassandraServiceManager {
         Ok(())
     }
 
+    /// Delete a lock only when its owner and acquisition generation still match.
+    /// Returns false when the lease expired, changed owner, or was reacquired.
+    #[tracing::instrument(skip(self))]
+    pub async fn delete_owned_lock(
+        &self,
+        lock_name: &str,
+        node_id: &str,
+        acquired_at: DateTime<Utc>,
+        lock_token: &str,
+    ) -> Result<bool> {
+        let session = self.get_session().await?;
+        with_cassandra_timing("delete_owned_lock", || async {
+            let mut prepared = session
+                .prepare(get_delete_owned_lock_stmt(&self.config.keyspace))
+                .await?;
+            prepared.set_consistency(Consistency::LocalQuorum);
+            prepared.set_serial_consistency(Some(SerialConsistency::Serial));
+            prepared.set_is_idempotent(false);
+            let result = session
+                .execute_unpaged(&prepared, (lock_name, node_id, acquired_at, lock_token))
+                .await?;
+            let rows = result.into_rows_result()?;
+            if rows.column_specs().get_by_name("node_id").is_some()
+                || rows.column_specs().get_by_name("acquired_at").is_some()
+            {
+                return Ok(false);
+            }
+            if let Some(row) = rows.rows::<(bool,)>()?.next() {
+                return Ok(row?.0);
+            }
+            Ok(false)
+        })
+        .await
+    }
+
     #[tracing::instrument(skip(self))]
     pub async fn insert_to_nodes(&self, node: &NodeHealth) -> Result<()> {
         let session = self.get_session().await?;
@@ -878,6 +933,7 @@ mod tests {
             lock_name: "test_lock".to_string(),
             node_id: Uuid::new_v4().to_string(),
             acquired_at: Utc::now(),
+            lock_token: Some(Uuid::new_v4().to_string()),
         }
     }
 
@@ -939,11 +995,50 @@ mod tests {
         let get_result = manager.get_lock(&lock.lock_name).await;
         assert!(get_result.is_ok());
         assert!(get_result.unwrap().is_none());
-        // Test insert again- it should succeed
-        let put_result = manager.put_lock(&lock, ttl_seconds).await;
+        // Reacquiring under the same node ID creates a new lock generation.
+        let reacquired_lock = DistributedLock {
+            lock_name: lock.lock_name.clone(),
+            node_id: lock.node_id.clone(),
+            acquired_at: Utc::now(),
+            lock_token: Some(Uuid::new_v4().to_string()),
+        };
+        let put_result = manager.put_lock(&reacquired_lock, ttl_seconds).await;
         assert!(put_result.is_ok());
         assert!(put_result.unwrap());
 
+        // A delayed guard for the previous generation must not delete this lease,
+        // even though the node ID is unchanged.
+        assert!(!manager
+            .delete_owned_lock(
+                &lock.lock_name,
+                &lock.node_id,
+                lock.acquired_at.clone(),
+                lock.lock_token
+                    .as_deref()
+                    .expect("test lock has a token"),
+            )
+            .await
+            .unwrap());
+        let current_lock = manager.get_lock(&lock.lock_name).await.unwrap().unwrap();
+        assert_eq!(
+            current_lock.acquired_at.timestamp_millis(),
+            reacquired_lock.acquired_at.timestamp_millis()
+        );
+
+        assert!(!manager
+            .delete_owned_lock(
+                &lock.lock_name,
+                "another-node",
+                reacquired_lock.acquired_at.clone(),
+                reacquired_lock
+                    .lock_token
+                    .as_deref()
+                    .expect("reacquired test lock has a token"),
+            )
+            .await
+            .unwrap());
+        assert!(manager.get_lock(&lock.lock_name).await.unwrap().is_some());
+
         // Test delete operation
         let delete_result = manager.delete_lock(&lock.lock_name).await;
         assert!(delete_result.is_ok());
@@ -985,12 +1080,17 @@ mod tests {
         // The original acquired_at TTL has expired, but the refresh must keep
         // the complete row readable while the renewed lease remains active.
         tokio::time::sleep(Duration::from_secs(2)).await;
-        assert!(manager.get_lock(&lock.lock_name).await.unwrap().is_some());
+        let refreshed_lock = manager.get_lock(&lock.lock_name).await.unwrap().unwrap();
+        assert_eq!(
+            refreshed_lock.acquired_at.timestamp_millis(),
+            lock.acquired_at.timestamp_millis()
+        );
 
         let contender = DistributedLock {
             lock_name: lock.lock_name.clone(),
             node_id: "contending-node".to_string(),
             acquired_at: Utc::now(),
+            lock_token: Some(Uuid::new_v4().to_string()),
         };
         assert!(!manager.put_lock(&contender, ttl_seconds).await.unwrap());
 

src/control-plane-services/function-autoscaler/crates/server/src/cassandra/distributed_lock.rs

@@ -19,9 +19,10 @@ use crate::cassandra::cassandra_service::CassandraServiceManager;
 use crate::metrics;
 use crate::models::DistributedLock;
 use anyhow::Result;
-use chrono::Utc;
+use chrono::{DateTime, Utc};
 use std::sync::Arc;
 use tracing;
+use uuid::Uuid;
 
 pub struct DistributedLockManager {
     cassandra_service: Arc<CassandraServiceManager>,
@@ -61,13 +62,16 @@ impl DistributedLockManager {
         lock_name: String,
         lock_duration_seconds: i32,
     ) -> Result<bool> {
+        let acquired_at = Utc::now();
+        let lock_token = Uuid::new_v4().to_string();
         let acquired = self
             .cassandra_service
             .put_lock(
                 &DistributedLock {
                     lock_name: lock_name.clone(),
                     node_id: self.node_id.clone(),
-                    acquired_at: Utc::now(),
+                    acquired_at: acquired_at.clone(),
+                    lock_token: Some(lock_token.clone()),
                 },
                 lock_duration_seconds,
             )
@@ -97,6 +101,9 @@ impl DistributedLockManager {
 
 pub struct DistributedLockGuard {
     lock_name: String,
+    node_id: String,
+    acquired_at: DateTime<Utc>,
+    lock_token: String,
     cassandra_service: Arc<CassandraServiceManager>,
     released: bool,
 }
@@ -108,12 +115,15 @@ impl DistributedLockGuard {
         cassandra_service: Arc<CassandraServiceManager>,
         lock_duration_seconds: i32,
     ) -> Result<Option<Self>> {
+        let acquired_at = Utc::now();
+        let lock_token = Uuid::new_v4().to_string();
         let lock_acquired = cassandra_service
             .put_lock(
                 &DistributedLock {
                     lock_name: lock_name.clone(),
-                    node_id,
-                    acquired_at: Utc::now(),
+                    node_id: node_id.clone(),
+                    acquired_at: acquired_at.clone(),
+                    lock_token: Some(lock_token.clone()),
                 },
                 lock_duration_seconds,
             )
@@ -122,6 +132,9 @@ impl DistributedLockGuard {
         if lock_acquired {
             Ok(Some(Self {
                 lock_name,
+                node_id,
+                acquired_at,
+                lock_token,
                 cassandra_service,
                 released: false,
             }))
@@ -136,34 +149,58 @@ impl Drop for DistributedLockGuard {
     fn drop(&mut self) {
         if !self.released {
             let lock_name = self.lock_name.clone();
+            let node_id = self.node_id.clone();
+            let acquired_at = self.acquired_at.clone();
+            let lock_token = self.lock_token.clone();
             let cassandra_service = self.cassandra_service.clone();
 
             tokio::spawn(async move {
-                match cassandra_service.delete_lock(&lock_name).await {
-                    Ok(()) => {
+                match cassandra_service
+                    .delete_owned_lock(&lock_name, &node_id, acquired_at.clone(), &lock_token)
+                    .await
+                {
+                    Ok(true) => {
                         tracing::info!("Released lock {} during drop", lock_name);
                     }
+                    Ok(false) => {
+                        tracing::debug!(
+                            "Lock {} was no longer owned by node {}; leaving current owner intact",
+                            lock_name,
+                            node_id
+                        );
+                    }
                     Err(e) => {
                         tracing::error!("Failed to release lock {} during drop: {}", lock_name, e);
 
-                        // If deletion fails, spawn a task to retry after the lock duration
+                        // If deletion fails, retry conditionally so a new owner is never removed.
                         let retry_lock_name = lock_name.clone();
+                        let retry_node_id = node_id.clone();
+                        let retry_acquired_at = acquired_at.clone();
+                        let retry_lock_token = lock_token.clone();
                         let retry_cassandra_service = cassandra_service.clone();
                         tokio::spawn(async move {
                             tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
-                            if let Err(e) =
-                                retry_cassandra_service.delete_lock(&retry_lock_name).await
+                            match retry_cassandra_service
+                                .delete_owned_lock(
+                                    &retry_lock_name,
+                                    &retry_node_id,
+                                    retry_acquired_at,
+                                    &retry_lock_token,
+                                )
+                                .await
                             {
-                                tracing::error!(
-                                    "Failed to release lock {} during retry: {}",
-                                    retry_lock_name,
-                                    e
-                                );
-                            } else {
-                                tracing::info!(
+                                Ok(true) => tracing::info!(
                                     "Successfully released lock {} during retry",
                                     retry_lock_name
-                                );
+                                ),
+                                Ok(false) => tracing::debug!(
+                                    "Lock {} changed owner before retry; leaving current owner intact",
+                                    retry_lock_name
+                                ),
+                                Err(e) => tracing::error!(
+                                    "Failed to release lock {} during retry: {}",
+                                    retry_lock_name, e
+                                ),
                             }
                         });
                     }