diff --git a/migrations/cassandra/keyspaces/nvcf_autoscaler/05_add_lock_token.up.sql b/migrations/cassandra/keyspaces/nvcf_autoscaler/05_add_lock_token.up.sql new file mode 100644 index 0000000000..cac9c4c864 --- /dev/null +++ b/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; diff --git a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/cassandra_service.rs b/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/cassandra_service.rs index 0aef0148b6..dd6a1b1d2c 100644 --- a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/cassandra_service.rs +++ b/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 { + 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, + lock_token: &str, + ) -> Result { + 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()); diff --git a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/distributed_lock.rs b/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/distributed_lock.rs index 05ff3d266d..110abff5f6 100644 --- a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/distributed_lock.rs +++ b/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, @@ -61,13 +62,16 @@ impl DistributedLockManager { lock_name: String, lock_duration_seconds: i32, ) -> Result { + 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, + lock_token: String, cassandra_service: Arc, released: bool, } @@ -108,12 +115,15 @@ impl DistributedLockGuard { cassandra_service: Arc, lock_duration_seconds: i32, ) -> Result> { + 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 + ), } }); } diff --git a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/statements.rs b/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/statements.rs index 47cc7dc90b..9fc5f0b48d 100644 --- a/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/statements.rs +++ b/src/control-plane-services/function-autoscaler/crates/server/src/cassandra/statements.rs @@ -18,7 +18,7 @@ // locks table pub(crate) fn get_select_locks_stmt(keyspace: &str) -> String { format!( - "SELECT lock_name, node_id, acquired_at FROM {}.locks WHERE lock_name = ?;", + "SELECT lock_name, node_id, acquired_at, lock_token FROM {}.locks WHERE lock_name = ?;", keyspace ) } @@ -27,6 +27,13 @@ pub(crate) fn get_delete_locks_stmt(keyspace: &str) -> String { format!("DELETE FROM {}.locks WHERE lock_name = ?;", keyspace) } +pub(crate) fn get_delete_owned_lock_stmt(keyspace: &str) -> String { + format!( + "DELETE FROM {}.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?;", + keyspace + ) +} + // healthy_nodes table pub(crate) fn get_select_all_nodes_stmt(keyspace: &str) -> String { format!( @@ -55,22 +62,23 @@ pub(crate) fn get_health_check_query_stmt(keyspace: &str) -> String { } // Inserts to the locks table must be done with a row TTL. -// Bind order: (lock_name, node_id, acquired_at, ttl_seconds) +// Bind order: (lock_name, node_id, acquired_at, lock_token, ttl_seconds) pub(crate) fn get_stmt_insert_to_locks(keyspace: &str) -> String { format!( - "INSERT INTO {}.locks (lock_name, node_id, acquired_at) VALUES (?, ?, ?) IF NOT EXISTS USING TTL ?", + "INSERT INTO {}.locks (lock_name, node_id, acquired_at, lock_token) VALUES (?, ?, ?, ?) IF NOT EXISTS USING TTL ?", keyspace, ) } -// LWT conditional update — only refreshes TTL if node_id still matches this node. +// LWT conditional update — only refreshes TTL if this acquisition still owns the lock. // Refresh every non-key lock column so Cassandra's per-cell TTL semantics -// cannot leave a partially expired row behind. -// Returns [applied]=true if the row was updated, false if another node now owns the lock. -// Bind order: (ttl_seconds, node_id, acquired_at, lock_name, node_id) +// cannot leave a partially expired row behind. Keep acquired_at unchanged: it is +// the generation token used to distinguish a reacquisition by the same node. +// Returns [applied]=true if the row was updated, false if ownership changed. +// Bind order: (ttl_seconds, node_id, acquired_at, lock_token, lock_name, node_id, acquired_at, lock_token) pub(crate) fn get_stmt_refresh_lock(keyspace: &str) -> String { format!( - "UPDATE {}.locks USING TTL ? SET node_id = ?, acquired_at = ? WHERE lock_name = ? IF node_id = ?", + "UPDATE {}.locks USING TTL ? SET node_id = ?, acquired_at = ?, lock_token = ? WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?", keyspace, ) } @@ -99,7 +107,7 @@ pub(crate) fn get_stmt_insert_to_recently_invoked_functions( #[cfg(test)] mod tests { use super::{ - get_select_recently_invoked_functions_in_token_range_stmt, + get_delete_owned_lock_stmt, get_select_recently_invoked_functions_in_token_range_stmt, get_stmt_insert_to_recently_invoked_functions, get_stmt_refresh_lock, }; @@ -107,8 +115,16 @@ mod tests { fn refresh_lock_renews_every_non_key_column() { let statement = get_stmt_refresh_lock("test_keyspace"); - assert!(statement.contains("SET node_id = ?, acquired_at = ?")); - assert!(statement.contains("IF node_id = ?")); + assert!(statement.contains("SET node_id = ?, acquired_at = ?, lock_token = ?")); + assert!(statement.contains("IF node_id = ? AND acquired_at = ? AND lock_token = ?")); + } + + #[test] + fn delete_owned_lock_checks_node_and_acquisition_generation() { + assert_eq!( + get_delete_owned_lock_stmt("test_keyspace"), + "DELETE FROM test_keyspace.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?;" + ); } #[test] diff --git a/src/control-plane-services/function-autoscaler/crates/server/src/models/mod.rs b/src/control-plane-services/function-autoscaler/crates/server/src/models/mod.rs index b9f086c8f0..412ec71094 100644 --- a/src/control-plane-services/function-autoscaler/crates/server/src/models/mod.rs +++ b/src/control-plane-services/function-autoscaler/crates/server/src/models/mod.rs @@ -53,6 +53,7 @@ pub struct DistributedLock { pub lock_name: String, pub node_id: String, pub acquired_at: DateTime, + pub lock_token: Option, } #[derive(Debug, Clone, DeserializeRow)] @@ -71,4 +72,5 @@ pub struct DistributedLockResult { pub lock_name: Option, pub node_id: Option, pub acquired_at: Option>, + pub lock_token: Option, } diff --git a/src/control-plane-services/function-autoscaler/local_env/cassandra/schema/0001_initial_schema.cql b/src/control-plane-services/function-autoscaler/local_env/cassandra/schema/0001_initial_schema.cql index 5e2c5a412d..496b08a905 100644 --- a/src/control-plane-services/function-autoscaler/local_env/cassandra/schema/0001_initial_schema.cql +++ b/src/control-plane-services/function-autoscaler/local_env/cassandra/schema/0001_initial_schema.cql @@ -28,7 +28,8 @@ CREATE TABLE IF NOT EXISTS nvcf_autoscaler.recently_invoked_functions ( CREATE TABLE IF NOT EXISTS nvcf_autoscaler.locks ( lock_name text PRIMARY KEY, node_id text, - acquired_at timestamp + acquired_at timestamp, + lock_token text ) WITH default_time_to_live = 600; CREATE TABLE IF NOT EXISTS nvcf_autoscaler.healthy_nodes (