From b1cd4818ddd912a00e0129fd7b23f6b27d7ab63e Mon Sep 17 00:00:00 2001 From: Pawel Iwan Date: Fri, 6 Feb 2026 18:21:46 +0100 Subject: [PATCH 1/5] fix: locks removed when clean up registry --- Cargo.lock | 2 +- Cargo.toml | 8 ++++++-- .../lua_scripts/clean_up_coin_registry.lua | 20 ++++++++++++++++--- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4938f39..942960d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4337,7 +4337,7 @@ dependencies = [ [[package]] name = "iota-gas-station" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "arc-swap", diff --git a/Cargo.toml b/Cargo.toml index 9a43054..7345af9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "iota-gas-station" -version = "0.4.0" +version = "0.5.0" edition = "2021" authors = ["Mysten Labs ", "IOTA Stiftung"] license = "Apache-2.0" @@ -14,7 +14,11 @@ iota-metrics = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", pa iota-config = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", package = "iota-config" } iota-json-rpc-types = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", package = "iota-json-rpc-types" } iota-sdk = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", package = "iota-sdk" } -iota-sdk-types = { package = "iota-sdk-types", git = "https://github.com/iotaledger/iota-rust-sdk.git", rev = "05608b7e4a5b96d85f84e1970a517f6f174beb8b", features = ["hash", "serde", "schemars"] } +iota-sdk-types = { package = "iota-sdk-types", git = "https://github.com/iotaledger/iota-rust-sdk.git", rev = "05608b7e4a5b96d85f84e1970a517f6f174beb8b", features = [ + "hash", + "serde", + "schemars", +] } iota-types = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", package = "iota-types" } telemetry-subscribers = { git = "https://github.com/iotaledger/iota", tag = "v1.15.0", package = "telemetry-subscribers" } diff --git a/src/storage/redis/lua_scripts/clean_up_coin_registry.lua b/src/storage/redis/lua_scripts/clean_up_coin_registry.lua index 9c54dc1..769e781 100644 --- a/src/storage/redis/lua_scripts/clean_up_coin_registry.lua +++ b/src/storage/redis/lua_scripts/clean_up_coin_registry.lua @@ -2,12 +2,16 @@ -- SPDX-License-Identifier: Apache-2.0 -- This script is used to clean up all data associated with a sponsor's coin registry. --- It deletes all keys with the namespace prefix. +-- It deletes all keys with the namespace prefix, EXCEPT for lock keys which must be preserved. -- The first argument is the namespace. local namespace = ARGV[1] local pattern = namespace .. ':*' +-- Lock keys that should NOT be deleted during cleanup +local t_maintenance_lock = namespace .. ':maintenance_lock' +local t_init_lock = namespace .. ':init_lock' + local cursor = "0" local deleted_count = 0 @@ -17,8 +21,18 @@ repeat local keys = result[2] if #keys > 0 then - redis.call('DEL', unpack(keys)) - deleted_count = deleted_count + #keys + -- Filter out lock keys that should be preserved + local keys_to_delete = {} + for i, key in ipairs(keys) do + if key ~= t_maintenance_lock and key ~= t_init_lock then + table.insert(keys_to_delete, key) + end + end + + if #keys_to_delete > 0 then + redis.call('DEL', unpack(keys_to_delete)) + deleted_count = deleted_count + #keys_to_delete + end end until cursor == "0" From 1c7144de46fd7c2eceb9044b4345accc92433b5e Mon Sep 17 00:00:00 2001 From: Pawel Iwan Date: Fri, 6 Feb 2026 18:22:45 +0100 Subject: [PATCH 2/5] fix: unnecessary consistency check every single time the scan is performed --- src/gas_station_initializer.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/gas_station_initializer.rs b/src/gas_station_initializer.rs index c2650a4..985ebb7 100644 --- a/src/gas_station_initializer.rs +++ b/src/gas_station_initializer.rs @@ -208,11 +208,14 @@ impl GasStationInitializer { force_full_rescan: bool, ignore_locks: bool, ) -> Self { + let sponsor_address = signer.get_address(); + Self::perform_consistency_check(&iota_client, &storage, sponsor_address).await; + if force_full_rescan { if storage .acquire_maintenance_lock(MAX_MAINTENANCE_DURATION_SEC) .await - .unwrap() + .expect("Failed to acquire maintenance lock for full rescan") { info!("Acquired maintenance lock for full rescan"); } else { @@ -365,8 +368,6 @@ impl GasStationInitializer { } } - Self::perform_consistency_check(&iota_client, storage, sponsor_address).await; - let start = Instant::now(); let balance_threshold = if matches!(mode, RunMode::Init) { info!("The pool has never been initialized. Initializing it for the first time"); @@ -468,7 +469,7 @@ impl GasStationInitializer { storage: &Arc, sponsor_address: IotaAddress, ) { - info!("Performing quick consistency check before rescan"); + info!("Performing quick consistency check"); let registry_coin_count = match storage.get_available_coin_count().await { Ok(count) => count as u64, Err(e) => { From f252b02b8202ea9f2c8d93f891637249483e105d Mon Sep 17 00:00:00 2001 From: Pawel Iwan Date: Fri, 6 Feb 2026 18:23:27 +0100 Subject: [PATCH 3/5] fix: cold params key fix --- src/config/cold_params.rs | 4 ++ src/storage/redis/mod.rs | 99 +++++++++++++++++++++++++++++++++++---- 2 files changed, 93 insertions(+), 10 deletions(-) diff --git a/src/config/cold_params.rs b/src/config/cold_params.rs index a6906a2..f881a8d 100644 --- a/src/config/cold_params.rs +++ b/src/config/cold_params.rs @@ -74,6 +74,10 @@ impl ColdParams { ))?; let changes = self.changes_details(&old_cold_params); + // TODO + // this is wrong, because if there are changes the it should be returned we might want to return it to the user + // and user might want to revert the changes and restart the gas station without the rescan + // So set_data should only happens when, changes are detected, but the flag --allow-reinit is set storage .set_data( key, diff --git a/src/storage/redis/mod.rs b/src/storage/redis/mod.rs index f2e0ea6..5678d7c 100644 --- a/src/storage/redis/mod.rs +++ b/src/storage/redis/mod.rs @@ -84,15 +84,15 @@ fn generate_namespace(namespace_prefix: &str) -> String { impl SetGetStorage for RedisStorage { async fn set_data(&self, key: &str, value: Vec) -> anyhow::Result<()> { let mut conn = self.conn_manager.clone(); - let full_key = format!("{}:{}", self.namespace, key); - conn.set::<_, _, ()>(full_key, value).await?; + // Key is expected to be absolute (already includes namespace if needed) + conn.set::<_, _, ()>(key, value).await?; Ok(()) } async fn get_data(&self, key: &str) -> anyhow::Result>> { let mut conn = self.conn_manager.clone(); - let full_key = format!("{}:{}", self.namespace, key); - let value: Option> = conn.get(full_key).await?; + // Key is expected to be absolute (already includes namespace if needed) + let value: Option> = conn.get(key).await?; Ok(value) } } @@ -419,11 +419,13 @@ mod tests { let test_struct = TestStruct { value: "test_value".to_string(), }; + // Use absolute key with namespace + let key = format!("{}:test_key", storage.namespace); storage - .set_data("test_key", serde_json::to_vec(&test_struct).unwrap()) + .set_data(&key, serde_json::to_vec(&test_struct).unwrap()) .await .unwrap(); - let value: Option> = storage.get_data("test_key").await.unwrap(); + let value: Option> = storage.get_data(&key).await.unwrap(); let value: TestStruct = serde_json::from_slice(&value.unwrap()).unwrap(); assert_eq!(value, test_struct); } @@ -513,12 +515,13 @@ mod tests { let total_balance = storage.get_available_coin_total_balance().await; assert_eq!(total_balance, 300); - // Set some data + // Set some data with absolute key + let test_key = format!("{}:test_key", storage.namespace); storage - .set_data("test_key", b"test_value".to_vec()) + .set_data(&test_key, b"test_value".to_vec()) .await .unwrap(); - let value = storage.get_data("test_key").await.unwrap(); + let value = storage.get_data(&test_key).await.unwrap(); assert!(value.is_some()); // Clean up the coin registry @@ -530,10 +533,86 @@ mod tests { assert_eq!(total_balance, 0); // Verify the test_key is also cleaned up - let value = storage.get_data("test_key").await.unwrap(); + let value = storage.get_data(&test_key).await.unwrap(); assert!(value.is_none()); } + #[tokio::test] + async fn test_clean_up_preserves_locks() { + let storage = setup_storage().await; + + // Add some coins and data + storage + .add_new_coins(vec![ + GasCoin { + balance: 100, + object_ref: random_object_ref(), + }, + GasCoin { + balance: 200, + object_ref: random_object_ref(), + }, + ]) + .await + .unwrap(); + let test_key = format!("{}:test_key", storage.namespace); + storage + .set_data(&test_key, b"test_value".to_vec()) + .await + .unwrap(); + + // Acquire maintenance lock BEFORE cleanup + let lock_acquired = storage + .acquire_maintenance_lock(300) // 300 seconds + .await + .unwrap(); + assert!(lock_acquired, "Should acquire maintenance lock"); + + // Verify we're in maintenance mode + assert!( + storage.is_maintenance_mode().await.unwrap(), + "Should be in maintenance mode" + ); + + // Clean up the coin registry (this should NOT delete the maintenance lock) + storage.clean_up_coin_registry().await.unwrap(); + + // Verify maintenance lock is STILL held after cleanup + assert!( + storage.is_maintenance_mode().await.unwrap(), + "Maintenance lock should be preserved after cleanup" + ); + + // Verify we cannot acquire the lock again (it's still held) + let lock_acquired_again = storage.acquire_maintenance_lock(300).await.unwrap(); + assert!( + !lock_acquired_again, + "Should not be able to acquire lock again while it's held" + ); + + // Verify other data was cleaned up + let (coin_count, total_balance) = storage.init_coin_stats_at_startup().await.unwrap(); + assert_eq!(coin_count, 0, "Coins should be cleaned up"); + assert_eq!(total_balance, 0, "Balance should be zero"); + + let value = storage.get_data(&test_key).await.unwrap(); + assert!(value.is_none(), "Test data should be cleaned up"); + + // Release the lock + storage.release_maintenance_lock().await.unwrap(); + assert!( + !storage.is_maintenance_mode().await.unwrap(), + "Should not be in maintenance mode after release" + ); + + // Now we should be able to acquire it again + let lock_acquired_final = storage.acquire_maintenance_lock(300).await.unwrap(); + assert!( + lock_acquired_final, + "Should be able to acquire lock after release" + ); + } + async fn setup_storage() -> RedisStorage { let sponsor = IotaAddress::ZERO; let namespace_prefix = "test"; From c9eaf17000beab9bf2a0f1083767ac37d37fe44a Mon Sep 17 00:00:00 2001 From: Pawel Iwan Date: Mon, 9 Feb 2026 11:27:22 +0100 Subject: [PATCH 4/5] fix: cold param check should automatically save cold params --- src/command.rs | 10 +++++++++- src/config/cold_params.rs | 25 +++++++++---------------- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/src/command.rs b/src/command.rs index 681eab6..c76c0e0 100644 --- a/src/command.rs +++ b/src/command.rs @@ -114,7 +114,7 @@ impl Command { let rescan_trigger_receiver = rescan_config.create_receiver(); let cold_params_changes = cold_params - .check_if_changed(&storage, &format!("{namespace_prefix}:cold_params")) + .check_if_changed(&storage, &get_cold_params_key(&namespace_prefix)) .await .expect("failed to check cold params changes"); @@ -149,6 +149,10 @@ impl Command { } else { None }; + cold_params + .save_to_storage(&storage, &get_cold_params_key(&namespace_prefix)) + .await + .expect("failed to save cold params to storage"); let core_metrics = GasStationCoreMetrics::new(&prometheus_registry); let stats_storage = connect_stats_storage(&gas_station_config, &namespace_prefix).await; let stats_tracker = StatsTracker::new(Arc::new(stats_storage)); @@ -186,3 +190,7 @@ impl Command { server.handle.await.unwrap(); } } + +fn get_cold_params_key(namespace_prefix: &str) -> String { + format!("{namespace_prefix}:cold_params") +} diff --git a/src/config/cold_params.rs b/src/config/cold_params.rs index f881a8d..b368343 100644 --- a/src/config/cold_params.rs +++ b/src/config/cold_params.rs @@ -56,28 +56,21 @@ impl ColdParams { .await .context("unable to get cold params from storage")?; - // If there are no cold params, it means this is the first run. - // In this case, there are no changes, so we only need to save the current cold params. if maybe_cold_params.is_none() { - storage - .set_data( - key, - serde_json::to_vec(&self) - .context("unable to serialize cold params and save to storage")?, - ) - .await?; return Ok(vec![]); } - let old_cold_params = serde_json::from_slice(&maybe_cold_params.unwrap()).context( format!("unable to deserialize cold params. The entry with the key {key} is not a valid cold params structure", ))?; - let changes = self.changes_details(&old_cold_params); - // TODO - // this is wrong, because if there are changes the it should be returned we might want to return it to the user - // and user might want to revert the changes and restart the gas station without the rescan - // So set_data should only happens when, changes are detected, but the flag --allow-reinit is set + Ok(changes) + } + + pub async fn save_to_storage( + &self, + storage: &Arc, + key: &str, + ) -> anyhow::Result<()> { storage .set_data( key, @@ -85,7 +78,7 @@ impl ColdParams { .context("unable to serialize cold params and save to storage")?, ) .await?; - Ok(changes) + Ok(()) } } From d2fe79e15c8cdb0238c1707f76a34b6eea60283c Mon Sep 17 00:00:00 2001 From: Pawel Iwan Date: Tue, 10 Feb 2026 12:49:40 +0100 Subject: [PATCH 5/5] fix: methods description --- src/storage/mod.rs | 5 +++++ src/storage/redis/mod.rs | 4 ---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/storage/mod.rs b/src/storage/mod.rs index ad22e97..7d4a480 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -92,7 +92,12 @@ pub trait Storage: SetGetStorage + Sync + Send { #[async_trait::async_trait] pub trait SetGetStorage: Sync + Send { + /// Set data in the storage. + /// Key is expected to be absolute (already includes namespace if needed) async fn set_data(&self, key: &str, value: Vec) -> anyhow::Result<()>; + + /// Get data from the storage. + /// Key is expected to be absolute (already includes namespace if needed) async fn get_data(&self, key: &str) -> anyhow::Result>>; } diff --git a/src/storage/redis/mod.rs b/src/storage/redis/mod.rs index 5678d7c..e49a73c 100644 --- a/src/storage/redis/mod.rs +++ b/src/storage/redis/mod.rs @@ -84,14 +84,11 @@ fn generate_namespace(namespace_prefix: &str) -> String { impl SetGetStorage for RedisStorage { async fn set_data(&self, key: &str, value: Vec) -> anyhow::Result<()> { let mut conn = self.conn_manager.clone(); - // Key is expected to be absolute (already includes namespace if needed) conn.set::<_, _, ()>(key, value).await?; Ok(()) } - async fn get_data(&self, key: &str) -> anyhow::Result>> { let mut conn = self.conn_manager.clone(); - // Key is expected to be absolute (already includes namespace if needed) let value: Option> = conn.get(key).await?; Ok(value) } @@ -419,7 +416,6 @@ mod tests { let test_struct = TestStruct { value: "test_value".to_string(), }; - // Use absolute key with namespace let key = format!("{}:test_key", storage.namespace); storage .set_data(&key, serde_json::to_vec(&test_struct).unwrap())