From 644f69903c1fd88e09ad4253868b112866680f05 Mon Sep 17 00:00:00 2001 From: Hossein Dehghan Date: Fri, 9 Oct 2026 15:17:04 +0000 Subject: [PATCH 1/6] fix(autoscaler): release distributed locks conditionally Use a Cassandra conditional delete in DistributedLockGuard::drop so an expired guard cannot delete a lock reacquired by another node. Keep the retry conditional and add statement and integration regression coverage. Signed-off-by: Hossein Dehghan --- .../crates/server/src/cassandra/statements.rs | 240 +++++++++++++++++- 1 file changed, 238 insertions(+), 2 deletions(-) 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..b6b36b0161 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 @@ -1,123 +1,359 @@ /* + * SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + + // 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 = ?;", + keyspace + ) -} + + } + + 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 = ?;", + + keyspace + + ) + } + + // healthy_nodes table + pub(crate) fn get_select_all_nodes_stmt(keyspace: &str) -> String { + format!( + "SELECT node_id, last_updated_at FROM {}.healthy_nodes;", + keyspace + ) + } + + pub(crate) fn get_delete_node_stmt(keyspace: &str) -> String { + format!("DELETE FROM {}.healthy_nodes WHERE node_id = ?;", keyspace) + } + + // recently_invoked_functions Table + // account_id is a regular column, not part of PK + pub(crate) fn get_select_recently_invoked_functions_in_token_range_stmt(keyspace: &str) -> String { + format!( + "SELECT function_id, function_version_id, account_id \ + FROM {}.recently_invoked_functions \ + WHERE token(function_id, function_version_id) >= ? AND token(function_id, function_version_id) <= ?;", + keyspace + ) + } + + pub(crate) fn get_health_check_query_stmt(keyspace: &str) -> String { + format!("SELECT now() from {}.healthy_nodes LIMIT 1;", keyspace) + } + + // Inserts to the locks table must be done with a row TTL. + // Bind order: (lock_name, node_id, acquired_at, 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 ?", + keyspace, + ) + } + + // LWT conditional update — only refreshes TTL if node_id still matches this node. + // 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) + 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 = ?", + keyspace, + ) + } + + // Inserts to the healthy_nodes table with a configurable row TTL (from CassandraSettings.node_health_ttl_seconds, default 180s). + // The node is pruned automatically after TTL expires if it stops reporting healthy. + pub(crate) fn get_stmt_insert_to_nodes(keyspace: &str, ttl_seconds: i32) -> String { + format!( + "INSERT INTO {}.healthy_nodes (node_id, last_updated_at) VALUES (?, ?) USING TTL {}", + keyspace, ttl_seconds + ) + } + + // Inserts to the recently_invoked_functions table with a configurable row TTL (from CassandraSettings.recently_invoked_ttl_seconds, default 1800s). + pub(crate) fn get_stmt_insert_to_recently_invoked_functions( + keyspace: &str, + ttl_seconds: i32, + ) -> String { + format!( + "INSERT INTO {}.recently_invoked_functions (function_id, function_version_id, account_id, last_updated_at) VALUES (?, ?, ?, ?) USING TTL {}", + keyspace, + ttl_seconds + ) + } + + #[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, + }; + + #[test] + 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 = ?")); + } + + #[test] + + fn delete_owned_lock_checks_node_id() { + + assert_eq!( + + get_delete_owned_lock_stmt("test_keyspace"), + + "DELETE FROM test_keyspace.locks WHERE lock_name = ? IF node_id = ?;" + + ); + + } + + + + #[test] + fn active_function_statements_preserve_table_and_ttl() { + let select = get_select_recently_invoked_functions_in_token_range_stmt("test_keyspace"); + let insert = get_stmt_insert_to_recently_invoked_functions("test_keyspace", 1800); + + assert!(select.contains("FROM test_keyspace.recently_invoked_functions")); + assert!(insert.contains("INTO test_keyspace.recently_invoked_functions")); + assert!(insert.ends_with("USING TTL 1800")); + } + } + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + From 37f708e054e587c41df4cf5b3067309d239b7ae2 Mon Sep 17 00:00:00 2001 From: Hossein Dehghan Date: Fri, 9 Oct 2026 15:18:44 +0000 Subject: [PATCH 2/6] fix(autoscaler): make lock release ownership-safe Carry the lock owner identity in DistributedLockGuard and use Cassandra LWT conditional deletes in drop and retry paths. Prevent an expired guard from removing a lock that another node acquired. Signed-off-by: Hossein Dehghan --- .../server/src/cassandra/cassandra_service.rs | 33 +++++++++++++++ .../server/src/cassandra/distributed_lock.rs | 41 ++++++++++++------- 2 files changed, 60 insertions(+), 14 deletions(-) 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..6fb8271a91 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 @@ -636,6 +636,33 @@ impl CassandraServiceManager { Ok(()) } + /// Delete a lock only when it is still owned by `node_id`. + /// Returns false when the lock has expired or another node owns it. + #[tracing::instrument(skip(self))] + pub async fn delete_owned_lock(&self, lock_name: &str, node_id: &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)) + .await?; + let rows = result.into_rows_result()?; + if rows.column_specs().get_by_name("node_id").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?; @@ -944,6 +971,12 @@ mod tests { assert!(put_result.is_ok()); assert!(put_result.unwrap()); + assert!(!manager + .delete_owned_lock(&lock.lock_name, "another-node") + .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()); 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..de65656bda 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 @@ -97,6 +97,7 @@ impl DistributedLockManager { pub struct DistributedLockGuard { lock_name: String, + node_id: String, cassandra_service: Arc, released: bool, } @@ -112,7 +113,7 @@ impl DistributedLockGuard { .put_lock( &DistributedLock { lock_name: lock_name.clone(), - node_id, + node_id: node_id.clone(), acquired_at: Utc::now(), }, lock_duration_seconds, @@ -122,6 +123,7 @@ impl DistributedLockGuard { if lock_acquired { Ok(Some(Self { lock_name, + node_id, cassandra_service, released: false, })) @@ -136,34 +138,45 @@ 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 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).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_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) + .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 + ), } }); } From 04f517117a1c3db45703fb536ac1433758c924fa Mon Sep 17 00:00:00 2001 From: Hossein Dehghan Date: Fri, 9 Oct 2026 15:21:27 +0000 Subject: [PATCH 3/6] style(autoscaler): normalize lock statement formatting Replace the editor-generated whitespace-only diff with the repository-formatted statements file. Signed-off-by: Hossein Dehghan --- .../crates/server/src/cassandra/statements.rs | 223 +----------------- 1 file changed, 1 insertion(+), 222 deletions(-) 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 b6b36b0161..3777094629 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 @@ -1,359 +1,138 @@ /* - * SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - - // 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 = ?;", - keyspace - ) - - } - - +} 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 = ?;", - keyspace - ) - } - - // healthy_nodes table - pub(crate) fn get_select_all_nodes_stmt(keyspace: &str) -> String { - format!( - "SELECT node_id, last_updated_at FROM {}.healthy_nodes;", - keyspace - ) - } - - pub(crate) fn get_delete_node_stmt(keyspace: &str) -> String { - format!("DELETE FROM {}.healthy_nodes WHERE node_id = ?;", keyspace) - } - - // recently_invoked_functions Table - // account_id is a regular column, not part of PK - pub(crate) fn get_select_recently_invoked_functions_in_token_range_stmt(keyspace: &str) -> String { - format!( - "SELECT function_id, function_version_id, account_id \ - FROM {}.recently_invoked_functions \ - WHERE token(function_id, function_version_id) >= ? AND token(function_id, function_version_id) <= ?;", - keyspace - ) - } - - pub(crate) fn get_health_check_query_stmt(keyspace: &str) -> String { - format!("SELECT now() from {}.healthy_nodes LIMIT 1;", keyspace) - } - - // Inserts to the locks table must be done with a row TTL. - // Bind order: (lock_name, node_id, acquired_at, 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 ?", - keyspace, - ) - } - - // LWT conditional update — only refreshes TTL if node_id still matches this node. - // 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) - 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 = ?", - keyspace, - ) - } - - // Inserts to the healthy_nodes table with a configurable row TTL (from CassandraSettings.node_health_ttl_seconds, default 180s). - // The node is pruned automatically after TTL expires if it stops reporting healthy. - pub(crate) fn get_stmt_insert_to_nodes(keyspace: &str, ttl_seconds: i32) -> String { - format!( - "INSERT INTO {}.healthy_nodes (node_id, last_updated_at) VALUES (?, ?) USING TTL {}", - keyspace, ttl_seconds - ) - } - - // Inserts to the recently_invoked_functions table with a configurable row TTL (from CassandraSettings.recently_invoked_ttl_seconds, default 1800s). - pub(crate) fn get_stmt_insert_to_recently_invoked_functions( - keyspace: &str, - ttl_seconds: i32, - ) -> String { - format!( - "INSERT INTO {}.recently_invoked_functions (function_id, function_version_id, account_id, last_updated_at) VALUES (?, ?, ?, ?) USING TTL {}", - keyspace, - ttl_seconds - ) - } - - #[cfg(test)] - mod tests { - use super::{ - 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, - }; - - #[test] - 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 = ?")); - } - - #[test] - fn delete_owned_lock_checks_node_id() { - assert_eq!( - get_delete_owned_lock_stmt("test_keyspace"), - "DELETE FROM test_keyspace.locks WHERE lock_name = ? IF node_id = ?;" - ); - } - - #[test] - fn active_function_statements_preserve_table_and_ttl() { - let select = get_select_recently_invoked_functions_in_token_range_stmt("test_keyspace"); - let insert = get_stmt_insert_to_recently_invoked_functions("test_keyspace", 1800); - - assert!(select.contains("FROM test_keyspace.recently_invoked_functions")); - assert!(insert.contains("INTO test_keyspace.recently_invoked_functions")); - assert!(insert.ends_with("USING TTL 1800")); - } - } - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - From e821a0c9b4c6d8bb047095cac0747a02310fd61e Mon Sep 17 00:00:00 2001 From: Hossein Dehghan Date: Fri, 9 Oct 2026 15:58:30 +0000 Subject: [PATCH 4/6] fix(autoscaler): release locks by acquisition generation Signed-off-by: Hossein Dehghan --- .../server/src/cassandra/cassandra_service.rs | 79 +++++++++++++++---- .../server/src/cassandra/distributed_lock.rs | 26 ++++-- .../crates/server/src/cassandra/statements.rs | 19 ++--- 3 files changed, 94 insertions(+), 30 deletions(-) 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 6fb8271a91..1d3019bcc8 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; @@ -582,6 +582,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 +599,25 @@ 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_name, + node_id, + acquired_at, + ), ) .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,10 +653,15 @@ impl CassandraServiceManager { Ok(()) } - /// Delete a lock only when it is still owned by `node_id`. - /// Returns false when the lock has expired or another node owns it. + /// 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) -> Result { + pub async fn delete_owned_lock( + &self, + lock_name: &str, + node_id: &str, + acquired_at: DateTime, + ) -> Result { let session = self.get_session().await?; with_cassandra_timing("delete_owned_lock", || async { let mut prepared = session @@ -649,10 +671,12 @@ impl CassandraServiceManager { prepared.set_serial_consistency(Some(SerialConsistency::Serial)); prepared.set_is_idempotent(false); let result = session - .execute_unpaged(&prepared, (lock_name, node_id)) + .execute_unpaged(&prepared, (lock_name, node_id, acquired_at)) .await?; let rows = result.into_rows_result()?; - if rows.column_specs().get_by_name("node_id").is_some() { + 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() { @@ -966,13 +990,34 @@ 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(), + }; + 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(),) + .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") + .delete_owned_lock( + &lock.lock_name, + "another-node", + reacquired_lock.acquired_at.clone(), + ) .await .unwrap()); assert!(manager.get_lock(&lock.lock_name).await.unwrap().is_some()); @@ -1018,7 +1063,11 @@ 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(), 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 de65656bda..b1b1dfb666 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,7 +19,7 @@ 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; @@ -61,13 +61,14 @@ impl DistributedLockManager { lock_name: String, lock_duration_seconds: i32, ) -> Result { + let acquired_at = Utc::now(); 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_duration_seconds, ) @@ -98,6 +99,7 @@ impl DistributedLockManager { pub struct DistributedLockGuard { lock_name: String, node_id: String, + acquired_at: DateTime, cassandra_service: Arc, released: bool, } @@ -109,12 +111,13 @@ impl DistributedLockGuard { cassandra_service: Arc, lock_duration_seconds: i32, ) -> Result> { + let acquired_at = Utc::now(); let lock_acquired = cassandra_service .put_lock( &DistributedLock { lock_name: lock_name.clone(), node_id: node_id.clone(), - acquired_at: Utc::now(), + acquired_at: acquired_at.clone(), }, lock_duration_seconds, ) @@ -124,6 +127,7 @@ impl DistributedLockGuard { Ok(Some(Self { lock_name, node_id, + acquired_at, cassandra_service, released: false, })) @@ -139,17 +143,22 @@ impl Drop for DistributedLockGuard { 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 cassandra_service = self.cassandra_service.clone(); tokio::spawn(async move { - match cassandra_service.delete_owned_lock(&lock_name, &node_id).await { + match cassandra_service + .delete_owned_lock(&lock_name, &node_id, acquired_at.clone()) + .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 + lock_name, + node_id ); } Err(e) => { @@ -158,11 +167,16 @@ impl Drop for DistributedLockGuard { // 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_cassandra_service = cassandra_service.clone(); tokio::spawn(async move { tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; match retry_cassandra_service - .delete_owned_lock(&retry_lock_name, &retry_node_id) + .delete_owned_lock( + &retry_lock_name, + &retry_node_id, + retry_acquired_at, + ) .await { Ok(true) => tracing::info!( 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 3777094629..c81418d0e5 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 @@ -29,7 +29,7 @@ pub(crate) fn get_delete_locks_stmt(keyspace: &str) -> String { pub(crate) fn get_delete_owned_lock_stmt(keyspace: &str) -> String { format!( - "DELETE FROM {}.locks WHERE lock_name = ? IF node_id = ?;", + "DELETE FROM {}.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ?;", keyspace ) } @@ -70,14 +70,15 @@ pub(crate) fn get_stmt_insert_to_locks(keyspace: &str) -> String { ) } -// 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_name, node_id, acquired_at) 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 = ? WHERE lock_name = ? IF node_id = ? AND acquired_at = ?", keyspace, ) } @@ -115,14 +116,14 @@ mod tests { 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("IF node_id = ? AND acquired_at = ?")); } #[test] - fn delete_owned_lock_checks_node_id() { + 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 = ?;" + "DELETE FROM test_keyspace.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ?;" ); } From 7e310bc855bcbbe260c7c490cfacabba723414c2 Mon Sep 17 00:00:00 2001 From: hossiendehghan989 <1.86807511e+08+hossiendehghan989@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:29:38 +0000 Subject: [PATCH 5/6] fix(autoscaler): guard locks with per-acquisition tokens --- .../nvcf_autoscaler/03_init_tables.up.sql | 3 ++- .../nvcf_autoscaler/05_add_lock_token.up.sql | 2 ++ .../server/src/cassandra/cassandra_service.rs | 17 +++++++++++++++-- .../server/src/cassandra/distributed_lock.rs | 12 +++++++++++- .../crates/server/src/cassandra/statements.rs | 18 +++++++++--------- .../crates/server/src/models/mod.rs | 2 ++ .../cassandra/schema/0001_initial_schema.cql | 3 ++- 7 files changed, 43 insertions(+), 14 deletions(-) create mode 100644 migrations/cassandra/keyspaces/nvcf_autoscaler/05_add_lock_token.up.sql diff --git a/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql b/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql index e38c166b52..2d8be06a55 100644 --- a/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql +++ b/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql @@ -62,7 +62,8 @@ CREATE TABLE IF NOT EXISTS nvcf_autoscaler.running_functions_without_invocations 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 = 3600 AND compaction = {'class': 'UnifiedCompactionStrategy', 'scaling_parameters': 'T4', 'target_sstable_size': '50MiB', 'base_shard_count': '4', 'expired_sstable_check_frequency_seconds': '300'} AND read_repair = 'NONE'; 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 1d3019bcc8..2b213d5110 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 @@ -524,6 +524,7 @@ impl CassandraServiceManager { &lock.lock_name, &lock.node_id, lock.acquired_at, + &lock.lock_token, ttl_seconds, ), ) @@ -606,9 +607,11 @@ impl CassandraServiceManager { ttl_seconds, node_id, acquired_at.clone(), + &lock.lock_token, lock_name, node_id, acquired_at, + &lock.lock_token, ), ) .await?; @@ -661,6 +664,7 @@ impl CassandraServiceManager { 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 { @@ -671,7 +675,7 @@ impl CassandraServiceManager { prepared.set_serial_consistency(Some(SerialConsistency::Serial)); prepared.set_is_idempotent(false); let result = session - .execute_unpaged(&prepared, (lock_name, node_id, acquired_at)) + .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() @@ -929,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()), } } @@ -995,6 +1000,7 @@ mod tests { 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()); @@ -1003,7 +1009,12 @@ mod tests { // 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(),) + .delete_owned_lock( + &lock.lock_name, + &lock.node_id, + lock.acquired_at.clone(), + &lock.lock_token, + ) .await .unwrap()); let current_lock = manager.get_lock(&lock.lock_name).await.unwrap().unwrap(); @@ -1017,6 +1028,7 @@ mod tests { &lock.lock_name, "another-node", reacquired_lock.acquired_at.clone(), + &reacquired_lock.lock_token, ) .await .unwrap()); @@ -1073,6 +1085,7 @@ mod tests { 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 b1b1dfb666..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 @@ -22,6 +22,7 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use std::sync::Arc; use tracing; +use uuid::Uuid; pub struct DistributedLockManager { cassandra_service: Arc, @@ -62,6 +63,7 @@ impl DistributedLockManager { 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( @@ -69,6 +71,7 @@ impl DistributedLockManager { lock_name: lock_name.clone(), node_id: self.node_id.clone(), acquired_at: acquired_at.clone(), + lock_token: Some(lock_token.clone()), }, lock_duration_seconds, ) @@ -100,6 +103,7 @@ pub struct DistributedLockGuard { lock_name: String, node_id: String, acquired_at: DateTime, + lock_token: String, cassandra_service: Arc, released: bool, } @@ -112,12 +116,14 @@ impl DistributedLockGuard { 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: node_id.clone(), acquired_at: acquired_at.clone(), + lock_token: Some(lock_token.clone()), }, lock_duration_seconds, ) @@ -128,6 +134,7 @@ impl DistributedLockGuard { lock_name, node_id, acquired_at, + lock_token, cassandra_service, released: false, })) @@ -144,11 +151,12 @@ impl Drop for DistributedLockGuard { 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_owned_lock(&lock_name, &node_id, acquired_at.clone()) + .delete_owned_lock(&lock_name, &node_id, acquired_at.clone(), &lock_token) .await { Ok(true) => { @@ -168,6 +176,7 @@ impl Drop for DistributedLockGuard { 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; @@ -176,6 +185,7 @@ impl Drop for DistributedLockGuard { &retry_lock_name, &retry_node_id, retry_acquired_at, + &retry_lock_token, ) .await { 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 c81418d0e5..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 ) } @@ -29,7 +29,7 @@ pub(crate) fn get_delete_locks_stmt(keyspace: &str) -> String { pub(crate) fn get_delete_owned_lock_stmt(keyspace: &str) -> String { format!( - "DELETE FROM {}.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ?;", + "DELETE FROM {}.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?;", keyspace ) } @@ -62,10 +62,10 @@ 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, ) } @@ -75,10 +75,10 @@ pub(crate) fn get_stmt_insert_to_locks(keyspace: &str) -> String { // 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_name, node_id, acquired_at) +// 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 = ? AND acquired_at = ?", + "UPDATE {}.locks USING TTL ? SET node_id = ?, acquired_at = ?, lock_token = ? WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?", keyspace, ) } @@ -115,15 +115,15 @@ 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 = ? AND acquired_at = ?")); + 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 = ?;" + "DELETE FROM test_keyspace.locks WHERE lock_name = ? IF node_id = ? AND acquired_at = ? AND lock_token = ?;" ); } 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 ( From c7dc9412f94f7b4963b0dd0b0c632e914e1da660 Mon Sep 17 00:00:00 2001 From: hossiendehghan989 <1.86807511e+08+hossiendehghan989@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:49:54 +0000 Subject: [PATCH 6/6] fix(autoscaler): address lock review findings Signed-off-by: hossiendehghan989 <1.86807511e+08+hossiendehghan989@users.noreply.github.com> --- .../keyspaces/nvcf_autoscaler/03_init_tables.up.sql | 3 +-- .../crates/server/src/cassandra/cassandra_service.rs | 9 +++++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql b/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql index 2d8be06a55..e38c166b52 100644 --- a/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql +++ b/migrations/cassandra/keyspaces/nvcf_autoscaler/03_init_tables.up.sql @@ -62,8 +62,7 @@ CREATE TABLE IF NOT EXISTS nvcf_autoscaler.running_functions_without_invocations CREATE TABLE IF NOT EXISTS nvcf_autoscaler.locks ( lock_name TEXT PRIMARY KEY, node_id TEXT, - acquired_at TIMESTAMP, - lock_token TEXT + acquired_at TIMESTAMP ) WITH default_time_to_live = 3600 AND compaction = {'class': 'UnifiedCompactionStrategy', 'scaling_parameters': 'T4', 'target_sstable_size': '50MiB', 'base_shard_count': '4', 'expired_sstable_check_frequency_seconds': '300'} AND read_repair = 'NONE'; 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 2b213d5110..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 @@ -1013,7 +1013,9 @@ mod tests { &lock.lock_name, &lock.node_id, lock.acquired_at.clone(), - &lock.lock_token, + lock.lock_token + .as_deref() + .expect("test lock has a token"), ) .await .unwrap()); @@ -1028,7 +1030,10 @@ mod tests { &lock.lock_name, "another-node", reacquired_lock.acquired_at.clone(), - &reacquired_lock.lock_token, + reacquired_lock + .lock_token + .as_deref() + .expect("reacquired test lock has a token"), ) .await .unwrap());