From f4691f8cda5ccbbd695427caa62833dbdfe3b8c2 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 12 Mar 2026 12:52:11 +0200 Subject: [PATCH 01/21] do not recycle presignatures --- .../node/src/protocol/presignature.rs | 34 +++++++----- .../node/src/protocol/signature.rs | 23 +++----- chain-signatures/node/src/protocol/triple.rs | 43 ++++++++------- .../node/src/storage/protocol_storage.rs | 55 ------------------- 4 files changed, 51 insertions(+), 104 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index b915b5671..acca5b7cb 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -682,6 +682,7 @@ impl PresignatureSpawner { mut cfg: watch::Receiver, ongoing_gen_tx: watch::Sender, ) { + let mut last_active_warn: Option = None; let mut stockpile_interval = time::interval(Duration::from_millis(100)); let mut expiration_interval = tokio::time::interval(Duration::from_secs(1)); let mut posits = self.msg.subscribe_presignature_posit().await; @@ -720,21 +721,26 @@ impl PresignatureSpawner { self.ongoing_owned.remove(&id); let _ = ongoing_gen_tx.send(self.ongoing.len()); } - _ = stockpile_interval.tick(), if active.len() >= self.threshold => { - self.stockpile(&active, &protocol).await; - let _ = ongoing_gen_tx.send(self.ongoing.len()); - - crate::metrics::storage::NUM_PRESIGNATURES_MINE - - .set(self.len_mine().await as i64); - crate::metrics::storage::NUM_PRESIGNATURES_TOTAL - - .set(self.len_generated().await as i64); - crate::metrics::protocols::NUM_PRESIGNATURE_GENERATORS_TOTAL - - .set( - self.len_potential().await as i64 - self.len_generated().await as i64, + _ = stockpile_interval.tick() => { + if active.len() >= self.threshold { + last_active_warn = None; + self.stockpile(&active, &protocol).await; + let _ = ongoing_gen_tx.send(self.ongoing.len()); + + crate::metrics::storage::NUM_PRESIGNATURES_MINE + .set(self.len_mine().await as i64); + crate::metrics::storage::NUM_PRESIGNATURES_TOTAL + .set(self.len_generated().await as i64); + crate::metrics::protocols::NUM_PRESIGNATURE_GENERATORS_TOTAL + .set(self.len_potential().await as i64 - self.len_generated().await as i64); + } else if last_active_warn.map_or(true, |i: Instant| i.elapsed() > Duration::from_secs(60)) { + tracing::warn!( + ?active, + threshold = self.threshold, + "not enough active participants to generate presignatures" ); + last_active_warn = Some(Instant::now()); + } } Ok(()) = cfg.changed() => { protocol = cfg.borrow().protocol.clone(); diff --git a/chain-signatures/node/src/protocol/signature.rs b/chain-signatures/node/src/protocol/signature.rs index dfcac4e02..ed378ffc7 100644 --- a/chain-signatures/node/src/protocol/signature.rs +++ b/chain-signatures/node/src/protocol/signature.rs @@ -272,14 +272,16 @@ impl SignOrganizer { let (presignature_id, presignature, active) = if is_proposer { tracing::info!(?sign_id, round = ?state.round, "proposer waiting for presignature"); let active = active.iter().copied().collect::>(); - let mut recycle = Vec::new(); let remaining = state.budget.remaining(); let fetch = tokio::time::timeout(remaining, async { loop { if let Some(taken) = ctx.presignatures.take_mine(ctx.me).await { let participants = intersect_vec(&[&taken.artifact.participants, &active]); if participants.len() < ctx.threshold { - recycle.push(taken); + tracing::warn!( + ?sign_id, + "discarding presignature due to inactive participants" + ); continue; } @@ -290,13 +292,6 @@ impl SignOrganizer { }) .await; - let presignatures = ctx.presignatures.clone(); - tokio::spawn(async move { - for taken in recycle { - presignatures.recycle_mine(me, taken).await; - } - }); - let (taken, participants) = match fetch { Ok(value) => value, Err(_) => { @@ -622,9 +617,8 @@ impl SignPositor { if counter.enough_rejects(ctx.threshold) { tracing::warn!(?sign_id, ?round, ?from, "received enough REJECTs, reorganizing"); - if let Some(taken) = presignature { - tracing::warn!(?sign_id, "recycling presignature due to REJECTs"); - ctx.presignatures.recycle_mine(ctx.me, taken).await; + if let Some(_taken) = presignature { + tracing::warn!(?sign_id, "discarding presignature due to REJECTs"); } state.bump_round(); return SignPhase::Organizing(SignOrganizer); @@ -664,9 +658,8 @@ impl SignPositor { ?round, "proposer posit deadline reached, expiring round" ); - if let Some(taken) = presignature { - tracing::warn!(?sign_id, "recycling presignature due to proposer timeout"); - ctx.presignatures.recycle_mine(ctx.me, taken).await; + if let Some(_taken) = presignature { + tracing::warn!(?sign_id, "discarding presignature due to proposer timeout"); } } else { tracing::warn!(?sign_id, me=?ctx.me, ?proposer, "deliberator posit timeout waiting for Start, reorganizing"); diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index 1ee0ae81a..ac5e1ecbb 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -556,10 +556,6 @@ impl TripleSpawner { /// Stockpile triples if the amount of unspent triples is below the minimum /// and the maximum number of all ongoing generation protocols is below the maximum. async fn stockpile(&mut self, participants: &[Participant], cfg: &ProtocolConfig) { - if participants.len() < self.threshold { - return; - } - let not_enough_triples = { // Stopgap to prevent too many triples in the system. This should be around min_triple*nodes*2 // for good measure so that we have enough triples to do presig generation while also maintain @@ -591,6 +587,7 @@ impl TripleSpawner { let mut active = mesh_state.borrow().active().keys_vec(); let mut protocol = cfg.borrow().protocol.clone(); + let mut last_active_warn = None; loop { tokio::select! { @@ -619,22 +616,28 @@ impl TripleSpawner { self.ongoing_introduced.remove(&id); let _ = ongoing_gen_tx.send(self.ongoing.len()); } - _ = stockpile_interval.tick(), if active.len() >= self.threshold => { - self.stockpile(&active, &protocol).await; - let _ = ongoing_gen_tx.send(self.ongoing.len()); - - crate::metrics::storage::NUM_TRIPLES_MINE - - .set(self.len_mine().await as i64); - crate::metrics::storage::NUM_TRIPLES_TOTAL - - .set(self.triple_storage.len_generated().await as i64); - crate::metrics::protocols::NUM_TRIPLE_GENERATORS_INTRODUCED - - .set(self.len_introduced() as i64); - crate::metrics::protocols::NUM_TRIPLE_GENERATORS_TOTAL - - .set(self.len_ongoing() as i64); + _ = stockpile_interval.tick() => { + if active.len() >= self.threshold { + last_active_warn = None; + self.stockpile(&active, &protocol).await; + let _ = ongoing_gen_tx.send(self.ongoing.len()); + + crate::metrics::storage::NUM_TRIPLES_MINE + .set(self.len_mine().await as i64); + crate::metrics::storage::NUM_TRIPLES_TOTAL + .set(self.triple_storage.len_generated().await as i64); + crate::metrics::protocols::NUM_TRIPLE_GENERATORS_INTRODUCED + .set(self.len_introduced() as i64); + crate::metrics::protocols::NUM_TRIPLE_GENERATORS_TOTAL + .set(self.len_ongoing() as i64); + } else if last_active_warn.map_or(true, |i: Instant| i.elapsed() > Duration::from_secs(60)) { + tracing::warn!( + ?active, + threshold = self.threshold, + "not enough active participants to generate triples" + ); + last_active_warn = Some(Instant::now()); + } } Ok(()) = cfg.changed() => { protocol = cfg.borrow().protocol.clone(); diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index a2ce74feb..80befdf0e 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -620,61 +620,6 @@ impl ProtocolStorage { } } - /// Return a taken artifact back to the available pool. - pub async fn recycle_mine(&self, me: Participant, taken: ArtifactTaken) -> bool { - const SCRIPT: &str = r#" - local artifact_key = KEYS[1] - local mine_key = KEYS[2] - local artifact_id = ARGV[1] - local artifact = ARGV[2] - - -- Add back to artifact hash map - redis.call("HSET", artifact_key, artifact_id, artifact) - - -- Add back to mine set - redis.call("SADD", mine_key, artifact_id) - - return 1 - "#; - - let start = Instant::now(); - let (artifact, mut dropper) = taken.take(); - // We manually handle the return, so we don't want the dropper to unreserve it. - dropper.dropper.take(); - - let id = artifact.id(); - let Some(mut conn) = self.connect().await else { - tracing::warn!(id, "failed to return artifact: connection failed"); - return false; - }; - - let result: Result = redis::Script::new(SCRIPT) - .key(&self.artifact_key) - .key(owner_key(&self.owner_keys, me)) - .arg(id) - .arg(artifact) - .invoke_async(&mut conn) - .await; - - let elapsed = start.elapsed(); - crate::metrics::storage::REDIS_LATENCY - .with_label_values(&[A::METRIC_LABEL, "return_mine"]) - .observe(elapsed.as_millis() as f64); - - match result { - Ok(_) => { - self.reserved.write().await.remove(&id); - self.used.write().await.remove(&id); - tracing::info!(id, ?elapsed, "returned mine artifact"); - true - } - Err(err) => { - tracing::warn!(id, ?err, ?elapsed, "failed to return mine artifact"); - false - } - } - } - /// Check if an artifact is reserved. pub async fn contains_reserved(&self, id: A::Id) -> bool { self.reserved.read().await.contains(&id) From 59808b834a62ece63aafe548a4a0d11140bfe4c9 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 12 Mar 2026 13:10:57 +0200 Subject: [PATCH 02/21] add .vscode to gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 2d8ff61c0..3106d8414 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ .direnv .DS_Store .idea +.vscode tmp *.log From feb9b136faa8077d9b7b41df4119d99e07f288e9 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 12 Mar 2026 13:12:29 +0200 Subject: [PATCH 03/21] clippy --- chain-signatures/node/src/protocol/presignature.rs | 2 +- chain-signatures/node/src/protocol/triple.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index acca5b7cb..b906b80a5 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -733,7 +733,7 @@ impl PresignatureSpawner { .set(self.len_generated().await as i64); crate::metrics::protocols::NUM_PRESIGNATURE_GENERATORS_TOTAL .set(self.len_potential().await as i64 - self.len_generated().await as i64); - } else if last_active_warn.map_or(true, |i: Instant| i.elapsed() > Duration::from_secs(60)) { + } else if last_active_warn.is_none_or(|i: Instant| i.elapsed() > Duration::from_secs(60)) { tracing::warn!( ?active, threshold = self.threshold, diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index ac5e1ecbb..35b5ba970 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -630,7 +630,7 @@ impl TripleSpawner { .set(self.len_introduced() as i64); crate::metrics::protocols::NUM_TRIPLE_GENERATORS_TOTAL .set(self.len_ongoing() as i64); - } else if last_active_warn.map_or(true, |i: Instant| i.elapsed() > Duration::from_secs(60)) { + } else if last_active_warn.is_none_or(|i: Instant| i.elapsed() > Duration::from_secs(60)) { tracing::warn!( ?active, threshold = self.threshold, From 7a8815bcc3eab351e5a4f426c10825cb4fcddf3f Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Tue, 17 Mar 2026 17:15:47 +0200 Subject: [PATCH 04/21] extend comment --- chain-signatures/node/src/protocol/signature.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/chain-signatures/node/src/protocol/signature.rs b/chain-signatures/node/src/protocol/signature.rs index 19544fd6f..2d2f7765e 100644 --- a/chain-signatures/node/src/protocol/signature.rs +++ b/chain-signatures/node/src/protocol/signature.rs @@ -288,6 +288,9 @@ impl SignOrganizer { if participants.len() < ctx.threshold { tracing::warn!( ?sign_id, + id = taken.artifact.id, + ?holders, + ?active, "discarding presignature due to inactive participants" ); continue; From 95268db19e95b199857c615e3dfb72b7ce8b53cd Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 10:14:18 +0300 Subject: [PATCH 05/21] replace reserved and used with using and generating --- .../node/src/protocol/presignature.rs | 12 +- chain-signatures/node/src/protocol/triple.rs | 10 +- .../node/src/storage/protocol_storage.rs | 202 +++++++----------- integration-tests/benches/store.rs | 10 +- integration-tests/src/containers.rs | 3 +- integration-tests/src/mpc_fixture/builder.rs | 7 +- integration-tests/tests/cases/helpers.rs | 6 +- integration-tests/tests/cases/store.rs | 82 ++----- 8 files changed, 112 insertions(+), 220 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index 35d0777d9..6d291323b 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -375,8 +375,8 @@ impl PresignatureSpawner { self.ongoing.contains_key(&id) } - pub async fn contains_used(&self, id: PresignatureId) -> bool { - self.presignatures.contains_used(id).await + pub async fn contains_using(&self, id: PresignatureId) -> bool { + self.presignatures.contains_using(id).await } /// Returns the number of unspent presignatures available in the manager. @@ -430,7 +430,7 @@ impl PresignatureSpawner { // TODO: we can potentially wait for the triples to exist first to then be able to accept. // whereas we just blatantly reject here. The problem with waiting is that the other side // might expire their posit first. - self.triples.contains_reserved(id.pair_id).await + self.triples.contains_generating(id.pair_id).await || self.triples.contains(id.pair_id).await } { tracing::warn!( @@ -566,11 +566,7 @@ impl PresignatureSpawner { "starting protocol to generate a new presignature", ); - let Some(slot) = self.presignatures.reserve(id.id).await else { - return Err(InitializationError::BadParameters(format!( - "id collision: presignature_id={id:?}" - ))); - }; + let slot = self.presignatures.create_slot(id.id).await; let mut participants = participants.to_vec(); participants.sort(); diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index d6b4d9c0c..e60a377a4 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -388,8 +388,8 @@ impl TripleSpawner { self.ongoing.contains_key(&id) } - pub async fn contains_used(&self, id: TripleId) -> bool { - self.triple_storage.contains_used(id).await + pub async fn contains_using(&self, id: TripleId) -> bool { + self.triple_storage.contains_using(id).await } /// Returns the number of unspent triples assigned to this node. @@ -528,11 +528,7 @@ impl TripleSpawner { timeout: Duration, ) -> Result<(), InitializationError> { // Check if the `id` is already in the system. Error out and have the next cycle try again. - let Some(slot) = self.triple_storage.reserve(id).await else { - return Err(InitializationError::BadParameters(format!( - "id collision: pair_id={id}" - ))); - }; + let slot = self.triple_storage.create_slot(id).await; tracing::info!(?id, "starting protocol to generate a new triple"); let generator = TripleGenerator::new( diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index ed352ca62..4469fe047 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -6,7 +6,6 @@ use std::collections::HashSet; use std::sync::Arc; use std::{fmt, time::Instant}; use tokio::sync::RwLock; -use tokio::task::JoinHandle; use tracing; use super::{owner_key, STORAGE_VERSION}; @@ -68,7 +67,8 @@ pub trait ProtocolArtifact: fn set_holders(&mut self, holders: Vec); } -/// A pre-reserved slot for an artifact that will eventually be inserted. +/// A handle for inserting a generated artifact into storage. +/// Tracks the artifact ID in the `generating` set until insertion or drop. pub struct ArtifactSlot { id: A::Id, storage: ProtocolStorage, @@ -80,25 +80,15 @@ impl ArtifactSlot { self.stored = self.storage.insert(artifact, owner).await; self.stored } - - pub fn unreserve(&self) -> Option> { - if self.stored { - return None; - } - - let storage = self.storage.clone(); - let id = self.id; - let task = tokio::spawn(async move { - tracing::info!(id, "unreserving artifact"); - storage.unreserve(id).await; - }); - Some(task) - } } impl Drop for ArtifactSlot { fn drop(&mut self) { - self.unreserve(); + let storage = self.storage.clone(); + let id = self.id; + tokio::spawn(async move { + storage.generating.write().await.remove(&id); + }); } } @@ -117,7 +107,7 @@ impl Drop for ArtifactTakenDropper { if let Some(storage) = self.dropper.take() { let id = self.id; tokio::spawn(async move { - storage.unreserve(id).await; + storage.using.write().await.remove(&id); }); } } @@ -143,8 +133,8 @@ impl ArtifactTaken { pub struct ProtocolStorage { redis_pool: Pool, artifact_key: String, - used: Arc>>, - reserved: Arc>>, + generating: Arc>>, + using: Arc>>, owner_keys: String, account_id: AccountId, _phantom: std::marker::PhantomData, @@ -155,8 +145,8 @@ impl Clone for ProtocolStorage { Self { redis_pool: self.redis_pool.clone(), artifact_key: self.artifact_key.clone(), - used: self.used.clone(), - reserved: self.reserved.clone(), + generating: self.generating.clone(), + using: self.using.clone(), owner_keys: self.owner_keys.clone(), account_id: self.account_id.clone(), _phantom: std::marker::PhantomData, @@ -167,15 +157,15 @@ impl Clone for ProtocolStorage { impl ProtocolStorage { pub fn new(pool: &Pool, account_id: &AccountId, base_prefix: &str) -> Self { let artifact_key = format!("{base_prefix}:{STORAGE_VERSION}:{account_id}"); - let used = Arc::new(RwLock::new(HashSet::new())); - let reserved = Arc::new(RwLock::new(HashSet::new())); + let generating = Arc::new(RwLock::new(HashSet::new())); + let using = Arc::new(RwLock::new(HashSet::new())); let owner_keys = format!("{base_prefix}_owners:{STORAGE_VERSION}:{account_id}"); Self { redis_pool: pool.clone(), artifact_key, - used, - reserved, + generating, + using, owner_keys, account_id: account_id.clone(), _phantom: std::marker::PhantomData, @@ -210,51 +200,25 @@ impl ProtocolStorage { Ok(owned.into_iter().collect()) } - pub async fn reserve(&self, id: A::Id) -> Option> { - let used = self.used.read().await; - if used.contains(&id) { - return None; - } - if !self.reserved.write().await.insert(id) { - return None; - } - drop(used); - - let start = Instant::now(); - let Some(mut conn) = self.connect().await else { - self.reserved.write().await.remove(&id); - return None; - }; - - // Check directly whether the artifact is already stored in Redis. - let artifact_exists: Result = conn.hexists(&self.artifact_key, id).await; - let elapsed = start.elapsed(); - crate::metrics::storage::REDIS_LATENCY - .with_label_values(&[A::METRIC_LABEL, "reserve"]) - .observe(elapsed.as_millis() as f64); - - match artifact_exists { - Ok(true) => { - // artifact already stored, reserve cannot be done, remove reservation - self.reserved.write().await.remove(&id); - None - } - // artifact does not exist, reservation successful - Ok(false) => Some(ArtifactSlot { + /// Create a slot for generating an artifact with the given ID. + /// Tracks the ID in the `generating` set until the slot is inserted or dropped. + pub async fn create_slot(&self, id: A::Id) -> ArtifactSlot { + if !self.generating.write().await.insert(id) { + tracing::error!( id, - storage: self.clone(), - stored: false, - }), - Err(err) => { - self.reserved.write().await.remove(&id); - tracing::warn!(id, ?err, ?elapsed, "failed to reserve artifact"); - None - } + "creating slot for artifact that is already being generated" + ); + } + ArtifactSlot { + id, + storage: self.clone(), + stored: false, } } - async fn unreserve(&self, id: A::Id) -> bool { - self.reserved.write().await.remove(&id) + /// Check if an artifact is currently being generated. + pub async fn contains_generating(&self, id: A::Id) -> bool { + self.generating.read().await.contains(&id) } pub async fn remove_outdated( @@ -336,20 +300,17 @@ impl ProtocolStorage { match result { Ok((outdated, not_found)) => { - if !outdated.is_empty() { - tracing::info!(?outdated, ?elapsed, "removed outdated artifacts"); - // remove outdated entries from our in-memory reserved set - let mut reserved = self.reserved.write().await; - for id in outdated.iter() { - reserved.remove(id); - } - drop(reserved); - // remove outdated entries from our in-memory used set - let mut used = self.used.write().await; - for id in outdated.iter() { - used.remove(id); - } - } + // Filter out artifacts that are on this node but not in Redis: + // - `generating`: being generated, not yet persisted + // - `using`: taken from Redis, actively consumed by a protocol + let generating = self.generating.read().await; + let using = self.using.read().await; + let not_found: Vec<_> = not_found + .into_iter() + .filter(|id| !generating.contains(id) && !using.contains(id)) + .collect(); + drop(generating); + drop(using); Ok(RemoveOutdatedResult::new(outdated, not_found)) } Err(err) => { @@ -362,7 +323,8 @@ impl ProtocolStorage { /// Insert an artifact into storage under `owner`'s ownership set. /// Holders must be set on the artifact before calling this; they are /// persisted as a dedicated Redis set for later holder-tracking. - pub async fn insert(&self, artifact: A, owner: Participant) -> bool { + /// Private: callers must use `create_slot()` + `ArtifactSlot::insert()`. + async fn insert(&self, artifact: A, owner: Participant) -> bool { const SCRIPT: &str = r#" local artifact_key = KEYS[1] local owner_keys = KEYS[2] @@ -385,11 +347,6 @@ impl ProtocolStorage { let start = Instant::now(); let id = artifact.id(); - let used = self.used.read().await; - if used.contains(&id) { - tracing::warn!(id, "artifact already marked used"); - return false; - } let holders: Vec = artifact .holders() @@ -412,7 +369,6 @@ impl ProtocolStorage { .arg(holders.as_slice()) .invoke_async(&mut conn) .await; - drop(used); let elapsed = start.elapsed(); crate::metrics::storage::REDIS_LATENCY @@ -420,10 +376,7 @@ impl ProtocolStorage { .observe(elapsed.as_millis() as f64); match outcome { - Ok(()) => { - self.reserved.write().await.remove(&id); - true - } + Ok(()) => true, Err(err) => { tracing::warn!(id, ?err, ?elapsed, "failed to insert artifact"); false @@ -462,8 +415,14 @@ impl ProtocolStorage { } } - pub async fn contains_used(&self, id: A::Id) -> bool { - self.used.read().await.contains(&id) + /// Check if an artifact is currently being consumed by an active protocol. + pub async fn contains_using(&self, id: A::Id) -> bool { + self.using.read().await.contains(&id) + } + + /// Returns the set of artifact IDs currently being consumed by active protocols. + pub async fn using_ids(&self) -> HashSet { + self.using.read().await.clone() } pub async fn take(&self, id: A::Id, owner: Participant) -> Option> { @@ -491,14 +450,14 @@ impl ProtocolStorage { "#; let start = Instant::now(); - if !self.used.write().await.insert(id) { - tracing::warn!(id, "taking artifact that is already used"); + if !self.using.write().await.insert(id) { + tracing::warn!(id, "taking artifact that is already in use"); return None; } let Some(mut conn) = self.connect().await else { tracing::warn!(id, "failed to take artifact: connection failed"); - self.used.write().await.remove(&id); + self.using.write().await.remove(&id); return None; }; let result: Result<(A, Vec), _> = redis::Script::new(SCRIPT) @@ -521,7 +480,7 @@ impl ProtocolStorage { Some(ArtifactTaken::new(artifact, self.clone())) } Err(err) => { - self.used.write().await.remove(&id); + self.using.write().await.remove(&id); tracing::warn!(id, ?err, ?elapsed, "failed to take artifact"); None } @@ -602,8 +561,8 @@ impl ProtocolStorage { .with_label_values(&[A::METRIC_LABEL, "clear"]) .observe(elapsed.as_millis() as f64); - self.reserved.write().await.clear(); - self.used.write().await.clear(); + self.generating.write().await.clear(); + self.using.write().await.clear(); // if the outcome is None, it means the script failed or there was an error. outcome.is_some() @@ -659,10 +618,8 @@ impl ProtocolStorage { Ok(Some((mut artifact, holders))) => { let holders = holders.into_iter().map(Participant::from).collect(); artifact.set_holders(holders); - // mark reserved and used in-memory so that it won't be reserved or reused locally let id = artifact.id(); - self.reserved.write().await.insert(id); - self.used.write().await.insert(id); + self.using.write().await.insert(id); let taken = ArtifactTaken::new(artifact, self.clone()); tracing::debug!(id, ?elapsed, "took mine artifact"); Some(taken) @@ -675,16 +632,13 @@ impl ProtocolStorage { } } - /// Check if an artifact is reserved. - pub async fn contains_reserved(&self, id: A::Id) -> bool { - self.reserved.read().await.contains(&id) - } - pub fn artifact_key(&self) -> &str { &self.artifact_key } - /// Batch remove a peer from holders for a set of artifact IDs, and prune artifacts below threshold if owned by `me`. + /// Batch remove a peer from holders for a set of artifact IDs, and prune + /// artifacts that fall below the holder threshold. + /// Assumes the given IDs are owned by `me` for ownership-set cleanup. /// Returns (Vec, Vec) pub async fn remove_holder_and_prune( &self, @@ -697,7 +651,8 @@ impl ProtocolStorage { return Ok((vec![], vec![])); } - // Lua script expects: KEYS[1]=artifact_key, KEYS[2]=owner_key, ARGV[1]=peer, ARGV[2]=threshold, ARGV[3...]=ids + // Lua script expects: KEYS[1]=artifact_key, KEYS[2]=owner_key, + // ARGV[1]=peer, ARGV[2]=threshold, ARGV[3...]=ids const SCRIPT: &str = r#" local artifact_key = KEYS[1] local owner_key = KEYS[2] @@ -707,22 +662,19 @@ impl ProtocolStorage { local updated = {} for i = 3, #ARGV do local id = ARGV[i] - -- Error if 'me' does not own this artifact - if redis.call('SISMEMBER', owner_key, id) == 0 then - return redis.error_reply('OWNERSHIP_VIOLATION:' .. id) - end - -- Remove peer from holders set local holders_key = artifact_key .. ':holders:' .. id - redis.call('SREM', holders_key, peer) - local count = redis.call('SCARD', holders_key) - if count < threshold then - -- Prune: remove artifact, holders set, and owner set entry - redis.call('HDEL', artifact_key, id) - redis.call('DEL', holders_key) - redis.call('SREM', owner_key, id) - table.insert(removed, id) - else - table.insert(updated, id) + -- Skip if holders set doesn't exist (artifact already taken/consumed) + if redis.call('EXISTS', holders_key) == 1 then + redis.call('SREM', holders_key, peer) + local count = redis.call('SCARD', holders_key) + if count < threshold then + redis.call('HDEL', artifact_key, id) + redis.call('DEL', holders_key) + redis.call('SREM', owner_key, id) + table.insert(removed, id) + else + table.insert(updated, id) + end end end return {removed, updated} diff --git a/integration-tests/benches/store.rs b/integration-tests/benches/store.rs index 2f14f6519..47e057fdf 100644 --- a/integration-tests/benches/store.rs +++ b/integration-tests/benches/store.rs @@ -157,12 +157,7 @@ fn bench_load_keys(c: &mut Criterion) { rt.block_on(async { for i in 0..1000 { let t = dummy_pair(i); - env.triples - .reserve(t.id) - .await - .unwrap() - .insert(t, env.me) - .await; + env.triples.create_slot(t.id).await.insert(t, env.me).await; } }); }) @@ -174,9 +169,8 @@ fn bench_load_keys(c: &mut Criterion) { for i in 0..1000 { let p = dummy_presignature(i); env.presignatures - .reserve(p.id) + .create_slot(p.id) .await - .unwrap() .insert(p, env.me) .await; } diff --git a/integration-tests/src/containers.rs b/integration-tests/src/containers.rs index 240785b78..d67bfd448 100644 --- a/integration-tests/src/containers.rs +++ b/integration-tests/src/containers.rs @@ -429,9 +429,8 @@ impl Redis { storage .get(me) .unwrap() - .reserve(pair_id) + .create_slot(pair_id) .await - .unwrap() .insert(pair, *owner) .await; } diff --git a/integration-tests/src/mpc_fixture/builder.rs b/integration-tests/src/mpc_fixture/builder.rs index b14db0622..6724bc1e1 100644 --- a/integration-tests/src/mpc_fixture/builder.rs +++ b/integration-tests/src/mpc_fixture/builder.rs @@ -564,7 +564,7 @@ impl MpcFixtureNodeBuilder { if pair.holders.is_none() { pair.holders = Some(pair.triple0.public.participants.clone()); } - let mut slot = triple_storage.reserve(pair_id).await.unwrap(); + let mut slot = triple_storage.create_slot(pair_id).await; slot.insert(pair, owner).await; } } @@ -582,9 +582,8 @@ impl MpcFixtureNodeBuilder { presignature_share.holders = Some(presignature_share.participants.clone()); } let mut slot = presignature_storage - .reserve(presignature_share.id) - .await - .unwrap(); + .create_slot(presignature_share.id) + .await; slot.insert(presignature_share, owner).await; } } diff --git a/integration-tests/tests/cases/helpers.rs b/integration-tests/tests/cases/helpers.rs index aa6bb6ab2..736fef149 100644 --- a/integration-tests/tests/cases/helpers.rs +++ b/integration-tests/tests/cases/helpers.rs @@ -67,9 +67,8 @@ pub(crate) async fn insert_triples_for_owner( let holders = holders.to_vec(); for id in ids { triples - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_pair_with_holders(id, holders.clone()), owner) .await; } @@ -84,9 +83,8 @@ pub(crate) async fn insert_presignatures_for_owner( let holders = holders.to_vec(); for id in ids { presignatures - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_presignature_with_holders(id, holders.clone()), owner) .await; } diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 0e3a9cf5f..240f21b23 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -38,15 +38,13 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 0); triple_storage - .reserve(triple_id1) + .create_slot(triple_id1) .await - .unwrap() .insert(dummy_pair(triple_id1), node1) .await; triple_storage - .reserve(triple_id2) + .create_slot(triple_id2) .await - .unwrap() .insert(dummy_pair(triple_id2), node1) .await; @@ -59,7 +57,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_storage.len_by_owner(node0).await, 0); assert_eq!(triple_spawner.len_potential().await, 2); - // Take triple pairs and check that they are removed from the storage and added to used set + // Take triple pairs and check that they are removed from the storage and marked as using triple_storage.take(triple_id1, node1).await.unwrap(); triple_storage.take(triple_id2, node1).await.unwrap(); assert!(!triple_spawner.contains(triple_id1).await); @@ -69,35 +67,21 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_storage.len_generated().await, 0); assert_eq!(triple_spawner.len_mine().await, 0); assert_eq!(triple_spawner.len_potential().await, 0); - assert!(triple_storage.contains_used(triple_id1).await); - assert!(triple_storage.contains_used(triple_id2).await); - - // Attempt to re-reserve used triples and check that it cannot be reserved since it is used. - assert!(triple_storage.reserve(triple_id1).await.is_none()); - assert!(triple_storage.reserve(triple_id2).await.is_none()); - assert!(!triple_spawner.contains(triple_id1).await); - assert!(!triple_spawner.contains(triple_id2).await); + assert!(triple_storage.contains_using(triple_id1).await); + assert!(triple_storage.contains_using(triple_id2).await); let id3 = 3; let id4: u64 = 4; - // check that reserve and unreserve works: - let slot = triple_storage.reserve(id3).await.unwrap(); - if let Some(task) = slot.unreserve() { - task.await.unwrap(); - } - // Add mine triple and check that it is in the storage triple_storage - .reserve(id3) + .create_slot(id3) .await - .unwrap() .insert(dummy_pair(id3), node0) .await; triple_storage - .reserve(id4) + .create_slot(id4) .await - .unwrap() .insert(dummy_pair(id4), node0) .await; assert!(triple_spawner.contains(id3).await); @@ -108,7 +92,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_mine().await, 2); assert_eq!(triple_spawner.len_potential().await, 2); - // Take mine triple pairs and check that they are removed from the storage and added to used set + // Take mine triple pairs and check that they are removed from the storage and marked as using triple_storage.take_mine(node0).await.unwrap(); triple_storage.take_mine(node0).await.unwrap(); assert!(!triple_spawner.contains(id3).await); @@ -119,22 +103,15 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_mine().await, 0); assert!(triple_storage.is_empty().await); assert_eq!(triple_spawner.len_potential().await, 0); - assert!(triple_storage.contains_used(id3).await); - assert!(triple_storage.contains_used(id4).await); - - // Attempt to re-insert used mine triples and check that it fails - assert!(triple_storage.reserve(id3).await.is_none()); - assert!(triple_storage.reserve(id4).await.is_none()); - assert!(!triple_spawner.contains(id3).await); - assert!(!triple_spawner.contains(id4).await); + assert!(triple_storage.contains_using(id3).await); + assert!(triple_storage.contains_using(id4).await); assert!(triple_storage.clear().await); // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { triple_storage - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_pair(id), node1) .await; } @@ -142,9 +119,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { triple_storage - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_pair(id), node0) .await; } @@ -204,18 +180,11 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert!(presignature_storage.is_empty().await); assert_eq!(presignature_spawner.len_potential().await, 0); - // check that reserve then dropping unreserves the slot: - let slot = presignature_storage.reserve(presignature.id).await.unwrap(); - if let Some(task) = slot.unreserve() { - task.await.unwrap(); - } - // Insert presignature owned by node1, with our node0 view being that it is a foreign presignature assert!( presignature_storage - .reserve(presignature.id) + .create_slot(presignature.id) .await - .unwrap() .insert(presignature, node1) .await ); @@ -227,18 +196,14 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_mine().await, 0); assert_eq!(presignature_spawner.len_potential().await, 1); - // Take presignature and check that it is removed from the storage and added to used set + // Take presignature and check that it is removed from the storage and marked as using presignature_storage.take(id, node1).await.unwrap(); assert!(!presignature_storage.contains(id).await); assert!(!presignature_spawner.contains_mine(id).await); assert_eq!(presignature_storage.len_generated().await, 0); assert_eq!(presignature_spawner.len_mine().await, 0); assert_eq!(presignature_spawner.len_potential().await, 0); - assert!(presignature_storage.contains_used(id).await); - - // Attempt to re-insert used presignature and check that it fails - assert!(presignature_storage.reserve(id).await.is_none()); - assert!(!presignature_spawner.contains(id).await); + assert!(presignature_storage.contains_using(id).await); let id2 = 2; let mine_presignature = dummy_presignature(id2); @@ -246,9 +211,8 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Add a presignature to our own node0 assert!( presignature_storage - .reserve(id2) + .create_slot(id2) .await - .unwrap() .insert(mine_presignature, node0) .await ); @@ -259,7 +223,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_mine().await, 1); assert_eq!(presignature_spawner.len_potential().await, 1); - // Take mine presignature and check that it is removed from the storage and added to used set + // Take mine presignature and check that it is removed from the storage and marked as using presignature_storage.take_mine(node0).await.unwrap(); assert!(!presignature_storage.contains(id2).await); assert!(!presignature_spawner.contains_mine(id2).await); @@ -267,19 +231,14 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_mine().await, 0); assert!(presignature_storage.is_empty().await); assert_eq!(presignature_spawner.len_potential().await, 0); - assert!(presignature_storage.contains_used(id2).await); - - // Attempt to re-insert used mine presignature and check that it fails - assert!(presignature_storage.reserve(id2).await.is_none()); - assert!(!presignature_spawner.contains(id2).await); + assert!(presignature_storage.contains_using(id2).await); presignature_storage.clear().await; // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { presignature_storage - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_presignature(id), node1) .await; } @@ -287,9 +246,8 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { presignature_storage - .reserve(id) + .create_slot(id) .await - .unwrap() .insert(dummy_presignature(id), node0) .await; } From ac06acbd5b0735633fe87d51e80d0869f4fb0ef9 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 11:33:23 +0300 Subject: [PATCH 06/21] fix dropper --- .../node/src/storage/protocol_storage.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 4469fe047..507dc8530 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -86,9 +86,11 @@ impl Drop for ArtifactSlot { fn drop(&mut self) { let storage = self.storage.clone(); let id = self.id; - tokio::spawn(async move { - storage.generating.write().await.remove(&id); - }); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + storage.generating.write().await.remove(&id); + }); + } } } @@ -106,9 +108,11 @@ impl Drop for ArtifactTakenDropper { fn drop(&mut self) { if let Some(storage) = self.dropper.take() { let id = self.id; - tokio::spawn(async move { - storage.using.write().await.remove(&id); - }); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + storage.using.write().await.remove(&id); + }); + } } } } From 1295a8f1d30eac52b37266d34404ff1248cea8a9 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 11:44:48 +0300 Subject: [PATCH 07/21] add checks to artifact creation --- .../node/src/protocol/presignature.rs | 7 +++++- chain-signatures/node/src/protocol/triple.rs | 6 ++++- .../node/src/storage/protocol_storage.rs | 19 +++++++++++--- integration-tests/benches/store.rs | 8 +++++- integration-tests/src/containers.rs | 1 + integration-tests/src/mpc_fixture/builder.rs | 5 ++-- integration-tests/tests/cases/helpers.rs | 2 ++ integration-tests/tests/cases/store.rs | 25 +++++++++++++++++++ 8 files changed, 64 insertions(+), 9 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index 6d291323b..9728244fb 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -566,7 +566,12 @@ impl PresignatureSpawner { "starting protocol to generate a new presignature", ); - let slot = self.presignatures.create_slot(id.id).await; + let Some(slot) = self.presignatures.create_slot(id.id).await else { + return Err(InitializationError::BadParameters(format!( + "presignature {} is already generating, in use, or stored", + id.id + ))); + }; let mut participants = participants.to_vec(); participants.sort(); diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index e60a377a4..f38db070d 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -528,7 +528,11 @@ impl TripleSpawner { timeout: Duration, ) -> Result<(), InitializationError> { // Check if the `id` is already in the system. Error out and have the next cycle try again. - let slot = self.triple_storage.create_slot(id).await; + let Some(slot) = self.triple_storage.create_slot(id).await else { + return Err(InitializationError::BadParameters(format!( + "triple {id} is already generating, in use, or stored" + ))); + }; tracing::info!(?id, "starting protocol to generate a new triple"); let generator = TripleGenerator::new( diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 507dc8530..53defa3f1 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -206,18 +206,29 @@ impl ProtocolStorage { /// Create a slot for generating an artifact with the given ID. /// Tracks the ID in the `generating` set until the slot is inserted or dropped. - pub async fn create_slot(&self, id: A::Id) -> ArtifactSlot { + /// Returns `None` if the ID is already generating, in use, or stored in Redis. + pub async fn create_slot(&self, id: A::Id) -> Option> { + if self.using.read().await.contains(&id) { + tracing::error!(id, "cannot create slot: artifact is currently in use"); + return None; + } if !self.generating.write().await.insert(id) { tracing::error!( id, - "creating slot for artifact that is already being generated" + "cannot create slot: artifact is already being generated" ); + return None; } - ArtifactSlot { + if self.contains(id).await { + self.generating.write().await.remove(&id); + tracing::error!(id, "cannot create slot: artifact already exists in storage"); + return None; + } + Some(ArtifactSlot { id, storage: self.clone(), stored: false, - } + }) } /// Check if an artifact is currently being generated. diff --git a/integration-tests/benches/store.rs b/integration-tests/benches/store.rs index 47e057fdf..a111bb733 100644 --- a/integration-tests/benches/store.rs +++ b/integration-tests/benches/store.rs @@ -157,7 +157,12 @@ fn bench_load_keys(c: &mut Criterion) { rt.block_on(async { for i in 0..1000 { let t = dummy_pair(i); - env.triples.create_slot(t.id).await.insert(t, env.me).await; + env.triples + .create_slot(t.id) + .await + .unwrap() + .insert(t, env.me) + .await; } }); }) @@ -171,6 +176,7 @@ fn bench_load_keys(c: &mut Criterion) { env.presignatures .create_slot(p.id) .await + .unwrap() .insert(p, env.me) .await; } diff --git a/integration-tests/src/containers.rs b/integration-tests/src/containers.rs index d67bfd448..0fa925c1c 100644 --- a/integration-tests/src/containers.rs +++ b/integration-tests/src/containers.rs @@ -431,6 +431,7 @@ impl Redis { .unwrap() .create_slot(pair_id) .await + .unwrap() .insert(pair, *owner) .await; } diff --git a/integration-tests/src/mpc_fixture/builder.rs b/integration-tests/src/mpc_fixture/builder.rs index 6724bc1e1..bc9127eaa 100644 --- a/integration-tests/src/mpc_fixture/builder.rs +++ b/integration-tests/src/mpc_fixture/builder.rs @@ -564,7 +564,7 @@ impl MpcFixtureNodeBuilder { if pair.holders.is_none() { pair.holders = Some(pair.triple0.public.participants.clone()); } - let mut slot = triple_storage.create_slot(pair_id).await; + let mut slot = triple_storage.create_slot(pair_id).await.unwrap(); slot.insert(pair, owner).await; } } @@ -583,7 +583,8 @@ impl MpcFixtureNodeBuilder { } let mut slot = presignature_storage .create_slot(presignature_share.id) - .await; + .await + .unwrap(); slot.insert(presignature_share, owner).await; } } diff --git a/integration-tests/tests/cases/helpers.rs b/integration-tests/tests/cases/helpers.rs index 736fef149..75ac01058 100644 --- a/integration-tests/tests/cases/helpers.rs +++ b/integration-tests/tests/cases/helpers.rs @@ -69,6 +69,7 @@ pub(crate) async fn insert_triples_for_owner( triples .create_slot(id) .await + .unwrap() .insert(dummy_pair_with_holders(id, holders.clone()), owner) .await; } @@ -85,6 +86,7 @@ pub(crate) async fn insert_presignatures_for_owner( presignatures .create_slot(id) .await + .unwrap() .insert(dummy_presignature_with_holders(id, holders.clone()), owner) .await; } diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 240f21b23..9072100d9 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -40,11 +40,14 @@ async fn test_triple_persistence() -> anyhow::Result<()> { triple_storage .create_slot(triple_id1) .await + .unwrap() .insert(dummy_pair(triple_id1), node1) .await; triple_storage .create_slot(triple_id2) .await + .unwrap() + .await .insert(dummy_pair(triple_id2), node1) .await; @@ -70,6 +73,10 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert!(triple_storage.contains_using(triple_id1).await); assert!(triple_storage.contains_using(triple_id2).await); + // Attempt to re-create slot for in-use triples and check that it fails + assert!(triple_storage.create_slot(triple_id1).await.is_none()); + assert!(triple_storage.create_slot(triple_id2).await.is_none()); + let id3 = 3; let id4: u64 = 4; @@ -77,11 +84,13 @@ async fn test_triple_persistence() -> anyhow::Result<()> { triple_storage .create_slot(id3) .await + .unwrap() .insert(dummy_pair(id3), node0) .await; triple_storage .create_slot(id4) .await + .unwrap() .insert(dummy_pair(id4), node0) .await; assert!(triple_spawner.contains(id3).await); @@ -106,12 +115,17 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert!(triple_storage.contains_using(id3).await); assert!(triple_storage.contains_using(id4).await); + // Attempt to re-create slot for in-use mine triples and check that it fails + assert!(triple_storage.create_slot(id3).await.is_none()); + assert!(triple_storage.create_slot(id4).await.is_none()); + assert!(triple_storage.clear().await); // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { triple_storage .create_slot(id) .await + .unwrap() .insert(dummy_pair(id), node1) .await; } @@ -121,6 +135,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { triple_storage .create_slot(id) .await + .unwrap() .insert(dummy_pair(id), node0) .await; } @@ -185,6 +200,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { presignature_storage .create_slot(presignature.id) .await + .unwrap() .insert(presignature, node1) .await ); @@ -205,6 +221,9 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 0); assert!(presignature_storage.contains_using(id).await); + // Attempt to re-create slot for in-use presignature and check that it fails + assert!(presignature_storage.create_slot(id).await.is_none()); + let id2 = 2; let mine_presignature = dummy_presignature(id2); @@ -213,6 +232,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { presignature_storage .create_slot(id2) .await + .unwrap() .insert(mine_presignature, node0) .await ); @@ -233,12 +253,16 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 0); assert!(presignature_storage.contains_using(id2).await); + // Attempt to re-create slot for in-use mine presignature and check that it fails + assert!(presignature_storage.create_slot(id2).await.is_none()); + presignature_storage.clear().await; // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { presignature_storage .create_slot(id) .await + .unwrap() .insert(dummy_presignature(id), node1) .await; } @@ -248,6 +272,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { presignature_storage .create_slot(id) .await + .unwrap() .insert(dummy_presignature(id), node0) .await; } From 744020c074b6ac4ee1bd7fdef7c1cf287fe598e1 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 11:48:36 +0300 Subject: [PATCH 08/21] clippy --- integration-tests/tests/cases/store.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 9072100d9..24166a2e1 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -47,7 +47,6 @@ async fn test_triple_persistence() -> anyhow::Result<()> { .create_slot(triple_id2) .await .unwrap() - .await .insert(dummy_pair(triple_id2), node1) .await; From c6244a2e1881bc30dd31555a155367f2c1feb711 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 11:54:54 +0300 Subject: [PATCH 09/21] check ownership before pruning --- chain-signatures/node/src/storage/protocol_storage.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 53defa3f1..fd6b63e51 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -677,9 +677,11 @@ impl ProtocolStorage { local updated = {} for i = 3, #ARGV do local id = ARGV[i] - local holders_key = artifact_key .. ':holders:' .. id - -- Skip if holders set doesn't exist (artifact already taken/consumed) - if redis.call('EXISTS', holders_key) == 1 then + -- Skip if not owned by me (defense against malicious/buggy peer responses) + if redis.call('SISMEMBER', owner_key, id) == 0 then + -- noop: not our artifact + elseif redis.call('EXISTS', artifact_key .. ':holders:' .. id) == 1 then + local holders_key = artifact_key .. ':holders:' .. id redis.call('SREM', holders_key, peer) local count = redis.call('SCARD', holders_key) if count < threshold then From 30b70ec7bc42ac47136f4d09aa68228e5e53ed7f Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 13:01:52 +0300 Subject: [PATCH 10/21] remove redundant store bool flag --- chain-signatures/node/src/storage/protocol_storage.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index fd6b63e51..1da116b91 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -72,13 +72,11 @@ pub trait ProtocolArtifact: pub struct ArtifactSlot { id: A::Id, storage: ProtocolStorage, - stored: bool, } impl ArtifactSlot { pub async fn insert(&mut self, artifact: A, owner: Participant) -> bool { - self.stored = self.storage.insert(artifact, owner).await; - self.stored + self.storage.insert(artifact, owner).await } } @@ -227,7 +225,6 @@ impl ProtocolStorage { Some(ArtifactSlot { id, storage: self.clone(), - stored: false, }) } From 34137407511aad3f6780a5a2a5901f476b7414de Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 15:18:03 +0300 Subject: [PATCH 11/21] use MissedTickBehavior::Skip when stockpiling --- chain-signatures/node/src/protocol/presignature.rs | 1 + chain-signatures/node/src/protocol/triple.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index 9728244fb..0ad9410e2 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -697,6 +697,7 @@ impl PresignatureSpawner { ) { let mut last_active_warn: Option = None; let mut stockpile_interval = time::interval(Duration::from_millis(100)); + stockpile_interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip); let mut expiration_interval = tokio::time::interval(Duration::from_secs(1)); let mut posits = self.msg.subscribe_presignature_posit().await; diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index f38db070d..0f7efc4eb 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -581,6 +581,7 @@ impl TripleSpawner { ongoing_gen_tx: watch::Sender, ) { let mut stockpile_interval = tokio::time::interval(Duration::from_millis(100)); + stockpile_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut expiration_interval = tokio::time::interval(Duration::from_secs(1)); let mut posits = self.msg.subscribe_triple_posit().await; From 6a091c6427919cdccc925bd276cfd168970fd2b9 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 15:21:40 +0300 Subject: [PATCH 12/21] remove redundant drops --- chain-signatures/node/src/storage/protocol_storage.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 1da116b91..0a05fe1f5 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -321,8 +321,6 @@ impl ProtocolStorage { .into_iter() .filter(|id| !generating.contains(id) && !using.contains(id)) .collect(); - drop(generating); - drop(using); Ok(RemoveOutdatedResult::new(outdated, not_found)) } Err(err) => { From f66cd74e5272ad1304583c1999680aa801049fa4 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Thu, 2 Apr 2026 16:11:54 +0300 Subject: [PATCH 13/21] merge using and generating into a single lock --- .../node/src/storage/protocol_storage.rs | 84 +++++++++++-------- 1 file changed, 51 insertions(+), 33 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 0a05fe1f5..11a0b8058 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -86,7 +86,7 @@ impl Drop for ArtifactSlot { let id = self.id; if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { - storage.generating.write().await.remove(&id); + storage.reserved.write().await.generating.remove(&id); }); } } @@ -108,7 +108,7 @@ impl Drop for ArtifactTakenDropper { let id = self.id; if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { - storage.using.write().await.remove(&id); + storage.reserved.write().await.using.remove(&id); }); } } @@ -131,12 +131,30 @@ impl ArtifactTaken { } } +/// In-memory tracking of artifact IDs that are not yet in Redis. +/// Protected by a single `RwLock` to avoid multi-lock ordering issues. +#[derive(Debug)] +struct InMemoryState { + /// IDs currently being generated (protocol running, not yet persisted). + generating: HashSet, + /// IDs taken from Redis and actively consumed by a protocol. + using: HashSet, +} + +impl InMemoryState { + fn new() -> Self { + Self { + generating: HashSet::new(), + using: HashSet::new(), + } + } +} + #[derive(Debug)] pub struct ProtocolStorage { redis_pool: Pool, artifact_key: String, - generating: Arc>>, - using: Arc>>, + reserved: Arc>>, owner_keys: String, account_id: AccountId, _phantom: std::marker::PhantomData, @@ -147,8 +165,7 @@ impl Clone for ProtocolStorage { Self { redis_pool: self.redis_pool.clone(), artifact_key: self.artifact_key.clone(), - generating: self.generating.clone(), - using: self.using.clone(), + reserved: self.reserved.clone(), owner_keys: self.owner_keys.clone(), account_id: self.account_id.clone(), _phantom: std::marker::PhantomData, @@ -159,15 +176,13 @@ impl Clone for ProtocolStorage { impl ProtocolStorage { pub fn new(pool: &Pool, account_id: &AccountId, base_prefix: &str) -> Self { let artifact_key = format!("{base_prefix}:{STORAGE_VERSION}:{account_id}"); - let generating = Arc::new(RwLock::new(HashSet::new())); - let using = Arc::new(RwLock::new(HashSet::new())); + let state = Arc::new(RwLock::new(InMemoryState::new())); let owner_keys = format!("{base_prefix}_owners:{STORAGE_VERSION}:{account_id}"); Self { redis_pool: pool.clone(), artifact_key, - generating, - using, + reserved: state, owner_keys, account_id: account_id.clone(), _phantom: std::marker::PhantomData, @@ -206,19 +221,22 @@ impl ProtocolStorage { /// Tracks the ID in the `generating` set until the slot is inserted or dropped. /// Returns `None` if the ID is already generating, in use, or stored in Redis. pub async fn create_slot(&self, id: A::Id) -> Option> { - if self.using.read().await.contains(&id) { - tracing::error!(id, "cannot create slot: artifact is currently in use"); - return None; - } - if !self.generating.write().await.insert(id) { - tracing::error!( - id, - "cannot create slot: artifact is already being generated" - ); - return None; + { + let mut state = self.reserved.write().await; + if state.using.contains(&id) { + tracing::error!(id, "cannot create slot: artifact is currently in use"); + return None; + } + if !state.generating.insert(id) { + tracing::error!( + id, + "cannot create slot: artifact is already being generated" + ); + return None; + } } if self.contains(id).await { - self.generating.write().await.remove(&id); + self.reserved.write().await.generating.remove(&id); tracing::error!(id, "cannot create slot: artifact already exists in storage"); return None; } @@ -230,7 +248,7 @@ impl ProtocolStorage { /// Check if an artifact is currently being generated. pub async fn contains_generating(&self, id: A::Id) -> bool { - self.generating.read().await.contains(&id) + self.reserved.read().await.generating.contains(&id) } pub async fn remove_outdated( @@ -315,11 +333,10 @@ impl ProtocolStorage { // Filter out artifacts that are on this node but not in Redis: // - `generating`: being generated, not yet persisted // - `using`: taken from Redis, actively consumed by a protocol - let generating = self.generating.read().await; - let using = self.using.read().await; + let state = self.reserved.read().await; let not_found: Vec<_> = not_found .into_iter() - .filter(|id| !generating.contains(id) && !using.contains(id)) + .filter(|id| !state.generating.contains(id) && !state.using.contains(id)) .collect(); Ok(RemoveOutdatedResult::new(outdated, not_found)) } @@ -427,12 +444,12 @@ impl ProtocolStorage { /// Check if an artifact is currently being consumed by an active protocol. pub async fn contains_using(&self, id: A::Id) -> bool { - self.using.read().await.contains(&id) + self.reserved.read().await.using.contains(&id) } /// Returns the set of artifact IDs currently being consumed by active protocols. pub async fn using_ids(&self) -> HashSet { - self.using.read().await.clone() + self.reserved.read().await.using.clone() } pub async fn take(&self, id: A::Id, owner: Participant) -> Option> { @@ -460,14 +477,14 @@ impl ProtocolStorage { "#; let start = Instant::now(); - if !self.using.write().await.insert(id) { + if !self.reserved.write().await.using.insert(id) { tracing::warn!(id, "taking artifact that is already in use"); return None; } let Some(mut conn) = self.connect().await else { tracing::warn!(id, "failed to take artifact: connection failed"); - self.using.write().await.remove(&id); + self.reserved.write().await.using.remove(&id); return None; }; let result: Result<(A, Vec), _> = redis::Script::new(SCRIPT) @@ -490,7 +507,7 @@ impl ProtocolStorage { Some(ArtifactTaken::new(artifact, self.clone())) } Err(err) => { - self.using.write().await.remove(&id); + self.reserved.write().await.using.remove(&id); tracing::warn!(id, ?err, ?elapsed, "failed to take artifact"); None } @@ -571,8 +588,9 @@ impl ProtocolStorage { .with_label_values(&[A::METRIC_LABEL, "clear"]) .observe(elapsed.as_millis() as f64); - self.generating.write().await.clear(); - self.using.write().await.clear(); + let mut state = self.reserved.write().await; + state.generating.clear(); + state.using.clear(); // if the outcome is None, it means the script failed or there was an error. outcome.is_some() @@ -629,7 +647,7 @@ impl ProtocolStorage { let holders = holders.into_iter().map(Participant::from).collect(); artifact.set_holders(holders); let id = artifact.id(); - self.using.write().await.insert(id); + self.reserved.write().await.using.insert(id); let taken = ArtifactTaken::new(artifact, self.clone()); tracing::debug!(id, ?elapsed, "took mine artifact"); Some(taken) From 4f65cc29c2f222c902e64d10061fd32b6f74338a Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Fri, 3 Apr 2026 10:31:33 +0300 Subject: [PATCH 14/21] fix storage tests --- integration-tests/tests/cases/store.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 24166a2e1..036b97cfd 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -60,8 +60,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 2); // Take triple pairs and check that they are removed from the storage and marked as using - triple_storage.take(triple_id1, node1).await.unwrap(); - triple_storage.take(triple_id2, node1).await.unwrap(); + let _taken1 = triple_storage.take(triple_id1, node1).await.unwrap(); + let _taken2 = triple_storage.take(triple_id2, node1).await.unwrap(); assert!(!triple_spawner.contains(triple_id1).await); assert!(!triple_spawner.contains(triple_id2).await); assert!(!triple_spawner.contains_mine(triple_id1).await); @@ -101,8 +101,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 2); // Take mine triple pairs and check that they are removed from the storage and marked as using - triple_storage.take_mine(node0).await.unwrap(); - triple_storage.take_mine(node0).await.unwrap(); + let _taken3 = triple_storage.take_mine(node0).await.unwrap(); + let _taken4 = triple_storage.take_mine(node0).await.unwrap(); assert!(!triple_spawner.contains(id3).await); assert!(!triple_spawner.contains(id4).await); assert!(!triple_spawner.contains_mine(id3).await); @@ -212,7 +212,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 1); // Take presignature and check that it is removed from the storage and marked as using - presignature_storage.take(id, node1).await.unwrap(); + let _taken_ps1 = presignature_storage.take(id, node1).await.unwrap(); assert!(!presignature_storage.contains(id).await); assert!(!presignature_spawner.contains_mine(id).await); assert_eq!(presignature_storage.len_generated().await, 0); @@ -243,7 +243,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 1); // Take mine presignature and check that it is removed from the storage and marked as using - presignature_storage.take_mine(node0).await.unwrap(); + let _taken_ps2 = presignature_storage.take_mine(node0).await.unwrap(); assert!(!presignature_storage.contains(id2).await); assert!(!presignature_spawner.contains_mine(id2).await); assert_eq!(presignature_storage.len_generated().await, 0); From 258fa46a8b22e5dbcb17f68550fa32493891c02f Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Fri, 3 Apr 2026 16:15:27 +0300 Subject: [PATCH 15/21] add test_sync_when_peer_generating --- integration-tests/tests/cases/mpc.rs | 128 ++++++++++++++++++++++++++- 1 file changed, 126 insertions(+), 2 deletions(-) diff --git a/integration-tests/tests/cases/mpc.rs b/integration-tests/tests/cases/mpc.rs index 3cf6409b3..b5d883205 100644 --- a/integration-tests/tests/cases/mpc.rs +++ b/integration-tests/tests/cases/mpc.rs @@ -28,8 +28,8 @@ const TRIPLE_PAIRS_PER_OWNER: usize = 50; const PRESIGNATURES_PER_OWNER: usize = 25; use super::helpers::{ - assert_presig_owned_state, assert_triples_owned_state, insert_presignatures_for_owner, - insert_triples_for_owner, + assert_presig_owned_state, assert_triples_owned_state, dummy_pair_with_holders, + dummy_presignature_with_holders, insert_presignatures_for_owner, insert_triples_for_owner, }; #[test(tokio::test(flavor = "multi_thread"))] @@ -303,6 +303,130 @@ async fn test_sync_remove_outdated_orphan() { } } +/// When a peer is still generating an artifact (create_slot called, not yet inserted), +/// sync should NOT report it as missing. After generation fails (slot dropped), +/// the artifact should be reported as missing on the next sync. +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_when_peer_generating() { + const GENERATING_ID: u64 = 600; + + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let node0 = &fixture.nodes[0]; + let node1 = &fixture.nodes[1]; + let all_participants = fixture.sorted_participants(); + + // Owner (node0): create_slot + insert → artifact is fully generated and in Redis. + node0 + .triple_storage + .create_slot(GENERATING_ID) + .await + .unwrap() + .insert( + dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), + node0.me, + ) + .await; + node0 + .presignature_storage + .create_slot(GENERATING_ID) + .await + .unwrap() + .insert( + dummy_presignature_with_holders(GENERATING_ID, all_participants), + node0.me, + ) + .await; + + // Non-owner (node1): create_slot only → artifact is still generating, not in Redis. + let _triple_slot = node1 + .triple_storage + .create_slot(GENERATING_ID) + .await + .unwrap(); + let _presig_slot = node1 + .presignature_storage + .create_slot(GENERATING_ID) + .await + .unwrap(); + + assert!( + node1 + .triple_storage + .contains_generating(GENERATING_ID) + .await + ); + assert!( + node1 + .presignature_storage + .contains_generating(GENERATING_ID) + .await + ); + + // Sync: node0 broadcasts its directory (includes GENERATING_ID). + // node1 should NOT report it as missing because it's currently generating it. + let response = node1 + .sync( + node0.me, + node0.owned_triples().await, + node0.owned_presignatures().await, + ) + .await; + + assert!( + !response.triples.contains(&GENERATING_ID), + "triple should not be reported as missing while generating" + ); + assert!( + !response.presignatures.contains(&GENERATING_ID), + "presignature should not be reported as missing while generating" + ); + + // Simulate generation failure: drop the slots → no longer generating, not in Redis. + drop(_triple_slot); + drop(_presig_slot); + // ArtifactSlot::Drop removes from `generating` asynchronously; wait for it. + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if !node1 + .triple_storage + .contains_generating(GENERATING_ID) + .await + && !node1 + .presignature_storage + .contains_generating(GENERATING_ID) + .await + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("generating state should be cleared after dropping slots"); + + // Sync again: node1 should now report GENERATING_ID as missing. + let response = node1 + .sync( + node0.me, + node0.owned_triples().await, + node0.owned_presignatures().await, + ) + .await; + + assert!( + response.triples.contains(&GENERATING_ID), + "triple should be reported as missing after generation failed" + ); + assert!( + response.presignatures.contains(&GENERATING_ID), + "presignature should be reported as missing after generation failed" + ); +} + #[test(tokio::test(flavor = "multi_thread"))] async fn test_basic_generate_keys() { let network = MpcFixtureBuilder::new(5, 4).build().await; From dcd19af8523c8066140481848b1edf756fcd0901 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Sun, 5 Apr 2026 20:34:48 +0300 Subject: [PATCH 16/21] add test_sync_when_owner_generating and fixes --- .../node/src/protocol/presignature.rs | 6 +- .../node/src/protocol/sync/mod.rs | 5 +- chain-signatures/node/src/protocol/triple.rs | 6 +- .../node/src/storage/protocol_storage.rs | 51 +++++-- integration-tests/benches/store.rs | 4 +- integration-tests/src/containers.rs | 2 +- integration-tests/src/mpc_fixture/builder.rs | 4 +- .../src/mpc_fixture/fixture_interface.rs | 22 +++ integration-tests/tests/cases/helpers.rs | 4 +- integration-tests/tests/cases/mpc.rs | 141 ++++++++++++++++-- integration-tests/tests/cases/store.rs | 38 +++-- 11 files changed, 229 insertions(+), 54 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index 0ad9410e2..a6500d71a 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -566,7 +566,11 @@ impl PresignatureSpawner { "starting protocol to generate a new presignature", ); - let Some(slot) = self.presignatures.create_slot(id.id).await else { + let Some(slot) = self + .presignatures + .create_slot(id.id, owner == self.me) + .await + else { return Err(InitializationError::BadParameters(format!( "presignature {} is already generating, in use, or stored", id.id diff --git a/chain-signatures/node/src/protocol/sync/mod.rs b/chain-signatures/node/src/protocol/sync/mod.rs index 106e4dc81..690be096c 100644 --- a/chain-signatures/node/src/protocol/sync/mod.rs +++ b/chain-signatures/node/src/protocol/sync/mod.rs @@ -236,9 +236,8 @@ impl SyncTask { } } - // TODO: use reserved values instead. Note that we cannot fetch our own triples via reserved async fn new_update(&self, me: Participant) -> Option { - let triples = match self.triples.fetch_owned(me).await { + let triples = match self.triples.fetch_owned_with_generating(me).await { Ok(ids) => ids, Err(err) => { tracing::warn!( @@ -248,7 +247,7 @@ impl SyncTask { return None; } }; - let presignatures = match self.presignatures.fetch_owned(me).await { + let presignatures = match self.presignatures.fetch_owned_with_generating(me).await { Ok(ids) => ids, Err(err) => { tracing::warn!( diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index 0f7efc4eb..4bedc768b 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -528,7 +528,11 @@ impl TripleSpawner { timeout: Duration, ) -> Result<(), InitializationError> { // Check if the `id` is already in the system. Error out and have the next cycle try again. - let Some(slot) = self.triple_storage.create_slot(id).await else { + let Some(slot) = self + .triple_storage + .create_slot(id, proposer == self.me) + .await + else { return Err(InitializationError::BadParameters(format!( "triple {id} is already generating, in use, or stored" ))); diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index f98cd7f69..e9b9ba8ca 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -2,7 +2,7 @@ use cait_sith::protocol::Participant; use deadpool_redis::{Connection, Pool}; use near_sdk::AccountId; use redis::{AsyncCommands, FromRedisValue, ToRedisArgs}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::{fmt, time::Instant}; use tokio::sync::RwLock; @@ -131,30 +131,34 @@ impl ArtifactTaken { } } -/// In-memory tracking of artifact IDs that are not yet in Redis. +/// Tracks artifact IDs that are in-flight but not yet in Redis. /// Protected by a single `RwLock` to avoid multi-lock ordering issues. #[derive(Debug)] -struct InMemoryState { - /// IDs currently being generated (protocol running, not yet persisted). - generating: HashSet, +struct ReservedState { + /// IDs currently being generated. Value is `true` if this node is the owner/proposer. + generating: HashMap, /// IDs taken from Redis and actively consumed by a protocol. using: HashSet, } -impl InMemoryState { +impl ReservedState { fn new() -> Self { Self { - generating: HashSet::new(), + generating: HashMap::new(), using: HashSet::new(), } } + + fn contains_generating(&self, id: &Id) -> bool { + self.generating.contains_key(id) + } } #[derive(Debug)] pub struct ProtocolStorage { redis_pool: Pool, artifact_key: String, - reserved: Arc>>, + reserved: Arc>>, owner_keys: String, account_id: AccountId, _phantom: std::marker::PhantomData, @@ -176,7 +180,7 @@ impl Clone for ProtocolStorage { impl ProtocolStorage { pub fn new(pool: &Pool, account_id: &AccountId, base_prefix: &str) -> Self { let artifact_key = format!("{base_prefix}:{STORAGE_VERSION}:{account_id}"); - let state = Arc::new(RwLock::new(InMemoryState::new())); + let state = Arc::new(RwLock::new(ReservedState::new())); let owner_keys = format!("{base_prefix}_owners:{STORAGE_VERSION}:{account_id}"); Self { @@ -219,21 +223,23 @@ impl ProtocolStorage { /// Create a slot for generating an artifact with the given ID. /// Tracks the ID in the `generating` set until the slot is inserted or dropped. + /// `mine` indicates whether this node is the owner/proposer of the artifact. /// Returns `None` if the ID is already generating, in use, or stored in Redis. - pub async fn create_slot(&self, id: A::Id) -> Option> { + pub async fn create_slot(&self, id: A::Id, mine: bool) -> Option> { { let mut state = self.reserved.write().await; if state.using.contains(&id) { tracing::error!(id, "cannot create slot: artifact is currently in use"); return None; } - if !state.generating.insert(id) { + if state.contains_generating(&id) { tracing::error!( id, "cannot create slot: artifact is already being generated" ); return None; } + state.generating.insert(id, mine); } if self.contains(id).await { self.reserved.write().await.generating.remove(&id); @@ -246,9 +252,26 @@ impl ProtocolStorage { }) } - /// Check if an artifact is currently being generated. + /// Check if an artifact is currently being generated (mine or peer). pub async fn contains_generating(&self, id: A::Id) -> bool { - self.reserved.read().await.generating.contains(&id) + self.reserved.read().await.contains_generating(&id) + } + + /// Owned artifacts in Redis plus artifacts currently being generated + /// where this node is the owner/proposer. + pub async fn fetch_owned_with_generating( + &self, + me: Participant, + ) -> Result, StorageError> { + let mut ids = self.fetch_owned(me).await?; + let state = self.reserved.read().await; + ids.extend( + state + .generating + .iter() + .filter_map(|(&id, &mine)| mine.then_some(id)), + ); + Ok(ids) } pub async fn remove_outdated( @@ -336,7 +359,7 @@ impl ProtocolStorage { let state = self.reserved.read().await; let not_found: Vec<_> = not_found .into_iter() - .filter(|id| !state.generating.contains(id) && !state.using.contains(id)) + .filter(|id| !state.contains_generating(id) && !state.using.contains(id)) .collect(); Ok(RemoveOutdatedResult::new(outdated, not_found)) } diff --git a/integration-tests/benches/store.rs b/integration-tests/benches/store.rs index a111bb733..27a4543e7 100644 --- a/integration-tests/benches/store.rs +++ b/integration-tests/benches/store.rs @@ -158,7 +158,7 @@ fn bench_load_keys(c: &mut Criterion) { for i in 0..1000 { let t = dummy_pair(i); env.triples - .create_slot(t.id) + .create_slot(t.id, true) .await .unwrap() .insert(t, env.me) @@ -174,7 +174,7 @@ fn bench_load_keys(c: &mut Criterion) { for i in 0..1000 { let p = dummy_presignature(i); env.presignatures - .create_slot(p.id) + .create_slot(p.id, true) .await .unwrap() .insert(p, env.me) diff --git a/integration-tests/src/containers.rs b/integration-tests/src/containers.rs index 0fa925c1c..9df54a663 100644 --- a/integration-tests/src/containers.rs +++ b/integration-tests/src/containers.rs @@ -429,7 +429,7 @@ impl Redis { storage .get(me) .unwrap() - .create_slot(pair_id) + .create_slot(pair_id, true) .await .unwrap() .insert(pair, *owner) diff --git a/integration-tests/src/mpc_fixture/builder.rs b/integration-tests/src/mpc_fixture/builder.rs index 05af3231f..0325d58b3 100644 --- a/integration-tests/src/mpc_fixture/builder.rs +++ b/integration-tests/src/mpc_fixture/builder.rs @@ -580,7 +580,7 @@ impl MpcFixtureNodeBuilder { if pair.holders.is_none() { pair.holders = Some(pair.triple0.public.participants.clone()); } - let mut slot = triple_storage.create_slot(pair_id).await.unwrap(); + let mut slot = triple_storage.create_slot(pair_id, true).await.unwrap(); slot.insert(pair, owner).await; } } @@ -598,7 +598,7 @@ impl MpcFixtureNodeBuilder { presignature_share.holders = Some(presignature_share.participants.clone()); } let mut slot = presignature_storage - .create_slot(presignature_share.id) + .create_slot(presignature_share.id, true) .await .unwrap(); slot.insert(presignature_share, owner).await; diff --git a/integration-tests/src/mpc_fixture/fixture_interface.rs b/integration-tests/src/mpc_fixture/fixture_interface.rs index 47c76d001..61ac74520 100644 --- a/integration-tests/src/mpc_fixture/fixture_interface.rs +++ b/integration-tests/src/mpc_fixture/fixture_interface.rs @@ -199,6 +199,28 @@ impl MpcFixtureNode { ids } + /// Owned + currently-generating-as-owner, sorted. + pub async fn owned_triples_with_generating(&self) -> Vec { + let mut ids = self + .triple_storage + .fetch_owned_with_generating(self.me) + .await + .unwrap(); + ids.sort(); + ids + } + + /// Owned + currently-generating-as-owner, sorted. + pub async fn owned_presignatures_with_generating(&self) -> Vec { + let mut ids = self + .presignature_storage + .fetch_owned_with_generating(self.me) + .await + .unwrap(); + ids.sort(); + ids + } + /// Simulate the caller side of /sync: process a peer's response by removing /// the peer from artifacts they don't have, pruning below threshold. pub async fn process_sync_response( diff --git a/integration-tests/tests/cases/helpers.rs b/integration-tests/tests/cases/helpers.rs index 75ac01058..b0a0fc238 100644 --- a/integration-tests/tests/cases/helpers.rs +++ b/integration-tests/tests/cases/helpers.rs @@ -67,7 +67,7 @@ pub(crate) async fn insert_triples_for_owner( let holders = holders.to_vec(); for id in ids { triples - .create_slot(id) + .create_slot(id, true) // mine flag is not important, since we are inserting right away .await .unwrap() .insert(dummy_pair_with_holders(id, holders.clone()), owner) @@ -84,7 +84,7 @@ pub(crate) async fn insert_presignatures_for_owner( let holders = holders.to_vec(); for id in ids { presignatures - .create_slot(id) + .create_slot(id, true) // mine flag is not important, since we are inserting right away .await .unwrap() .insert(dummy_presignature_with_holders(id, holders.clone()), owner) diff --git a/integration-tests/tests/cases/mpc.rs b/integration-tests/tests/cases/mpc.rs index b5d883205..9ba69d818 100644 --- a/integration-tests/tests/cases/mpc.rs +++ b/integration-tests/tests/cases/mpc.rs @@ -57,7 +57,7 @@ async fn test_sync_noop_when_fully_synced() { let node0_triples = node0.owned_triples().await; let node0_presigs = node0.owned_presignatures().await; - // Responder side: node1 receives node0's directory and reports what it's missing. + // Responder side: node1 receives node0's sync update and reports what it's missing. let response = node1 .sync(node0.me, node0_triples.clone(), node0_presigs.clone()) .await; @@ -209,7 +209,7 @@ async fn test_sync_prune_below_threshold() { } /// Orphaned artifact: owner doesn't have id=77 but other nodes do. -/// When owner broadcasts its directory (without id=77), responders remove it +/// When owner broadcasts its sync update (without id=77), responders remove it /// via remove_outdated. #[test(tokio::test(flavor = "multi_thread"))] async fn test_sync_remove_outdated_orphan() { @@ -252,7 +252,7 @@ async fn test_sync_remove_outdated_orphan() { "presig 77 should have all participants as holders before sync" ); - // node0 broadcasts its directory (which does NOT include id=77 since node0 doesn't have it). + // node0 broadcasts its sync update (which does NOT include id=77 since node0 doesn't have it). // Responder runs remove_outdated, which removes id=77. let node0_triples = node0.owned_triples().await; let node0_presigs = node0.owned_presignatures().await; @@ -322,7 +322,7 @@ async fn test_sync_when_peer_generating() { // Owner (node0): create_slot + insert → artifact is fully generated and in Redis. node0 .triple_storage - .create_slot(GENERATING_ID) + .create_slot(GENERATING_ID, true) .await .unwrap() .insert( @@ -332,7 +332,7 @@ async fn test_sync_when_peer_generating() { .await; node0 .presignature_storage - .create_slot(GENERATING_ID) + .create_slot(GENERATING_ID, true) .await .unwrap() .insert( @@ -344,12 +344,12 @@ async fn test_sync_when_peer_generating() { // Non-owner (node1): create_slot only → artifact is still generating, not in Redis. let _triple_slot = node1 .triple_storage - .create_slot(GENERATING_ID) + .create_slot(GENERATING_ID, false) .await .unwrap(); let _presig_slot = node1 .presignature_storage - .create_slot(GENERATING_ID) + .create_slot(GENERATING_ID, false) .await .unwrap(); @@ -366,13 +366,13 @@ async fn test_sync_when_peer_generating() { .await ); - // Sync: node0 broadcasts its directory (includes GENERATING_ID). + // Sync: node0 broadcasts its update (includes GENERATING_ID). // node1 should NOT report it as missing because it's currently generating it. let response = node1 .sync( node0.me, - node0.owned_triples().await, - node0.owned_presignatures().await, + node0.owned_triples_with_generating().await, + node0.owned_presignatures_with_generating().await, ) .await; @@ -412,8 +412,8 @@ async fn test_sync_when_peer_generating() { let response = node1 .sync( node0.me, - node0.owned_triples().await, - node0.owned_presignatures().await, + node0.owned_triples_with_generating().await, + node0.owned_presignatures_with_generating().await, ) .await; @@ -427,6 +427,123 @@ async fn test_sync_when_peer_generating() { ); } +/// When the owner is still generating an artifact and initiates sync, +/// it should include generating in the sync update so peer nodes do not remove it. +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_when_owner_generating() { + const GENERATING_ID: u64 = 601; + + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let node0 = &fixture.nodes[0]; + let node1 = &fixture.nodes[1]; + let all_participants = fixture.sorted_participants(); + + // Owner (node0): create_slot only → still generating, not in Redis. + let mut triple_slot = node0 + .triple_storage + .create_slot(GENERATING_ID, true) + .await + .unwrap(); + let mut presig_slot = node0 + .presignature_storage + .create_slot(GENERATING_ID, true) + .await + .unwrap(); + + assert!( + node0 + .triple_storage + .contains_generating(GENERATING_ID) + .await + ); + assert!( + node0 + .presignature_storage + .contains_generating(GENERATING_ID) + .await + ); + + // Peer (node1): already finished generation, artifact is in Redis. + node1 + .triple_storage + .create_slot(GENERATING_ID, false) + .await + .unwrap() + .insert( + dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), + node0.me, + ) + .await; + node1 + .presignature_storage + .create_slot(GENERATING_ID, false) + .await + .unwrap() + .insert( + dummy_presignature_with_holders(GENERATING_ID, all_participants.clone()), + node0.me, + ) + .await; + + // Owner (node0) initiates sync — GENERATING_ID is not in Redis but IS in mine_generating. + let node0_triples = node0.owned_triples_with_generating().await; + let node0_presigs = node0.owned_presignatures_with_generating().await; + assert!(node0_triples.contains(&GENERATING_ID)); + assert!(node0_presigs.contains(&GENERATING_ID)); + + let response = node1.sync(node0.me, node0_triples, node0_presigs).await; + + // Peer (node1) sees GENERATING_ID in the update and already has it — nothing to report. + assert!( + !response.triples.contains(&GENERATING_ID), + "generating artifact should not appear in sync response" + ); + assert!( + !response.presignatures.contains(&GENERATING_ID), + "generating artifact should not appear in sync response" + ); + + // node1 must still have the artifact — remove_outdated must NOT delete it. + assert_triples_owned_state(&node1.triple_storage, node0.me, &[GENERATING_ID], &[]).await; + assert_presig_owned_state(&node1.presignature_storage, node0.me, &[GENERATING_ID], &[]).await; + + // Owner completes generation: insert into Redis. + triple_slot + .insert( + dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), + node0.me, + ) + .await; + presig_slot + .insert( + dummy_presignature_with_holders(GENERATING_ID, all_participants), + node0.me, + ) + .await; + + // Sync again: GENERATING_ID is now in Redis. node1 already has it, so nothing to report. + let response = node1 + .sync( + node0.me, + node0.owned_triples_with_generating().await, + node0.owned_presignatures_with_generating().await, + ) + .await; + + assert!( + !response.triples.contains(&GENERATING_ID), + "triple should not be reported as missing since node1 already has it" + ); + assert!( + !response.presignatures.contains(&GENERATING_ID), + "presignature should not be reported as missing since node1 already has it" + ); +} + #[test(tokio::test(flavor = "multi_thread"))] async fn test_basic_generate_keys() { let network = MpcFixtureBuilder::new(5, 4).build().await; diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 036b97cfd..4850ddaab 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -38,13 +38,13 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 0); triple_storage - .create_slot(triple_id1) + .create_slot(triple_id1, false) .await .unwrap() .insert(dummy_pair(triple_id1), node1) .await; triple_storage - .create_slot(triple_id2) + .create_slot(triple_id2, false) .await .unwrap() .insert(dummy_pair(triple_id2), node1) @@ -73,21 +73,27 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert!(triple_storage.contains_using(triple_id2).await); // Attempt to re-create slot for in-use triples and check that it fails - assert!(triple_storage.create_slot(triple_id1).await.is_none()); - assert!(triple_storage.create_slot(triple_id2).await.is_none()); + assert!(triple_storage + .create_slot(triple_id1, false) + .await + .is_none()); + assert!(triple_storage + .create_slot(triple_id2, false) + .await + .is_none()); let id3 = 3; let id4: u64 = 4; // Add mine triple and check that it is in the storage triple_storage - .create_slot(id3) + .create_slot(id3, true) .await .unwrap() .insert(dummy_pair(id3), node0) .await; triple_storage - .create_slot(id4) + .create_slot(id4, true) .await .unwrap() .insert(dummy_pair(id4), node0) @@ -115,14 +121,14 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert!(triple_storage.contains_using(id4).await); // Attempt to re-create slot for in-use mine triples and check that it fails - assert!(triple_storage.create_slot(id3).await.is_none()); - assert!(triple_storage.create_slot(id4).await.is_none()); + assert!(triple_storage.create_slot(id3, true).await.is_none()); + assert!(triple_storage.create_slot(id4, true).await.is_none()); assert!(triple_storage.clear().await); // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { triple_storage - .create_slot(id) + .create_slot(id, false) .await .unwrap() .insert(dummy_pair(id), node1) @@ -132,7 +138,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { triple_storage - .create_slot(id) + .create_slot(id, true) .await .unwrap() .insert(dummy_pair(id), node0) @@ -197,7 +203,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Insert presignature owned by node1, with our node0 view being that it is a foreign presignature assert!( presignature_storage - .create_slot(presignature.id) + .create_slot(presignature.id, false) .await .unwrap() .insert(presignature, node1) @@ -221,7 +227,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert!(presignature_storage.contains_using(id).await); // Attempt to re-create slot for in-use presignature and check that it fails - assert!(presignature_storage.create_slot(id).await.is_none()); + assert!(presignature_storage.create_slot(id, false).await.is_none()); let id2 = 2; let mine_presignature = dummy_presignature(id2); @@ -229,7 +235,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Add a presignature to our own node0 assert!( presignature_storage - .create_slot(id2) + .create_slot(id2, true) .await .unwrap() .insert(mine_presignature, node0) @@ -253,13 +259,13 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert!(presignature_storage.contains_using(id2).await); // Attempt to re-create slot for in-use mine presignature and check that it fails - assert!(presignature_storage.create_slot(id2).await.is_none()); + assert!(presignature_storage.create_slot(id2, true).await.is_none()); presignature_storage.clear().await; // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { presignature_storage - .create_slot(id) + .create_slot(id, false) .await .unwrap() .insert(dummy_presignature(id), node1) @@ -269,7 +275,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { presignature_storage - .create_slot(id) + .create_slot(id, true) .await .unwrap() .insert(dummy_presignature(id), node0) From 20008f0bae8e2cecd2101381874fdc3138cba96e Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Mon, 6 Apr 2026 16:11:05 +0300 Subject: [PATCH 17/21] add matrix sync test, distinguish owned using --- .../node/src/protocol/presignature.rs | 2 +- .../node/src/protocol/signature.rs | 2 +- .../node/src/protocol/sync/mod.rs | 4 +- .../node/src/storage/protocol_storage.rs | 37 +- .../src/mpc_fixture/fixture_interface.rs | 12 +- integration-tests/tests/cases/mpc.rs | 417 ++++++++---------- integration-tests/tests/cases/store.rs | 6 +- 7 files changed, 234 insertions(+), 246 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index a6500d71a..c2f08e0e8 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -865,7 +865,7 @@ impl PendingTriples { let mut interval = tokio::time::interval(Duration::from_millis(200)); loop { interval.tick().await; - if let Some(triples) = storage.take(pair_id, owner).await { + if let Some(triples) = storage.take(pair_id, owner, false).await { break triples; }; } diff --git a/chain-signatures/node/src/protocol/signature.rs b/chain-signatures/node/src/protocol/signature.rs index 7f89a32fb..cd4fad540 100644 --- a/chain-signatures/node/src/protocol/signature.rs +++ b/chain-signatures/node/src/protocol/signature.rs @@ -1471,7 +1471,7 @@ impl PendingPresignature { let mut interval = tokio::time::interval(Duration::from_millis(250)); loop { interval.tick().await; - if let Some(presignature) = storage.take(id, owner).await { + if let Some(presignature) = storage.take(id, owner, false).await { break presignature; }; } diff --git a/chain-signatures/node/src/protocol/sync/mod.rs b/chain-signatures/node/src/protocol/sync/mod.rs index 690be096c..cb5407892 100644 --- a/chain-signatures/node/src/protocol/sync/mod.rs +++ b/chain-signatures/node/src/protocol/sync/mod.rs @@ -237,7 +237,7 @@ impl SyncTask { } async fn new_update(&self, me: Participant) -> Option { - let triples = match self.triples.fetch_owned_with_generating(me).await { + let triples = match self.triples.fetch_owned_with_reserved(me).await { Ok(ids) => ids, Err(err) => { tracing::warn!( @@ -247,7 +247,7 @@ impl SyncTask { return None; } }; - let presignatures = match self.presignatures.fetch_owned_with_generating(me).await { + let presignatures = match self.presignatures.fetch_owned_with_reserved(me).await { Ok(ids) => ids, Err(err) => { tracing::warn!( diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index e9b9ba8ca..0469f70de 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -138,14 +138,15 @@ struct ReservedState { /// IDs currently being generated. Value is `true` if this node is the owner/proposer. generating: HashMap, /// IDs taken from Redis and actively consumed by a protocol. - using: HashSet, + /// Value is `true` if this node is the owner of the artifact. + using: HashMap, } impl ReservedState { fn new() -> Self { Self { generating: HashMap::new(), - using: HashSet::new(), + using: HashMap::new(), } } @@ -228,7 +229,7 @@ impl ProtocolStorage { pub async fn create_slot(&self, id: A::Id, mine: bool) -> Option> { { let mut state = self.reserved.write().await; - if state.using.contains(&id) { + if state.using.contains_key(&id) { tracing::error!(id, "cannot create slot: artifact is currently in use"); return None; } @@ -257,9 +258,10 @@ impl ProtocolStorage { self.reserved.read().await.contains_generating(&id) } - /// Owned artifacts in Redis plus artifacts currently being generated - /// where this node is the owner/proposer. - pub async fn fetch_owned_with_generating( + /// Owned artifacts in Redis plus owned using and owned generating. + /// This is the full set that should be advertised during state sync to prevent + /// peers from pruning artifacts that are still actively in use. + pub async fn fetch_owned_with_reserved( &self, me: Participant, ) -> Result, StorageError> { @@ -271,6 +273,12 @@ impl ProtocolStorage { .iter() .filter_map(|(&id, &mine)| mine.then_some(id)), ); + ids.extend( + state + .using + .iter() + .filter_map(|(&id, &mine)| mine.then_some(id)), + ); Ok(ids) } @@ -359,7 +367,7 @@ impl ProtocolStorage { let state = self.reserved.read().await; let not_found: Vec<_> = not_found .into_iter() - .filter(|id| !state.contains_generating(id) && !state.using.contains(id)) + .filter(|id| !state.contains_generating(id) && !state.using.contains_key(id)) .collect(); Ok(RemoveOutdatedResult::new(outdated, not_found)) } @@ -467,15 +475,20 @@ impl ProtocolStorage { /// Check if an artifact is currently being consumed by an active protocol. pub async fn contains_using(&self, id: A::Id) -> bool { - self.reserved.read().await.using.contains(&id) + self.reserved.read().await.using.contains_key(&id) } /// Returns the set of artifact IDs currently being consumed by active protocols. pub async fn using_ids(&self) -> HashSet { - self.reserved.read().await.using.clone() + self.reserved.read().await.using.keys().copied().collect() } - pub async fn take(&self, id: A::Id, owner: Participant) -> Option> { + pub async fn take( + &self, + id: A::Id, + owner: Participant, + mine: bool, + ) -> Option> { const SCRIPT: &str = r#" local artifact_key = KEYS[1] local owner_key = KEYS[2] @@ -500,7 +513,7 @@ impl ProtocolStorage { "#; let start = Instant::now(); - if !self.reserved.write().await.using.insert(id) { + if self.reserved.write().await.using.insert(id, mine).is_some() { tracing::warn!(id, "taking artifact that is already in use"); return None; } @@ -670,7 +683,7 @@ impl ProtocolStorage { let holders = holders.into_iter().map(Participant::from).collect(); artifact.set_holders(holders); let id = artifact.id(); - self.reserved.write().await.using.insert(id); + self.reserved.write().await.using.insert(id, true); let taken = ArtifactTaken::new(artifact, self.clone()); tracing::debug!(id, ?elapsed, "took mine artifact"); Some(taken) diff --git a/integration-tests/src/mpc_fixture/fixture_interface.rs b/integration-tests/src/mpc_fixture/fixture_interface.rs index 61ac74520..dd95b4b79 100644 --- a/integration-tests/src/mpc_fixture/fixture_interface.rs +++ b/integration-tests/src/mpc_fixture/fixture_interface.rs @@ -199,22 +199,22 @@ impl MpcFixtureNode { ids } - /// Owned + currently-generating-as-owner, sorted. - pub async fn owned_triples_with_generating(&self) -> Vec { + /// Owned + owned using + owned generating, sorted. + pub async fn owned_triples_with_reserved(&self) -> Vec { let mut ids = self .triple_storage - .fetch_owned_with_generating(self.me) + .fetch_owned_with_reserved(self.me) .await .unwrap(); ids.sort(); ids } - /// Owned + currently-generating-as-owner, sorted. - pub async fn owned_presignatures_with_generating(&self) -> Vec { + /// Owned + owned using + owned generating, sorted. + pub async fn owned_presignatures_with_reserved(&self) -> Vec { let mut ids = self .presignature_storage - .fetch_owned_with_generating(self.me) + .fetch_owned_with_reserved(self.me) .await .unwrap(); ids.sort(); diff --git a/integration-tests/tests/cases/mpc.rs b/integration-tests/tests/cases/mpc.rs index 9ba69d818..3b0e6783e 100644 --- a/integration-tests/tests/cases/mpc.rs +++ b/integration-tests/tests/cases/mpc.rs @@ -28,8 +28,8 @@ const TRIPLE_PAIRS_PER_OWNER: usize = 50; const PRESIGNATURES_PER_OWNER: usize = 25; use super::helpers::{ - assert_presig_owned_state, assert_triples_owned_state, dummy_pair_with_holders, - dummy_presignature_with_holders, insert_presignatures_for_owner, insert_triples_for_owner, + assert_presig_owned_state, assert_triples_owned_state, insert_presignatures_for_owner, + insert_triples_for_owner, }; #[test(tokio::test(flavor = "multi_thread"))] @@ -303,245 +303,220 @@ async fn test_sync_remove_outdated_orphan() { } } -/// When a peer is still generating an artifact (create_slot called, not yet inserted), -/// sync should NOT report it as missing. After generation fails (slot dropped), -/// the artifact should be reported as missing on the next sync. #[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_when_peer_generating() { - const GENERATING_ID: u64 = 600; - - let fixture = MpcFixtureBuilder::default() - .only_generate_signatures() - .build() - .await; - - let node0 = &fixture.nodes[0]; - let node1 = &fixture.nodes[1]; - let all_participants = fixture.sorted_participants(); - - // Owner (node0): create_slot + insert → artifact is fully generated and in Redis. - node0 - .triple_storage - .create_slot(GENERATING_ID, true) - .await - .unwrap() - .insert( - dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), - node0.me, - ) - .await; - node0 - .presignature_storage - .create_slot(GENERATING_ID, true) - .await - .unwrap() - .insert( - dummy_presignature_with_holders(GENERATING_ID, all_participants), - node0.me, - ) - .await; - - // Non-owner (node1): create_slot only → artifact is still generating, not in Redis. - let _triple_slot = node1 - .triple_storage - .create_slot(GENERATING_ID, false) - .await - .unwrap(); - let _presig_slot = node1 - .presignature_storage - .create_slot(GENERATING_ID, false) - .await - .unwrap(); - - assert!( - node1 - .triple_storage - .contains_generating(GENERATING_ID) - .await - ); - assert!( - node1 - .presignature_storage - .contains_generating(GENERATING_ID) - .await - ); - - // Sync: node0 broadcasts its update (includes GENERATING_ID). - // node1 should NOT report it as missing because it's currently generating it. - let response = node1 - .sync( - node0.me, - node0.owned_triples_with_generating().await, - node0.owned_presignatures_with_generating().await, - ) - .await; - - assert!( - !response.triples.contains(&GENERATING_ID), - "triple should not be reported as missing while generating" - ); - assert!( - !response.presignatures.contains(&GENERATING_ID), - "presignature should not be reported as missing while generating" - ); +async fn test_sync_matrix() { + #[derive(Debug, Clone, Copy)] + enum ArtifactState { + Generating, + Stored, + Using, + None, + } - // Simulate generation failure: drop the slots → no longer generating, not in Redis. - drop(_triple_slot); - drop(_presig_slot); - // ArtifactSlot::Drop removes from `generating` asynchronously; wait for it. - tokio::time::timeout(Duration::from_secs(2), async { - loop { - if !node1 - .triple_storage - .contains_generating(GENERATING_ID) - .await - && !node1 - .presignature_storage - .contains_generating(GENERATING_ID) - .await - { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("generating state should be cleared after dropping slots"); + struct ExpectedCallerState { + /// Should the artifact appear in the caller's sync update? + in_update: bool, + /// Should the caller still have it in Redis after process_sync_response? + stored_after: bool, + } - // Sync again: node1 should now report GENERATING_ID as missing. - let response = node1 - .sync( - node0.me, - node0.owned_triples_with_generating().await, - node0.owned_presignatures_with_generating().await, - ) - .await; + struct ExpectedResponderState { + /// Should the responder report the artifact as not_found (missing)? + missing: bool, + /// Should the responder still have it in Redis after sync? + stored_after: bool, + } - assert!( - response.triples.contains(&GENERATING_ID), - "triple should be reported as missing after generation failed" - ); - assert!( - response.presignatures.contains(&GENERATING_ID), - "presignature should be reported as missing after generation failed" - ); -} + struct Case { + caller: ArtifactState, + responder: ArtifactState, + expected_caller: ExpectedCallerState, + expected_responder: ExpectedResponderState, + } -/// When the owner is still generating an artifact and initiates sync, -/// it should include generating in the sync update so peer nodes do not remove it. -#[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_when_owner_generating() { - const GENERATING_ID: u64 = 601; + // Caller is always the owner (mine=true), responder is the peer (mine=false). + #[rustfmt::skip] // cargo fmt makes the matrix unreadable + let test_matrix = [ + // Caller: Stored (in update) ────────────────────────────── + Case { caller: ArtifactState::Stored, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: Generating (in update) ────────────────────────── + Case { caller: ArtifactState::Generating, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: Using (in update) ────────────────────────────── + Case { caller: ArtifactState::Using, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: None (NOT in update) ──────────────────────────── + Case { caller: ArtifactState::None, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + ]; let fixture = MpcFixtureBuilder::default() .only_generate_signatures() .build() .await; - let node0 = &fixture.nodes[0]; - let node1 = &fixture.nodes[1]; + let caller = &fixture.nodes[0]; + let responder = &fixture.nodes[1]; let all_participants = fixture.sorted_participants(); - // Owner (node0): create_slot only → still generating, not in Redis. - let mut triple_slot = node0 - .triple_storage - .create_slot(GENERATING_ID, true) - .await - .unwrap(); - let mut presig_slot = node0 - .presignature_storage - .create_slot(GENERATING_ID, true) - .await - .unwrap(); - - assert!( - node0 - .triple_storage - .contains_generating(GENERATING_ID) - .await - ); - assert!( - node0 - .presignature_storage - .contains_generating(GENERATING_ID) - .await - ); - - // Peer (node1): already finished generation, artifact is in Redis. - node1 - .triple_storage - .create_slot(GENERATING_ID, false) - .await - .unwrap() - .insert( - dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), - node0.me, - ) - .await; - node1 - .presignature_storage - .create_slot(GENERATING_ID, false) - .await - .unwrap() - .insert( - dummy_presignature_with_holders(GENERATING_ID, all_participants.clone()), - node0.me, - ) - .await; + for (i, case) in test_matrix.iter().enumerate() { + let id: u64 = 800 + i as u64; + tracing::info!(id, ?case.caller, ?case.responder, "=== case {i} ==="); + + // Hold slots/taken artifacts alive until assertions are done. + let mut caller_slot = None; + let mut caller_taken = None; + let mut responder_slot = None; + let mut responder_taken = None; + + // --- Set up caller (owner) state --- + match case.caller { + ArtifactState::Stored => { + insert_triples_for_owner( + &caller.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + } + ArtifactState::Generating => { + caller_slot = Some(caller.triple_storage.create_slot(id, true).await.unwrap()); + } + ArtifactState::Using => { + insert_triples_for_owner( + &caller.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + caller_taken = Some( + caller + .triple_storage + .take(id, caller.me, true) + .await + .unwrap(), + ); + } + ArtifactState::None => {} + } - // Owner (node0) initiates sync — GENERATING_ID is not in Redis but IS in mine_generating. - let node0_triples = node0.owned_triples_with_generating().await; - let node0_presigs = node0.owned_presignatures_with_generating().await; - assert!(node0_triples.contains(&GENERATING_ID)); - assert!(node0_presigs.contains(&GENERATING_ID)); + // --- Set up responder (peer) state --- + match case.responder { + ArtifactState::Stored => { + insert_triples_for_owner( + &responder.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + } + ArtifactState::Generating => { + responder_slot = Some( + responder + .triple_storage + .create_slot(id, false) + .await + .unwrap(), + ); + } + ArtifactState::Using => { + insert_triples_for_owner( + &responder.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + responder_taken = Some( + responder + .triple_storage + .take(id, caller.me, false) + .await + .unwrap(), + ); + } + ArtifactState::None => {} + } - let response = node1.sync(node0.me, node0_triples, node0_presigs).await; + // --- Build caller's sync update --- + let caller_update = caller.owned_triples_with_reserved().await; + assert_eq!( + caller_update.contains(&id), + case.expected_caller.in_update, + "case {i}: caller={:?} → expected in_update={}", + case.caller, + case.expected_caller.in_update, + ); - // Peer (node1) sees GENERATING_ID in the update and already has it — nothing to report. - assert!( - !response.triples.contains(&GENERATING_ID), - "generating artifact should not appear in sync response" - ); - assert!( - !response.presignatures.contains(&GENERATING_ID), - "generating artifact should not appear in sync response" - ); + // --- Responder processes the sync update --- + let response = responder + .sync( + caller.me, + caller_update, + responder.owned_presignatures().await, + ) + .await; - // node1 must still have the artifact — remove_outdated must NOT delete it. - assert_triples_owned_state(&node1.triple_storage, node0.me, &[GENERATING_ID], &[]).await; - assert_presig_owned_state(&node1.presignature_storage, node0.me, &[GENERATING_ID], &[]).await; + // Verify the full SyncUpdate response from the responder. + assert_eq!( + response.from, responder.me, + "case {i}: response.from should be the responder", + ); + assert_eq!( + response.triples.contains(&id), + case.expected_responder.missing, + "case {i}: caller={:?}, responder={:?} → expected missing={}", + case.caller, + case.responder, + case.expected_responder.missing, + ); - // Owner completes generation: insert into Redis. - triple_slot - .insert( - dummy_pair_with_holders(GENERATING_ID, all_participants.clone()), - node0.me, - ) - .await; - presig_slot - .insert( - dummy_presignature_with_holders(GENERATING_ID, all_participants), - node0.me, - ) - .await; + // --- Verify responder's Redis state after sync --- + assert_eq!( + responder + .triple_storage + .contains_by_owner(id, caller.me) + .await, + case.expected_responder.stored_after, + "case {i}: caller={:?}, responder={:?} → expected responder stored_after={}", + case.caller, + case.responder, + case.expected_responder.stored_after, + ); - // Sync again: GENERATING_ID is now in Redis. node1 already has it, so nothing to report. - let response = node1 - .sync( - node0.me, - node0.owned_triples_with_generating().await, - node0.owned_presignatures_with_generating().await, - ) - .await; + // --- Caller processes the sync response (remove holder / prune) --- + caller + .process_sync_response(responder.me, 2, &response) + .await; + assert_eq!( + caller.triple_storage.contains_by_owner(id, caller.me).await, + case.expected_caller.stored_after, + "case {i}: caller={:?}, responder={:?} → expected caller stored_after={}", + case.caller, + case.responder, + case.expected_caller.stored_after, + ); - assert!( - !response.triples.contains(&GENERATING_ID), - "triple should not be reported as missing since node1 already has it" - ); - assert!( - !response.presignatures.contains(&GENERATING_ID), - "presignature should not be reported as missing since node1 already has it" - ); + // Clean up held slots/taken before next iteration. + drop(caller_slot); + drop(caller_taken); + drop(responder_slot); + drop(responder_taken); + // Wait for async Drop cleanup. + tokio::time::sleep(Duration::from_millis(50)).await; + } } #[test(tokio::test(flavor = "multi_thread"))] diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index 4850ddaab..abbad1e51 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -60,8 +60,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 2); // Take triple pairs and check that they are removed from the storage and marked as using - let _taken1 = triple_storage.take(triple_id1, node1).await.unwrap(); - let _taken2 = triple_storage.take(triple_id2, node1).await.unwrap(); + let _taken1 = triple_storage.take(triple_id1, node1, false).await.unwrap(); + let _taken2 = triple_storage.take(triple_id2, node1, false).await.unwrap(); assert!(!triple_spawner.contains(triple_id1).await); assert!(!triple_spawner.contains(triple_id2).await); assert!(!triple_spawner.contains_mine(triple_id1).await); @@ -218,7 +218,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 1); // Take presignature and check that it is removed from the storage and marked as using - let _taken_ps1 = presignature_storage.take(id, node1).await.unwrap(); + let _taken_ps1 = presignature_storage.take(id, node1, false).await.unwrap(); assert!(!presignature_storage.contains(id).await); assert!(!presignature_spawner.contains_mine(id).await); assert_eq!(presignature_storage.len_generated().await, 0); From 62d3456096314ac83f956e9db60cc21e15c0b0a4 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Mon, 6 Apr 2026 22:04:20 +0300 Subject: [PATCH 18/21] keep me in storage --- .../node/src/protocol/consensus.rs | 6 +++ .../node/src/protocol/presignature.rs | 10 ++-- .../node/src/protocol/signature.rs | 4 +- .../node/src/protocol/sync/mod.rs | 11 ++-- chain-signatures/node/src/protocol/triple.rs | 6 +-- .../node/src/storage/protocol_storage.rs | 54 +++++++++++-------- integration-tests/benches/store.rs | 12 ++--- integration-tests/src/cluster/mod.rs | 12 ----- integration-tests/src/containers.rs | 40 +++++++++----- integration-tests/src/lib.rs | 6 --- integration-tests/src/mpc_fixture/builder.rs | 6 ++- .../src/mpc_fixture/fixture_interface.rs | 16 +++--- integration-tests/tests/cases/helpers.rs | 4 +- integration-tests/tests/cases/mpc.rs | 30 +++++------ integration-tests/tests/cases/store.rs | 50 ++++++++--------- integration-tests/tests/cases/sync.rs | 20 +++---- 16 files changed, 141 insertions(+), 146 deletions(-) diff --git a/chain-signatures/node/src/protocol/consensus.rs b/chain-signatures/node/src/protocol/consensus.rs index bce96aa8a..11d8f62fb 100644 --- a/chain-signatures/node/src/protocol/consensus.rs +++ b/chain-signatures/node/src/protocol/consensus.rs @@ -97,6 +97,9 @@ impl ConsensusProtocol for StartedState { ); let threshold = contract_state.threshold; + // Initialize identity for storage; this is an entry point into Running. + ctx.triple_storage.set_me(me); + ctx.presignature_storage.set_me(me); let triple_task = TripleSpawnerTask::run(me, threshold, epoch, ctx); let presign_task = PresignatureSpawnerTask::run( me, @@ -399,6 +402,9 @@ impl ConsensusProtocol for WaitingForConsensusState { return NodeState::WaitingForConsensus(self); }; + // Initialize identity for storage; this is an entry point into Running. + ctx.triple_storage.set_me(me); + ctx.presignature_storage.set_me(me); let triple_task = TripleSpawnerTask::run(me, self.threshold, self.epoch, ctx); let presign_task = PresignatureSpawnerTask::run( me, diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index c2f08e0e8..04aa58bcd 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -480,7 +480,7 @@ impl PresignatureSpawner { // to use the same triple as any other node. // TODO: have all this part be a separate task such that finding a pair of triples is done in parallel instead // of waiting for storage to respond here. - let Some(triples) = self.triples.take_mine(self.me).await else { + let Some(triples) = self.triples.take_mine().await else { return; }; @@ -566,11 +566,7 @@ impl PresignatureSpawner { "starting protocol to generate a new presignature", ); - let Some(slot) = self - .presignatures - .create_slot(id.id, owner == self.me) - .await - else { + let Some(slot) = self.presignatures.create_slot(id.id, owner).await else { return Err(InitializationError::BadParameters(format!( "presignature {} is already generating, in use, or stored", id.id @@ -865,7 +861,7 @@ impl PendingTriples { let mut interval = tokio::time::interval(Duration::from_millis(200)); loop { interval.tick().await; - if let Some(triples) = storage.take(pair_id, owner, false).await { + if let Some(triples) = storage.take(pair_id, owner).await { break triples; }; } diff --git a/chain-signatures/node/src/protocol/signature.rs b/chain-signatures/node/src/protocol/signature.rs index cd4fad540..6ec47964e 100644 --- a/chain-signatures/node/src/protocol/signature.rs +++ b/chain-signatures/node/src/protocol/signature.rs @@ -328,7 +328,7 @@ impl SignOrganizer { let remaining = state.budget.remaining(); let fetch = tokio::time::timeout(remaining, async { loop { - if let Some(taken) = ctx.presignatures.take_mine(ctx.me).await { + if let Some(taken) = ctx.presignatures.take_mine().await { let Some(holders) = taken.artifact.holders() else { tracing::error!( id = taken.artifact.id, @@ -1471,7 +1471,7 @@ impl PendingPresignature { let mut interval = tokio::time::interval(Duration::from_millis(250)); loop { interval.tick().await; - if let Some(presignature) = storage.take(id, owner, false).await { + if let Some(presignature) = storage.take(id, owner).await { break presignature; }; } diff --git a/chain-signatures/node/src/protocol/sync/mod.rs b/chain-signatures/node/src/protocol/sync/mod.rs index cb5407892..d8b5391cc 100644 --- a/chain-signatures/node/src/protocol/sync/mod.rs +++ b/chain-signatures/node/src/protocol/sync/mod.rs @@ -219,7 +219,7 @@ impl SyncTask { match handle.await { Ok(responses) => { // Process sync responses: update artifact participants based on not_found data - if let Err(err) = self.process_sync_responses(responses, me, threshold).await { + if let Err(err) = self.process_sync_responses(responses, threshold).await { tracing::warn!(?err, "failed to process sync responses"); } tracing::debug!(elapsed = ?start.elapsed(), "processed broadcast"); @@ -237,7 +237,7 @@ impl SyncTask { } async fn new_update(&self, me: Participant) -> Option { - let triples = match self.triples.fetch_owned_with_reserved(me).await { + let triples = match self.triples.fetch_owned_with_reserved().await { Ok(ids) => ids, Err(err) => { tracing::warn!( @@ -247,7 +247,7 @@ impl SyncTask { return None; } }; - let presignatures = match self.presignatures.fetch_owned_with_reserved(me).await { + let presignatures = match self.presignatures.fetch_owned_with_reserved().await { Ok(ids) => ids, Err(err) => { tracing::warn!( @@ -271,7 +271,6 @@ impl SyncTask { async fn process_sync_responses( &self, responses: Vec<(Participant, SyncPeerResponse)>, - me: Participant, threshold: usize, ) -> Result<(), String> { for (peer, result) in responses { @@ -293,13 +292,13 @@ impl SyncTask { // Batch remove peer from all triples and prune let triple_res = self .triples - .remove_holder_and_prune(me, peer, threshold, &response.triples) + .remove_holder_and_prune(peer, threshold, &response.triples) .await; // Batch remove peer from all presignatures and prune let presig_res = self .presignatures - .remove_holder_and_prune(me, peer, threshold, &response.presignatures) + .remove_holder_and_prune(peer, threshold, &response.presignatures) .await; match (triple_res, presig_res) { diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index 4bedc768b..e66a3f7da 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -528,11 +528,7 @@ impl TripleSpawner { timeout: Duration, ) -> Result<(), InitializationError> { // Check if the `id` is already in the system. Error out and have the next cycle try again. - let Some(slot) = self - .triple_storage - .create_slot(id, proposer == self.me) - .await - else { + let Some(slot) = self.triple_storage.create_slot(id, proposer).await else { return Err(InitializationError::BadParameters(format!( "triple {id} is already generating, in use, or stored" ))); diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 0469f70de..2959323fe 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -3,7 +3,7 @@ use deadpool_redis::{Connection, Pool}; use near_sdk::AccountId; use redis::{AsyncCommands, FromRedisValue, ToRedisArgs}; use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::{fmt, time::Instant}; use tokio::sync::RwLock; use tracing; @@ -16,6 +16,8 @@ pub enum StorageError { ConnectionFailed, #[error("redis operation failed: {0}")] RedisFailed(String), + #[error("ProtocolStorage::set_me() was not called")] + NotInitialized, } #[derive(Debug, Clone)] @@ -162,6 +164,7 @@ pub struct ProtocolStorage { reserved: Arc>>, owner_keys: String, account_id: AccountId, + me: Arc>, _phantom: std::marker::PhantomData, } @@ -173,6 +176,7 @@ impl Clone for ProtocolStorage { reserved: self.reserved.clone(), owner_keys: self.owner_keys.clone(), account_id: self.account_id.clone(), + me: self.me.clone(), _phantom: std::marker::PhantomData, } } @@ -190,12 +194,23 @@ impl ProtocolStorage { reserved: state, owner_keys, account_id: account_id.clone(), + me: Arc::new(OnceLock::new()), _phantom: std::marker::PhantomData, } } } impl ProtocolStorage { + /// Set this node's participant identity. Must be called before using + /// methods that depend on ownership + pub fn set_me(&self, me: Participant) { + let _ = self.me.set(me); + } + + fn me(&self) -> Result { + self.me.get().copied().ok_or(StorageError::NotInitialized) + } + async fn connect(&self) -> Option { self.redis_pool .get() @@ -206,13 +221,17 @@ impl ProtocolStorage { .ok() } - pub async fn fetch_owned(&self, me: Participant) -> Result, StorageError> { + pub async fn fetch_owned(&self) -> Result, StorageError> { + self.fetch_owned_by(self.me()?).await + } + + pub async fn fetch_owned_by(&self, owner: Participant) -> Result, StorageError> { let Some(mut conn) = self.connect().await else { return Err(StorageError::ConnectionFailed); }; let owned: HashSet = conn - .smembers(owner_key(&self.owner_keys, me)) + .smembers(owner_key(&self.owner_keys, owner)) .await .map_err(|err| { tracing::warn!(?err, "failed to fetch my owned artifacts"); @@ -224,9 +243,8 @@ impl ProtocolStorage { /// Create a slot for generating an artifact with the given ID. /// Tracks the ID in the `generating` set until the slot is inserted or dropped. - /// `mine` indicates whether this node is the owner/proposer of the artifact. /// Returns `None` if the ID is already generating, in use, or stored in Redis. - pub async fn create_slot(&self, id: A::Id, mine: bool) -> Option> { + pub async fn create_slot(&self, id: A::Id, owner: Participant) -> Option> { { let mut state = self.reserved.write().await; if state.using.contains_key(&id) { @@ -240,7 +258,8 @@ impl ProtocolStorage { ); return None; } - state.generating.insert(id, mine); + let me = self.me().ok()?; + state.generating.insert(id, owner == me); } if self.contains(id).await { self.reserved.write().await.generating.remove(&id); @@ -261,11 +280,8 @@ impl ProtocolStorage { /// Owned artifacts in Redis plus owned using and owned generating. /// This is the full set that should be advertised during state sync to prevent /// peers from pruning artifacts that are still actively in use. - pub async fn fetch_owned_with_reserved( - &self, - me: Participant, - ) -> Result, StorageError> { - let mut ids = self.fetch_owned(me).await?; + pub async fn fetch_owned_with_reserved(&self) -> Result, StorageError> { + let mut ids = self.fetch_owned().await?; let state = self.reserved.read().await; ids.extend( state @@ -483,12 +499,7 @@ impl ProtocolStorage { self.reserved.read().await.using.keys().copied().collect() } - pub async fn take( - &self, - id: A::Id, - owner: Participant, - mine: bool, - ) -> Option> { + pub async fn take(&self, id: A::Id, owner: Participant) -> Option> { const SCRIPT: &str = r#" local artifact_key = KEYS[1] local owner_key = KEYS[2] @@ -513,6 +524,7 @@ impl ProtocolStorage { "#; let start = Instant::now(); + let mine = owner == self.me().ok()?; if self.reserved.write().await.using.insert(id, mine).is_some() { tracing::warn!(id, "taking artifact that is already in use"); return None; @@ -635,7 +647,7 @@ impl ProtocolStorage { /// Take one artifact owned by the given participant. /// It is very important to NOT reuse the same artifact twice for two different /// protocols. - pub async fn take_mine(&self, me: Participant) -> Option> { + pub async fn take_mine(&self) -> Option> { const SCRIPT: &str = r#" local artifact_key = KEYS[1] local mine_key = KEYS[2] @@ -667,6 +679,7 @@ impl ProtocolStorage { let start = Instant::now(); let mut conn = self.connect().await?; + let me = self.me().ok()?; let result: Result)>, _> = redis::Script::new(SCRIPT) .key(&self.artifact_key) .key(owner_key(&self.owner_keys, me)) @@ -702,11 +715,10 @@ impl ProtocolStorage { /// Batch remove a peer from holders for a set of artifact IDs, and prune /// artifacts that fall below the holder threshold. - /// Assumes the given IDs are owned by `me` for ownership-set cleanup. + /// Assumes the given IDs are owned by this node for ownership-set cleanup. /// Returns (Vec, Vec) pub async fn remove_holder_and_prune( &self, - me: Participant, peer: Participant, threshold: usize, ids: &[A::Id], @@ -752,7 +764,7 @@ impl ProtocolStorage { type SyncResult = Result<(Vec, Vec), redis::RedisError>; let result: SyncResult = redis::Script::new(SCRIPT) .key(&self.artifact_key) - .key(owner_key(&self.owner_keys, me)) + .key(owner_key(&self.owner_keys, self.me()?)) .arg(Into::::into(peer)) .arg(threshold as i64) .arg(ids) diff --git a/integration-tests/benches/store.rs b/integration-tests/benches/store.rs index 27a4543e7..311b48a90 100644 --- a/integration-tests/benches/store.rs +++ b/integration-tests/benches/store.rs @@ -87,8 +87,8 @@ fn env() -> (Runtime, SyncEnv) { .unwrap(); let redis = spawner.spawn_redis().await; - let triples = redis.triple_storage(&node_id); - let presignatures = redis.presignature_storage(&node_id); + let triples = redis.triple_storage(&node_id, me); + let presignatures = redis.presignature_storage(&node_id, me); { let participants = into_contract_participants(&participants); redis @@ -158,7 +158,7 @@ fn bench_load_keys(c: &mut Criterion) { for i in 0..1000 { let t = dummy_pair(i); env.triples - .create_slot(t.id, true) + .create_slot(t.id, env.me) .await .unwrap() .insert(t, env.me) @@ -174,7 +174,7 @@ fn bench_load_keys(c: &mut Criterion) { for i in 0..1000 { let p = dummy_presignature(i); env.presignatures - .create_slot(p.id, true) + .create_slot(p.id, env.me) .await .unwrap() .insert(p, env.me) @@ -187,7 +187,7 @@ fn bench_load_keys(c: &mut Criterion) { c.bench_function("load 1024 mine triple keys", |b| { b.iter(|| { let task = || async { - let _ = env.triples.fetch_owned(env.me).await; + let _ = env.triples.fetch_owned().await; }; rt.block_on(task()); @@ -197,7 +197,7 @@ fn bench_load_keys(c: &mut Criterion) { c.bench_function("load 1024 mine presignature keys", |b| { b.iter(|| { let task = || async { - let _ = env.presignatures.fetch_owned(env.me).await; + let _ = env.presignatures.fetch_owned().await; }; rt.block_on(task()); diff --git a/integration-tests/src/cluster/mod.rs b/integration-tests/src/cluster/mod.rs index 21a45a1b7..489e8889c 100644 --- a/integration-tests/src/cluster/mod.rs +++ b/integration-tests/src/cluster/mod.rs @@ -3,7 +3,6 @@ pub mod spawner; use std::collections::{HashMap, HashSet}; use mpc_contract::primitives::Participants; -use mpc_node::storage::{PresignatureStorage, TripleStorage}; use mpc_primitives::{Chain, Checkpoint}; use near_workspaces::network::Sandbox; use near_workspaces::types::{Finality, NearToken}; @@ -95,17 +94,6 @@ impl Cluster { self.nodes.contract() } - pub fn triples(&self, idx: usize) -> TripleStorage { - self.nodes.ctx().redis.triple_storage(self.account_id(idx)) - } - - pub fn presignatures(&self, idx: usize) -> PresignatureStorage { - self.nodes - .ctx() - .redis - .presignature_storage(self.account_id(idx)) - } - pub async fn contract_state(&self) -> anyhow::Result { let state: ProtocolContractState = self .contract() diff --git a/integration-tests/src/containers.rs b/integration-tests/src/containers.rs index 9df54a663..5cc9ecc91 100644 --- a/integration-tests/src/containers.rs +++ b/integration-tests/src/containers.rs @@ -374,12 +374,24 @@ impl Redis { .unwrap() } - pub fn triple_storage(&self, id: &AccountId) -> mpc_node::storage::TripleStorage { - TriplePair::storage(&self.pool(), id) + pub fn triple_storage( + &self, + id: &AccountId, + me: Participant, + ) -> mpc_node::storage::TripleStorage { + let storage = TriplePair::storage(&self.pool(), id); + storage.set_me(me); + storage } - pub fn presignature_storage(&self, id: &AccountId) -> mpc_node::storage::PresignatureStorage { - Presignature::storage(&self.pool(), id) + pub fn presignature_storage( + &self, + id: &AccountId, + me: Participant, + ) -> mpc_node::storage::PresignatureStorage { + let storage = Presignature::storage(&self.pool(), id); + storage.set_me(me); + storage } pub async fn stockpile_triples(&self, cfg: &NodeConfig, participants: &Participants, mul: u32) { @@ -388,15 +400,15 @@ impl Redis { .participants .keys() .map(|account_id| { - ( - Participant::from( - *participants - .account_to_participant_id - .get(account_id) - .unwrap(), - ), - TriplePair::storage(&pool, account_id), - ) + let me = Participant::from( + *participants + .account_to_participant_id + .get(account_id) + .unwrap(), + ); + let storage = TriplePair::storage(&pool, account_id); + storage.set_me(me); + (me, storage) }) .collect::>(); @@ -429,7 +441,7 @@ impl Redis { storage .get(me) .unwrap() - .create_slot(pair_id, true) + .create_slot(pair_id, *owner) .await .unwrap() .insert(pair, *owner) diff --git a/integration-tests/src/lib.rs b/integration-tests/src/lib.rs index 4b8efce2a..45ce23baf 100644 --- a/integration-tests/src/lib.rs +++ b/integration-tests/src/lib.rs @@ -16,7 +16,6 @@ use crate::containers::DockerClient; use anyhow::Context as _; use cluster::spawner::ClusterSpawner; -use deadpool_redis::Pool; use ethers::types::{Address, U256}; use mpc_contract::config::{PresignatureConfig, ProtocolConfig, TripleConfig}; use mpc_contract::primitives::CandidateInfo; @@ -24,7 +23,6 @@ use mpc_node::gcp::GcpService; use mpc_node::indexer_eth::EthConfig; use mpc_node::indexer_hydration::HydrationConfig; use mpc_node::indexer_sol::SolConfig; -use mpc_node::storage::triple_storage::{TriplePair, TripleStorage}; use mpc_node::{logs, mesh, node_client, storage}; use mpc_primitives::{Chain, Checkpoint}; use near_workspaces::network::Sandbox; @@ -238,10 +236,6 @@ impl Nodes { Ok(()) } - pub async fn triple_storage(&self, redis_pool: &Pool, account_id: &AccountId) -> TripleStorage { - TriplePair::storage(redis_pool, account_id) - } - pub async fn gcp_services(&self) -> anyhow::Result> { let mut gcp_services = Vec::new(); match self { diff --git a/integration-tests/src/mpc_fixture/builder.rs b/integration-tests/src/mpc_fixture/builder.rs index 0325d58b3..d535b33b8 100644 --- a/integration-tests/src/mpc_fixture/builder.rs +++ b/integration-tests/src/mpc_fixture/builder.rs @@ -570,6 +570,7 @@ impl MpcFixtureNodeBuilder { let triple_storage = TriplePair::storage(&context.redis_pool, &self.participant_info.account_id); + triple_storage.set_me(self.me); if fixture_config.use_preshared_triples { // removing here because we can't clone a triple @@ -580,7 +581,7 @@ impl MpcFixtureNodeBuilder { if pair.holders.is_none() { pair.holders = Some(pair.triple0.public.participants.clone()); } - let mut slot = triple_storage.create_slot(pair_id, true).await.unwrap(); + let mut slot = triple_storage.create_slot(pair_id, owner).await.unwrap(); slot.insert(pair, owner).await; } } @@ -588,6 +589,7 @@ impl MpcFixtureNodeBuilder { let presignature_storage = Presignature::storage(&context.redis_pool, &self.participant_info.account_id); + presignature_storage.set_me(self.me); if fixture_config.use_preshared_presignatures { // removing here because we can't clone a presignature @@ -598,7 +600,7 @@ impl MpcFixtureNodeBuilder { presignature_share.holders = Some(presignature_share.participants.clone()); } let mut slot = presignature_storage - .create_slot(presignature_share.id, true) + .create_slot(presignature_share.id, owner) .await .unwrap(); slot.insert(presignature_share, owner).await; diff --git a/integration-tests/src/mpc_fixture/fixture_interface.rs b/integration-tests/src/mpc_fixture/fixture_interface.rs index dd95b4b79..28acba891 100644 --- a/integration-tests/src/mpc_fixture/fixture_interface.rs +++ b/integration-tests/src/mpc_fixture/fixture_interface.rs @@ -183,18 +183,14 @@ impl MpcFixtureNode { /// Get the list of triple IDs this node owns in storage (sorted). pub async fn owned_triples(&self) -> Vec { - let mut ids = self.triple_storage.fetch_owned(self.me).await.unwrap(); + let mut ids = self.triple_storage.fetch_owned().await.unwrap(); ids.sort(); ids } /// Get the list of presignature IDs this node owns in storage (sorted). pub async fn owned_presignatures(&self) -> Vec { - let mut ids = self - .presignature_storage - .fetch_owned(self.me) - .await - .unwrap(); + let mut ids = self.presignature_storage.fetch_owned().await.unwrap(); ids.sort(); ids } @@ -203,7 +199,7 @@ impl MpcFixtureNode { pub async fn owned_triples_with_reserved(&self) -> Vec { let mut ids = self .triple_storage - .fetch_owned_with_reserved(self.me) + .fetch_owned_with_reserved() .await .unwrap(); ids.sort(); @@ -214,7 +210,7 @@ impl MpcFixtureNode { pub async fn owned_presignatures_with_reserved(&self) -> Vec { let mut ids = self .presignature_storage - .fetch_owned_with_reserved(self.me) + .fetch_owned_with_reserved() .await .unwrap(); ids.sort(); @@ -230,11 +226,11 @@ impl MpcFixtureNode { response: &mpc_node::protocol::sync::SyncUpdate, ) { self.triple_storage - .remove_holder_and_prune(self.me, peer, threshold, &response.triples) + .remove_holder_and_prune(peer, threshold, &response.triples) .await .expect("remove_holder_and_prune triples failed"); self.presignature_storage - .remove_holder_and_prune(self.me, peer, threshold, &response.presignatures) + .remove_holder_and_prune(peer, threshold, &response.presignatures) .await .expect("remove_holder_and_prune presignatures failed"); } diff --git a/integration-tests/tests/cases/helpers.rs b/integration-tests/tests/cases/helpers.rs index b0a0fc238..0e6f2c9ea 100644 --- a/integration-tests/tests/cases/helpers.rs +++ b/integration-tests/tests/cases/helpers.rs @@ -67,7 +67,7 @@ pub(crate) async fn insert_triples_for_owner( let holders = holders.to_vec(); for id in ids { triples - .create_slot(id, true) // mine flag is not important, since we are inserting right away + .create_slot(id, owner) .await .unwrap() .insert(dummy_pair_with_holders(id, holders.clone()), owner) @@ -84,7 +84,7 @@ pub(crate) async fn insert_presignatures_for_owner( let holders = holders.to_vec(); for id in ids { presignatures - .create_slot(id, true) // mine flag is not important, since we are inserting right away + .create_slot(id, owner) .await .unwrap() .insert(dummy_presignature_with_holders(id, holders.clone()), owner) diff --git a/integration-tests/tests/cases/mpc.rs b/integration-tests/tests/cases/mpc.rs index 3b0e6783e..352ccc477 100644 --- a/integration-tests/tests/cases/mpc.rs +++ b/integration-tests/tests/cases/mpc.rs @@ -390,7 +390,13 @@ async fn test_sync_matrix() { .await; } ArtifactState::Generating => { - caller_slot = Some(caller.triple_storage.create_slot(id, true).await.unwrap()); + caller_slot = Some( + caller + .triple_storage + .create_slot(id, caller.me) + .await + .unwrap(), + ); } ArtifactState::Using => { insert_triples_for_owner( @@ -400,13 +406,7 @@ async fn test_sync_matrix() { id..=id, ) .await; - caller_taken = Some( - caller - .triple_storage - .take(id, caller.me, true) - .await - .unwrap(), - ); + caller_taken = Some(caller.triple_storage.take(id, caller.me).await.unwrap()); } ArtifactState::None => {} } @@ -426,7 +426,7 @@ async fn test_sync_matrix() { responder_slot = Some( responder .triple_storage - .create_slot(id, false) + .create_slot(id, caller.me) .await .unwrap(), ); @@ -439,13 +439,7 @@ async fn test_sync_matrix() { id..=id, ) .await; - responder_taken = Some( - responder - .triple_storage - .take(id, caller.me, false) - .await - .unwrap(), - ); + responder_taken = Some(responder.triple_storage.take(id, caller.me).await.unwrap()); } ArtifactState::None => {} } @@ -589,7 +583,7 @@ async fn test_basic_generate_triples() { for node in &network.nodes { let mut nodes_shares = BTreeMap::new(); for peer in &network.nodes { - let triple_ids = node.triple_storage.fetch_owned(peer.me).await.unwrap(); + let triple_ids = node.triple_storage.fetch_owned_by(peer.me).await.unwrap(); let mut peer_triples = Vec::with_capacity(triple_ids.len()); for triple_id in triple_ids { let pair = conn @@ -643,7 +637,7 @@ async fn test_basic_generate_presignature() { for peer in &network.nodes { let presignature_ids = node .presignature_storage - .fetch_owned(peer.me) + .fetch_owned_by(peer.me) .await .unwrap(); let mut peer_presignatures = Vec::with_capacity(presignature_ids.len()); diff --git a/integration-tests/tests/cases/store.rs b/integration-tests/tests/cases/store.rs index abbad1e51..67f36e8d0 100644 --- a/integration-tests/tests/cases/store.rs +++ b/integration-tests/tests/cases/store.rs @@ -22,7 +22,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { let (_, _, msg) = MessageChannel::new(); let node0_id = "party0.near".parse().unwrap(); let redis = containers::Redis::run(&spawner).await; - let triple_storage = redis.triple_storage(&node0_id); + let triple_storage = redis.triple_storage(&node0_id, node0); let triple_spawner = TripleSpawner::new(node0, 5, 123, &triple_storage, msg, node0_id.to_string()); @@ -38,13 +38,13 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 0); triple_storage - .create_slot(triple_id1, false) + .create_slot(triple_id1, node1) .await .unwrap() .insert(dummy_pair(triple_id1), node1) .await; triple_storage - .create_slot(triple_id2, false) + .create_slot(triple_id2, node1) .await .unwrap() .insert(dummy_pair(triple_id2), node1) @@ -60,8 +60,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 2); // Take triple pairs and check that they are removed from the storage and marked as using - let _taken1 = triple_storage.take(triple_id1, node1, false).await.unwrap(); - let _taken2 = triple_storage.take(triple_id2, node1, false).await.unwrap(); + let _taken1 = triple_storage.take(triple_id1, node1).await.unwrap(); + let _taken2 = triple_storage.take(triple_id2, node1).await.unwrap(); assert!(!triple_spawner.contains(triple_id1).await); assert!(!triple_spawner.contains(triple_id2).await); assert!(!triple_spawner.contains_mine(triple_id1).await); @@ -74,11 +74,11 @@ async fn test_triple_persistence() -> anyhow::Result<()> { // Attempt to re-create slot for in-use triples and check that it fails assert!(triple_storage - .create_slot(triple_id1, false) + .create_slot(triple_id1, node1) .await .is_none()); assert!(triple_storage - .create_slot(triple_id2, false) + .create_slot(triple_id2, node1) .await .is_none()); @@ -87,13 +87,13 @@ async fn test_triple_persistence() -> anyhow::Result<()> { // Add mine triple and check that it is in the storage triple_storage - .create_slot(id3, true) + .create_slot(id3, node0) .await .unwrap() .insert(dummy_pair(id3), node0) .await; triple_storage - .create_slot(id4, true) + .create_slot(id4, node0) .await .unwrap() .insert(dummy_pair(id4), node0) @@ -107,8 +107,8 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert_eq!(triple_spawner.len_potential().await, 2); // Take mine triple pairs and check that they are removed from the storage and marked as using - let _taken3 = triple_storage.take_mine(node0).await.unwrap(); - let _taken4 = triple_storage.take_mine(node0).await.unwrap(); + let _taken3 = triple_storage.take_mine().await.unwrap(); + let _taken4 = triple_storage.take_mine().await.unwrap(); assert!(!triple_spawner.contains(id3).await); assert!(!triple_spawner.contains(id4).await); assert!(!triple_spawner.contains_mine(id3).await); @@ -121,14 +121,14 @@ async fn test_triple_persistence() -> anyhow::Result<()> { assert!(triple_storage.contains_using(id4).await); // Attempt to re-create slot for in-use mine triples and check that it fails - assert!(triple_storage.create_slot(id3, true).await.is_none()); - assert!(triple_storage.create_slot(id4, true).await.is_none()); + assert!(triple_storage.create_slot(id3, node0).await.is_none()); + assert!(triple_storage.create_slot(id4, node0).await.is_none()); assert!(triple_storage.clear().await); // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { triple_storage - .create_slot(id, false) + .create_slot(id, node1) .await .unwrap() .insert(dummy_pair(id), node1) @@ -138,7 +138,7 @@ async fn test_triple_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { triple_storage - .create_slot(id, true) + .create_slot(id, node0) .await .unwrap() .insert(dummy_pair(id), node0) @@ -175,8 +175,8 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { let (_, _, msg) = MessageChannel::new(); let node0_id = "party0.near".parse().unwrap(); let redis = containers::Redis::run(&spawner).await; - let triple_storage = redis.triple_storage(&node0_id); - let presignature_storage = redis.presignature_storage(&node0_id); + let triple_storage = redis.triple_storage(&node0_id, node0); + let presignature_storage = redis.presignature_storage(&node0_id, node0); let presignature_spawner = PresignatureSpawner::new( Participant::from(0), 5, @@ -203,7 +203,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Insert presignature owned by node1, with our node0 view being that it is a foreign presignature assert!( presignature_storage - .create_slot(presignature.id, false) + .create_slot(presignature.id, node1) .await .unwrap() .insert(presignature, node1) @@ -218,7 +218,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 1); // Take presignature and check that it is removed from the storage and marked as using - let _taken_ps1 = presignature_storage.take(id, node1, false).await.unwrap(); + let _taken_ps1 = presignature_storage.take(id, node1).await.unwrap(); assert!(!presignature_storage.contains(id).await); assert!(!presignature_spawner.contains_mine(id).await); assert_eq!(presignature_storage.len_generated().await, 0); @@ -227,7 +227,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert!(presignature_storage.contains_using(id).await); // Attempt to re-create slot for in-use presignature and check that it fails - assert!(presignature_storage.create_slot(id, false).await.is_none()); + assert!(presignature_storage.create_slot(id, node1).await.is_none()); let id2 = 2; let mine_presignature = dummy_presignature(id2); @@ -235,7 +235,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Add a presignature to our own node0 assert!( presignature_storage - .create_slot(id2, true) + .create_slot(id2, node0) .await .unwrap() .insert(mine_presignature, node0) @@ -249,7 +249,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert_eq!(presignature_spawner.len_potential().await, 1); // Take mine presignature and check that it is removed from the storage and marked as using - let _taken_ps2 = presignature_storage.take_mine(node0).await.unwrap(); + let _taken_ps2 = presignature_storage.take_mine().await.unwrap(); assert!(!presignature_storage.contains(id2).await); assert!(!presignature_spawner.contains_mine(id2).await); assert_eq!(presignature_storage.len_generated().await, 0); @@ -259,13 +259,13 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { assert!(presignature_storage.contains_using(id2).await); // Attempt to re-create slot for in-use mine presignature and check that it fails - assert!(presignature_storage.create_slot(id2, true).await.is_none()); + assert!(presignature_storage.create_slot(id2, node0).await.is_none()); presignature_storage.clear().await; // Have our node0 observe shares for triples 10 to 15 where node1 is owner. for id in 10..=15 { presignature_storage - .create_slot(id, false) + .create_slot(id, node1) .await .unwrap() .insert(dummy_presignature(id), node1) @@ -275,7 +275,7 @@ async fn test_presignature_persistence() -> anyhow::Result<()> { // Have our node0 own 16 to 20 for id in 16..=20 { presignature_storage - .create_slot(id, true) + .create_slot(id, node0) .await .unwrap() .insert(dummy_presignature(id), node0) diff --git a/integration-tests/tests/cases/sync.rs b/integration-tests/tests/cases/sync.rs index 526d2607b..110e42430 100644 --- a/integration-tests/tests/cases/sync.rs +++ b/integration-tests/tests/cases/sync.rs @@ -26,10 +26,10 @@ async fn test_state_sync_e2e_large_outdated_stockpile() { let redis = spawner.prespawn_redis().await; // immediately add to triples/presignatures storage the triples/presignatures we want to invalidate. - let node0_triples = redis.triple_storage(&node0_account_id); - let node0_presignatures = redis.presignature_storage(&node0_account_id); - let node1_triples = redis.triple_storage(&node1_account_id); - let node1_presignatures = redis.presignature_storage(&node1_account_id); + let node0_triples = redis.triple_storage(&node0_account_id, node0); + let node0_presignatures = redis.presignature_storage(&node0_account_id, node0); + let node1_triples = redis.triple_storage(&node1_account_id, node1); + let node1_presignatures = redis.presignature_storage(&node1_account_id, node1); // insert triples that will be invalidated after a sync, since nobody else has them. // node0 is saying that they have 0 to 5, but node1 will sync and say they have 4 and 5 only. @@ -110,12 +110,12 @@ async fn test_state_sync_e2e() { tokio::time::sleep(Duration::from_secs(1)).await; // Get triple/presignature storage for each node. - let node0_triples = redis.triple_storage(&node0_account_id); - let node0_presignatures = redis.presignature_storage(&node0_account_id); - let node1_triples = redis.triple_storage(&node1_account_id); - let node1_presignatures = redis.presignature_storage(&node1_account_id); - let node2_triples = redis.triple_storage(&node2_account_id); - let node2_presignatures = redis.presignature_storage(&node2_account_id); + let node0_triples = redis.triple_storage(&node0_account_id, node0); + let node0_presignatures = redis.presignature_storage(&node0_account_id, node0); + let node1_triples = redis.triple_storage(&node1_account_id, node1); + let node1_presignatures = redis.presignature_storage(&node1_account_id, node1); + let node2_triples = redis.triple_storage(&node2_account_id, node2); + let node2_presignatures = redis.presignature_storage(&node2_account_id, node2); // Populate 3 triples and 3 presignatures: each node owns 1, all nodes hold shares. for storage in [&node0_triples, &node1_triples, &node2_triples] { From a48fabb9aa9b0846c898ef1f162c057bbc3b5eec Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Mon, 6 Apr 2026 22:32:20 +0300 Subject: [PATCH 19/21] move state sync tests into a separate file --- integration-tests/tests/cases/mod.rs | 1 + integration-tests/tests/cases/mpc.rs | 486 ------------------- integration-tests/tests/cases/state_sync.rs | 490 ++++++++++++++++++++ 3 files changed, 491 insertions(+), 486 deletions(-) create mode 100644 integration-tests/tests/cases/state_sync.rs diff --git a/integration-tests/tests/cases/mod.rs b/integration-tests/tests/cases/mod.rs index 441b5b77f..5588b164e 100644 --- a/integration-tests/tests/cases/mod.rs +++ b/integration-tests/tests/cases/mod.rs @@ -24,6 +24,7 @@ pub mod mpc; pub mod nightly; pub mod solana; pub mod solana_stream; +pub mod state_sync; pub mod store; pub mod sync; diff --git a/integration-tests/tests/cases/mpc.rs b/integration-tests/tests/cases/mpc.rs index 352ccc477..5a0897870 100644 --- a/integration-tests/tests/cases/mpc.rs +++ b/integration-tests/tests/cases/mpc.rs @@ -27,492 +27,6 @@ const PRESIGNATURES_FILE: &str = "tmp/presignatures.json"; const TRIPLE_PAIRS_PER_OWNER: usize = 50; const PRESIGNATURES_PER_OWNER: usize = 25; -use super::helpers::{ - assert_presig_owned_state, assert_triples_owned_state, insert_presignatures_for_owner, - insert_triples_for_owner, -}; - -#[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_noop_when_fully_synced() { - let fixture = MpcFixtureBuilder::default() - .only_generate_signatures() - .build() - .await; - - let node0 = &fixture.nodes[0]; - let node1 = &fixture.nodes[1]; - - // Snapshot node1's owned artifacts before sync. - let node1_triples_before = node1.owned_triples().await; - let node1_presigs_before = node1.owned_presignatures().await; - assert!( - !node1_triples_before.is_empty(), - "node1 should own triples from fixture" - ); - assert!( - !node1_presigs_before.is_empty(), - "node1 should own presignatures from fixture" - ); - - let node0_triples = node0.owned_triples().await; - let node0_presigs = node0.owned_presignatures().await; - - // Responder side: node1 receives node0's sync update and reports what it's missing. - let response = node1 - .sync(node0.me, node0_triples.clone(), node0_presigs.clone()) - .await; - assert!( - response.triples.is_empty(), - "node1 should have all of node0's triples" - ); - assert!( - response.presignatures.is_empty(), - "node1 should have all of node0's presignatures" - ); - - // Caller side: node0 processes the response (nothing to remove since response is empty). - node0.process_sync_response(node1.me, 2, &response).await; - - // Verify node1's artifact state is exactly unchanged after sync. - let node1_triples_after = node1.owned_triples().await; - let node1_presigs_after = node1.owned_presignatures().await; - assert_eq!( - node1_triples_after, node1_triples_before, - "node1 triples should be unchanged after sync" - ); - assert_eq!( - node1_presigs_after, node1_presigs_before, - "node1 presignatures should be unchanged after sync" - ); - - // Verify node0's artifact state is also unchanged. - let node0_triples_after = node0.owned_triples().await; - let node0_presigs_after = node0.owned_presignatures().await; - assert_eq!( - node0_triples_after, node0_triples, - "node0 triples should be unchanged after sync" - ); - assert_eq!( - node0_presigs_after, node0_presigs, - "node0 presignatures should be unchanged after sync" - ); - - // Verify holders are full (all 3 nodes) for every artifact on both nodes. - let all_participants = fixture.sorted_participants(); - for node in [node0, node1] { - for id in &node.owned_triples().await { - let holders = node.triple_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "triple {id} on {:?} should have all participants as holders", - node.me - ); - } - for id in &node.owned_presignatures().await { - let holders = node.presignature_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "presignature {id} on {:?} should have all participants as holders", - node.me - ); - } - } -} - -#[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_prune_below_threshold() { - let fixture = MpcFixtureBuilder::default() - .only_generate_signatures() - .build() - .await; - - let node0 = &fixture.nodes[0]; - let node1 = &fixture.nodes[1]; - let node2 = &fixture.nodes[2]; - let all_participants = fixture.sorted_participants(); - - // Snapshot fixture artifacts before any mutations. - let node0_triples_before = node0.owned_triples().await; - let node0_presigs_before = node0.owned_presignatures().await; - - // Insert artifact id=99 only on node0's storage (claimed by all 3 holders). - insert_triples_for_owner(&node0.triple_storage, node0.me, &all_participants, 99..=99).await; - insert_presignatures_for_owner( - &node0.presignature_storage, - node0.me, - &all_participants, - 99..=99, - ) - .await; - - // node0 tells node1 "I own id=99", node1 responds "I don't have it". - let response = node1.sync(node0.me, vec![99], vec![99]).await; - assert_eq!(response.triples, vec![99]); - assert_eq!(response.presignatures, vec![99]); - - // node2 also doesn't have it. - let response2 = node2.sync(node0.me, vec![99], vec![99]).await; - assert_eq!(response2.triples, vec![99]); - assert_eq!(response2.presignatures, vec![99]); - - // Process node1's response first (removes node1, 2 holders remain = threshold, survives). - node0.process_sync_response(node1.me, 2, &response).await; - assert_triples_owned_state(&node0.triple_storage, node0.me, &[99], &[]).await; - assert_presig_owned_state(&node0.presignature_storage, node0.me, &[99], &[]).await; - - // Verify holders of 99: node1 removed, only node0 + node2 remain. - let mut expected_holders = vec![node0.me, node2.me]; - expected_holders.sort(); - let holders_99 = node0.triple_storage.fetch_holders(99).await; - assert_eq!( - holders_99, expected_holders, - "triple 99 should have node0+node2 as holders after first sync" - ); - let holders_99 = node0.presignature_storage.fetch_holders(99).await; - assert_eq!( - holders_99, expected_holders, - "presig 99 should have node0+node2 as holders after first sync" - ); - - // Process node2's response (removes node2, 1 holder < threshold → pruned). - node0.process_sync_response(node2.me, 2, &response2).await; - assert_triples_owned_state(&node0.triple_storage, node0.me, &[], &[99]).await; - assert_presig_owned_state(&node0.presignature_storage, node0.me, &[], &[99]).await; - - // Verify fixture artifacts on node0 are unchanged (only id=99 was pruned). - let node0_triples_after = node0.owned_triples().await; - let node0_presigs_after = node0.owned_presignatures().await; - assert_eq!( - node0_triples_after, node0_triples_before, - "node0 fixture triples should be unchanged" - ); - assert_eq!( - node0_presigs_after, node0_presigs_before, - "node0 fixture presignatures should be unchanged" - ); - - // Verify holders of fixture artifacts are still full on node0. - for id in &node0_triples_after { - let holders = node0.triple_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "fixture triple {id} should still have all holders" - ); - } - for id in &node0_presigs_after { - let holders = node0.presignature_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "fixture presig {id} should still have all holders" - ); - } -} - -/// Orphaned artifact: owner doesn't have id=77 but other nodes do. -/// When owner broadcasts its sync update (without id=77), responders remove it -/// via remove_outdated. -#[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_remove_outdated_orphan() { - let fixture = MpcFixtureBuilder::default() - .only_generate_signatures() - .build() - .await; - - let node0 = &fixture.nodes[0]; - let node1 = &fixture.nodes[1]; - let all_participants = fixture.sorted_participants(); - - // Snapshot fixture artifacts on node1 before mutations. - let node1_triples_before = node1.owned_triples().await; - let node1_presigs_before = node1.owned_presignatures().await; - - // Insert id=77 owned by node0, but only on node1 (NOT on node0). - insert_triples_for_owner(&node1.triple_storage, node0.me, &all_participants, 77..=77).await; - insert_presignatures_for_owner( - &node1.presignature_storage, - node0.me, - &all_participants, - 77..=77, - ) - .await; - - // Verify node1 has id=77 before sync. - assert_triples_owned_state(&node1.triple_storage, node0.me, &[77], &[]).await; - assert_presig_owned_state(&node1.presignature_storage, node0.me, &[77], &[]).await; - - // Verify holders of id=77 are all participants. - let holders = node1.triple_storage.fetch_holders(77).await; - assert_eq!( - holders, all_participants, - "triple 77 should have all participants as holders before sync" - ); - let holders = node1.presignature_storage.fetch_holders(77).await; - assert_eq!( - holders, all_participants, - "presig 77 should have all participants as holders before sync" - ); - - // node0 broadcasts its sync update (which does NOT include id=77 since node0 doesn't have it). - // Responder runs remove_outdated, which removes id=77. - let node0_triples = node0.owned_triples().await; - let node0_presigs = node0.owned_presignatures().await; - assert!(!node0_triples.contains(&77), "node0 should not own id=77"); - - let response = node1 - .sync(node0.me, node0_triples.clone(), node0_presigs.clone()) - .await; - assert!( - response.triples.is_empty(), - "node1 should not report any missing triples (remove_outdated handles orphans)" - ); - assert!( - response.presignatures.is_empty(), - "node1 should not report any missing presignatures" - ); - - // After sync, id=77 should be removed from node1 (via remove_outdated). - assert_triples_owned_state(&node1.triple_storage, node0.me, &[], &[77]).await; - assert_presig_owned_state(&node1.presignature_storage, node0.me, &[], &[77]).await; - - // Verify fixture artifacts on node1 are unchanged (only id=77 was removed). - let node1_triples_after = node1.owned_triples().await; - let node1_presigs_after = node1.owned_presignatures().await; - assert_eq!( - node1_triples_after, node1_triples_before, - "node1 fixture triples should be unchanged" - ); - assert_eq!( - node1_presigs_after, node1_presigs_before, - "node1 fixture presignatures should be unchanged" - ); - - // Verify holders of fixture artifacts are still full on node1. - for id in &node1_triples_after { - let holders = node1.triple_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "fixture triple {id} should still have all holders" - ); - } - for id in &node1_presigs_after { - let holders = node1.presignature_storage.fetch_holders(*id).await; - assert_eq!( - holders, all_participants, - "fixture presig {id} should still have all holders" - ); - } -} - -#[test(tokio::test(flavor = "multi_thread"))] -async fn test_sync_matrix() { - #[derive(Debug, Clone, Copy)] - enum ArtifactState { - Generating, - Stored, - Using, - None, - } - - struct ExpectedCallerState { - /// Should the artifact appear in the caller's sync update? - in_update: bool, - /// Should the caller still have it in Redis after process_sync_response? - stored_after: bool, - } - - struct ExpectedResponderState { - /// Should the responder report the artifact as not_found (missing)? - missing: bool, - /// Should the responder still have it in Redis after sync? - stored_after: bool, - } - - struct Case { - caller: ArtifactState, - responder: ArtifactState, - expected_caller: ExpectedCallerState, - expected_responder: ExpectedResponderState, - } - - // Caller is always the owner (mine=true), responder is the peer (mine=false). - #[rustfmt::skip] // cargo fmt makes the matrix unreadable - let test_matrix = [ - // Caller: Stored (in update) ────────────────────────────── - Case { caller: ArtifactState::Stored, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, - Case { caller: ArtifactState::Stored, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Stored, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Stored, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, - // Caller: Generating (in update) ────────────────────────── - Case { caller: ArtifactState::Generating, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, - Case { caller: ArtifactState::Generating, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Generating, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Generating, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, - // Caller: Using (in update) ────────────────────────────── - Case { caller: ArtifactState::Using, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, - Case { caller: ArtifactState::Using, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Using, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::Using, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, - // Caller: None (NOT in update) ──────────────────────────── - Case { caller: ArtifactState::None, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::None, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::None, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - Case { caller: ArtifactState::None, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, - ]; - - let fixture = MpcFixtureBuilder::default() - .only_generate_signatures() - .build() - .await; - - let caller = &fixture.nodes[0]; - let responder = &fixture.nodes[1]; - let all_participants = fixture.sorted_participants(); - - for (i, case) in test_matrix.iter().enumerate() { - let id: u64 = 800 + i as u64; - tracing::info!(id, ?case.caller, ?case.responder, "=== case {i} ==="); - - // Hold slots/taken artifacts alive until assertions are done. - let mut caller_slot = None; - let mut caller_taken = None; - let mut responder_slot = None; - let mut responder_taken = None; - - // --- Set up caller (owner) state --- - match case.caller { - ArtifactState::Stored => { - insert_triples_for_owner( - &caller.triple_storage, - caller.me, - &all_participants, - id..=id, - ) - .await; - } - ArtifactState::Generating => { - caller_slot = Some( - caller - .triple_storage - .create_slot(id, caller.me) - .await - .unwrap(), - ); - } - ArtifactState::Using => { - insert_triples_for_owner( - &caller.triple_storage, - caller.me, - &all_participants, - id..=id, - ) - .await; - caller_taken = Some(caller.triple_storage.take(id, caller.me).await.unwrap()); - } - ArtifactState::None => {} - } - - // --- Set up responder (peer) state --- - match case.responder { - ArtifactState::Stored => { - insert_triples_for_owner( - &responder.triple_storage, - caller.me, - &all_participants, - id..=id, - ) - .await; - } - ArtifactState::Generating => { - responder_slot = Some( - responder - .triple_storage - .create_slot(id, caller.me) - .await - .unwrap(), - ); - } - ArtifactState::Using => { - insert_triples_for_owner( - &responder.triple_storage, - caller.me, - &all_participants, - id..=id, - ) - .await; - responder_taken = Some(responder.triple_storage.take(id, caller.me).await.unwrap()); - } - ArtifactState::None => {} - } - - // --- Build caller's sync update --- - let caller_update = caller.owned_triples_with_reserved().await; - assert_eq!( - caller_update.contains(&id), - case.expected_caller.in_update, - "case {i}: caller={:?} → expected in_update={}", - case.caller, - case.expected_caller.in_update, - ); - - // --- Responder processes the sync update --- - let response = responder - .sync( - caller.me, - caller_update, - responder.owned_presignatures().await, - ) - .await; - - // Verify the full SyncUpdate response from the responder. - assert_eq!( - response.from, responder.me, - "case {i}: response.from should be the responder", - ); - assert_eq!( - response.triples.contains(&id), - case.expected_responder.missing, - "case {i}: caller={:?}, responder={:?} → expected missing={}", - case.caller, - case.responder, - case.expected_responder.missing, - ); - - // --- Verify responder's Redis state after sync --- - assert_eq!( - responder - .triple_storage - .contains_by_owner(id, caller.me) - .await, - case.expected_responder.stored_after, - "case {i}: caller={:?}, responder={:?} → expected responder stored_after={}", - case.caller, - case.responder, - case.expected_responder.stored_after, - ); - - // --- Caller processes the sync response (remove holder / prune) --- - caller - .process_sync_response(responder.me, 2, &response) - .await; - assert_eq!( - caller.triple_storage.contains_by_owner(id, caller.me).await, - case.expected_caller.stored_after, - "case {i}: caller={:?}, responder={:?} → expected caller stored_after={}", - case.caller, - case.responder, - case.expected_caller.stored_after, - ); - - // Clean up held slots/taken before next iteration. - drop(caller_slot); - drop(caller_taken); - drop(responder_slot); - drop(responder_taken); - // Wait for async Drop cleanup. - tokio::time::sleep(Duration::from_millis(50)).await; - } -} - #[test(tokio::test(flavor = "multi_thread"))] async fn test_basic_generate_keys() { let network = MpcFixtureBuilder::new(5, 4).build().await; diff --git a/integration-tests/tests/cases/state_sync.rs b/integration-tests/tests/cases/state_sync.rs new file mode 100644 index 000000000..8175be786 --- /dev/null +++ b/integration-tests/tests/cases/state_sync.rs @@ -0,0 +1,490 @@ +use integration_tests::mpc_fixture::MpcFixtureBuilder; +use test_log::test; + +use std::time::Duration; + +use super::helpers::{ + assert_presig_owned_state, assert_triples_owned_state, insert_presignatures_for_owner, + insert_triples_for_owner, +}; + +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_noop_when_fully_synced() { + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let node0 = &fixture.nodes[0]; + let node1 = &fixture.nodes[1]; + + // Snapshot node1's owned artifacts before sync. + let node1_triples_before = node1.owned_triples().await; + let node1_presigs_before = node1.owned_presignatures().await; + assert!( + !node1_triples_before.is_empty(), + "node1 should own triples from fixture" + ); + assert!( + !node1_presigs_before.is_empty(), + "node1 should own presignatures from fixture" + ); + + let node0_triples = node0.owned_triples().await; + let node0_presigs = node0.owned_presignatures().await; + + // Responder side: node1 receives node0's sync update and reports what it's missing. + let response = node1 + .sync(node0.me, node0_triples.clone(), node0_presigs.clone()) + .await; + assert!( + response.triples.is_empty(), + "node1 should have all of node0's triples" + ); + assert!( + response.presignatures.is_empty(), + "node1 should have all of node0's presignatures" + ); + + // Caller side: node0 processes the response (nothing to remove since response is empty). + node0.process_sync_response(node1.me, 2, &response).await; + + // Verify node1's artifact state is exactly unchanged after sync. + let node1_triples_after = node1.owned_triples().await; + let node1_presigs_after = node1.owned_presignatures().await; + assert_eq!( + node1_triples_after, node1_triples_before, + "node1 triples should be unchanged after sync" + ); + assert_eq!( + node1_presigs_after, node1_presigs_before, + "node1 presignatures should be unchanged after sync" + ); + + // Verify node0's artifact state is also unchanged. + let node0_triples_after = node0.owned_triples().await; + let node0_presigs_after = node0.owned_presignatures().await; + assert_eq!( + node0_triples_after, node0_triples, + "node0 triples should be unchanged after sync" + ); + assert_eq!( + node0_presigs_after, node0_presigs, + "node0 presignatures should be unchanged after sync" + ); + + // Verify holders are full (all 3 nodes) for every artifact on both nodes. + let all_participants = fixture.sorted_participants(); + for node in [node0, node1] { + for id in &node.owned_triples().await { + let holders = node.triple_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "triple {id} on {:?} should have all participants as holders", + node.me + ); + } + for id in &node.owned_presignatures().await { + let holders = node.presignature_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "presignature {id} on {:?} should have all participants as holders", + node.me + ); + } + } +} + +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_prune_below_threshold() { + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let node0 = &fixture.nodes[0]; + let node1 = &fixture.nodes[1]; + let node2 = &fixture.nodes[2]; + let all_participants = fixture.sorted_participants(); + + // Snapshot fixture artifacts before any mutations. + let node0_triples_before = node0.owned_triples().await; + let node0_presigs_before = node0.owned_presignatures().await; + + // Insert artifact id=99 only on node0's storage (claimed by all 3 holders). + insert_triples_for_owner(&node0.triple_storage, node0.me, &all_participants, 99..=99).await; + insert_presignatures_for_owner( + &node0.presignature_storage, + node0.me, + &all_participants, + 99..=99, + ) + .await; + + // node0 tells node1 "I own id=99", node1 responds "I don't have it". + let response = node1.sync(node0.me, vec![99], vec![99]).await; + assert_eq!(response.triples, vec![99]); + assert_eq!(response.presignatures, vec![99]); + + // node2 also doesn't have it. + let response2 = node2.sync(node0.me, vec![99], vec![99]).await; + assert_eq!(response2.triples, vec![99]); + assert_eq!(response2.presignatures, vec![99]); + + // Process node1's response first (removes node1, 2 holders remain = threshold, survives). + node0.process_sync_response(node1.me, 2, &response).await; + assert_triples_owned_state(&node0.triple_storage, node0.me, &[99], &[]).await; + assert_presig_owned_state(&node0.presignature_storage, node0.me, &[99], &[]).await; + + // Verify holders of 99: node1 removed, only node0 + node2 remain. + let mut expected_holders = vec![node0.me, node2.me]; + expected_holders.sort(); + let holders_99 = node0.triple_storage.fetch_holders(99).await; + assert_eq!( + holders_99, expected_holders, + "triple 99 should have node0+node2 as holders after first sync" + ); + let holders_99 = node0.presignature_storage.fetch_holders(99).await; + assert_eq!( + holders_99, expected_holders, + "presig 99 should have node0+node2 as holders after first sync" + ); + + // Process node2's response (removes node2, 1 holder < threshold → pruned). + node0.process_sync_response(node2.me, 2, &response2).await; + assert_triples_owned_state(&node0.triple_storage, node0.me, &[], &[99]).await; + assert_presig_owned_state(&node0.presignature_storage, node0.me, &[], &[99]).await; + + // Verify fixture artifacts on node0 are unchanged (only id=99 was pruned). + let node0_triples_after = node0.owned_triples().await; + let node0_presigs_after = node0.owned_presignatures().await; + assert_eq!( + node0_triples_after, node0_triples_before, + "node0 fixture triples should be unchanged" + ); + assert_eq!( + node0_presigs_after, node0_presigs_before, + "node0 fixture presignatures should be unchanged" + ); + + // Verify holders of fixture artifacts are still full on node0. + for id in &node0_triples_after { + let holders = node0.triple_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "fixture triple {id} should still have all holders" + ); + } + for id in &node0_presigs_after { + let holders = node0.presignature_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "fixture presig {id} should still have all holders" + ); + } +} + +/// Orphaned artifact: owner doesn't have id=77 but other nodes do. +/// When owner broadcasts its sync update (without id=77), responders remove it +/// via remove_outdated. +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_remove_outdated_orphan() { + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let node0 = &fixture.nodes[0]; + let node1 = &fixture.nodes[1]; + let all_participants = fixture.sorted_participants(); + + // Snapshot fixture artifacts on node1 before mutations. + let node1_triples_before = node1.owned_triples().await; + let node1_presigs_before = node1.owned_presignatures().await; + + // Insert id=77 owned by node0, but only on node1 (NOT on node0). + insert_triples_for_owner(&node1.triple_storage, node0.me, &all_participants, 77..=77).await; + insert_presignatures_for_owner( + &node1.presignature_storage, + node0.me, + &all_participants, + 77..=77, + ) + .await; + + // Verify node1 has id=77 before sync. + assert_triples_owned_state(&node1.triple_storage, node0.me, &[77], &[]).await; + assert_presig_owned_state(&node1.presignature_storage, node0.me, &[77], &[]).await; + + // Verify holders of id=77 are all participants. + let holders = node1.triple_storage.fetch_holders(77).await; + assert_eq!( + holders, all_participants, + "triple 77 should have all participants as holders before sync" + ); + let holders = node1.presignature_storage.fetch_holders(77).await; + assert_eq!( + holders, all_participants, + "presig 77 should have all participants as holders before sync" + ); + + // node0 broadcasts its sync update (which does NOT include id=77 since node0 doesn't have it). + // Responder runs remove_outdated, which removes id=77. + let node0_triples = node0.owned_triples().await; + let node0_presigs = node0.owned_presignatures().await; + assert!(!node0_triples.contains(&77), "node0 should not own id=77"); + + let response = node1 + .sync(node0.me, node0_triples.clone(), node0_presigs.clone()) + .await; + assert!( + response.triples.is_empty(), + "node1 should not report any missing triples (remove_outdated handles orphans)" + ); + assert!( + response.presignatures.is_empty(), + "node1 should not report any missing presignatures" + ); + + // After sync, id=77 should be removed from node1 (via remove_outdated). + assert_triples_owned_state(&node1.triple_storage, node0.me, &[], &[77]).await; + assert_presig_owned_state(&node1.presignature_storage, node0.me, &[], &[77]).await; + + // Verify fixture artifacts on node1 are unchanged (only id=77 was removed). + let node1_triples_after = node1.owned_triples().await; + let node1_presigs_after = node1.owned_presignatures().await; + assert_eq!( + node1_triples_after, node1_triples_before, + "node1 fixture triples should be unchanged" + ); + assert_eq!( + node1_presigs_after, node1_presigs_before, + "node1 fixture presignatures should be unchanged" + ); + + // Verify holders of fixture artifacts are still full on node1. + for id in &node1_triples_after { + let holders = node1.triple_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "fixture triple {id} should still have all holders" + ); + } + for id in &node1_presigs_after { + let holders = node1.presignature_storage.fetch_holders(*id).await; + assert_eq!( + holders, all_participants, + "fixture presig {id} should still have all holders" + ); + } +} + +#[test(tokio::test(flavor = "multi_thread"))] +async fn test_sync_matrix() { + #[derive(Debug, Clone, Copy)] + enum ArtifactState { + Generating, + Stored, + Using, + None, + } + + struct ExpectedCallerState { + /// Should the artifact appear in the caller's sync update? + in_update: bool, + /// Should the caller still have it in Redis after process_sync_response? + stored_after: bool, + } + + struct ExpectedResponderState { + /// Should the responder report the artifact as not_found (missing)? + missing: bool, + /// Should the responder still have it in Redis after sync? + stored_after: bool, + } + + struct Case { + caller: ArtifactState, + responder: ArtifactState, + expected_caller: ExpectedCallerState, + expected_responder: ExpectedResponderState, + } + + // Caller is always the owner (mine=true), responder is the peer (mine=false). + #[rustfmt::skip] // cargo fmt makes the matrix unreadable + let test_matrix = [ + // Caller: Stored (in update) ────────────────────────────── + Case { caller: ArtifactState::Stored, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Stored, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: true }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: Generating (in update) ────────────────────────── + Case { caller: ArtifactState::Generating, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Generating, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: Using (in update) ────────────────────────────── + Case { caller: ArtifactState::Using, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: true } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::Using, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: true, stored_after: false }, expected_responder: ExpectedResponderState { missing: true, stored_after: false } }, + // Caller: None (NOT in update) ──────────────────────────── + Case { caller: ArtifactState::None, responder: ArtifactState::Stored, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::Generating, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::Using, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + Case { caller: ArtifactState::None, responder: ArtifactState::None, expected_caller: ExpectedCallerState { in_update: false, stored_after: false }, expected_responder: ExpectedResponderState { missing: false, stored_after: false } }, + ]; + + let fixture = MpcFixtureBuilder::default() + .only_generate_signatures() + .build() + .await; + + let caller = &fixture.nodes[0]; + let responder = &fixture.nodes[1]; + let all_participants = fixture.sorted_participants(); + + for (i, case) in test_matrix.iter().enumerate() { + let id: u64 = 800 + i as u64; + tracing::info!(id, ?case.caller, ?case.responder, "=== case {i} ==="); + + // Hold slots/taken artifacts alive until assertions are done. + let mut caller_slot = None; + let mut caller_taken = None; + let mut responder_slot = None; + let mut responder_taken = None; + + // --- Set up caller (owner) state --- + match case.caller { + ArtifactState::Stored => { + insert_triples_for_owner( + &caller.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + } + ArtifactState::Generating => { + caller_slot = Some( + caller + .triple_storage + .create_slot(id, caller.me) + .await + .unwrap(), + ); + } + ArtifactState::Using => { + insert_triples_for_owner( + &caller.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + caller_taken = Some(caller.triple_storage.take(id, caller.me).await.unwrap()); + } + ArtifactState::None => {} + } + + // --- Set up responder (peer) state --- + match case.responder { + ArtifactState::Stored => { + insert_triples_for_owner( + &responder.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + } + ArtifactState::Generating => { + responder_slot = Some( + responder + .triple_storage + .create_slot(id, caller.me) + .await + .unwrap(), + ); + } + ArtifactState::Using => { + insert_triples_for_owner( + &responder.triple_storage, + caller.me, + &all_participants, + id..=id, + ) + .await; + responder_taken = Some(responder.triple_storage.take(id, caller.me).await.unwrap()); + } + ArtifactState::None => {} + } + + // --- Build caller's sync update --- + let caller_update = caller.owned_triples_with_reserved().await; + assert_eq!( + caller_update.contains(&id), + case.expected_caller.in_update, + "case {i}: caller={:?} → expected in_update={}", + case.caller, + case.expected_caller.in_update, + ); + + // --- Responder processes the sync update --- + let response = responder + .sync( + caller.me, + caller_update, + responder.owned_presignatures().await, + ) + .await; + + // Verify the full SyncUpdate response from the responder. + assert_eq!( + response.from, responder.me, + "case {i}: response.from should be the responder", + ); + assert_eq!( + response.triples.contains(&id), + case.expected_responder.missing, + "case {i}: caller={:?}, responder={:?} → expected missing={}", + case.caller, + case.responder, + case.expected_responder.missing, + ); + + // --- Verify responder's Redis state after sync --- + assert_eq!( + responder + .triple_storage + .contains_by_owner(id, caller.me) + .await, + case.expected_responder.stored_after, + "case {i}: caller={:?}, responder={:?} → expected responder stored_after={}", + case.caller, + case.responder, + case.expected_responder.stored_after, + ); + + // --- Caller processes the sync response (remove holder / prune) --- + caller + .process_sync_response(responder.me, 2, &response) + .await; + assert_eq!( + caller.triple_storage.contains_by_owner(id, caller.me).await, + case.expected_caller.stored_after, + "case {i}: caller={:?}, responder={:?} → expected caller stored_after={}", + case.caller, + case.responder, + case.expected_caller.stored_after, + ); + + // Clean up held slots/taken before next iteration. + drop(caller_slot); + drop(caller_taken); + drop(responder_slot); + drop(responder_taken); + // Wait for async Drop cleanup. + tokio::time::sleep(Duration::from_millis(50)).await; + } +} From 671ccb207fd9ce8871c527b27f70f045673faace Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Tue, 7 Apr 2026 00:16:58 +0300 Subject: [PATCH 20/21] fix: check for reserved, not only generating --- chain-signatures/node/src/protocol/presignature.rs | 6 +----- chain-signatures/node/src/protocol/triple.rs | 4 ---- chain-signatures/node/src/storage/protocol_storage.rs | 10 +++++++++- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/chain-signatures/node/src/protocol/presignature.rs b/chain-signatures/node/src/protocol/presignature.rs index 04aa58bcd..e95d3c903 100644 --- a/chain-signatures/node/src/protocol/presignature.rs +++ b/chain-signatures/node/src/protocol/presignature.rs @@ -375,10 +375,6 @@ impl PresignatureSpawner { self.ongoing.contains_key(&id) } - pub async fn contains_using(&self, id: PresignatureId) -> bool { - self.presignatures.contains_using(id).await - } - /// Returns the number of unspent presignatures available in the manager. pub async fn len_generated(&self) -> usize { self.presignatures.len_generated().await @@ -430,7 +426,7 @@ impl PresignatureSpawner { // TODO: we can potentially wait for the triples to exist first to then be able to accept. // whereas we just blatantly reject here. The problem with waiting is that the other side // might expire their posit first. - self.triples.contains_generating(id.pair_id).await + self.triples.contains_reserved(id.pair_id).await || self.triples.contains(id.pair_id).await } { tracing::warn!( diff --git a/chain-signatures/node/src/protocol/triple.rs b/chain-signatures/node/src/protocol/triple.rs index e66a3f7da..237397440 100644 --- a/chain-signatures/node/src/protocol/triple.rs +++ b/chain-signatures/node/src/protocol/triple.rs @@ -388,10 +388,6 @@ impl TripleSpawner { self.ongoing.contains_key(&id) } - pub async fn contains_using(&self, id: TripleId) -> bool { - self.triple_storage.contains_using(id).await - } - /// Returns the number of unspent triples assigned to this node. pub async fn len_mine(&self) -> usize { self.triple_storage.len_by_owner(self.me).await diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 2959323fe..3c18b183f 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -155,6 +155,10 @@ impl ReservedState { fn contains_generating(&self, id: &Id) -> bool { self.generating.contains_key(id) } + + fn contains_reserved(&self, id: &Id) -> bool { + self.generating.contains_key(id) || self.using.contains_key(id) + } } #[derive(Debug)] @@ -277,6 +281,10 @@ impl ProtocolStorage { self.reserved.read().await.contains_generating(&id) } + pub async fn contains_reserved(&self, id: A::Id) -> bool { + self.reserved.read().await.contains_reserved(&id) + } + /// Owned artifacts in Redis plus owned using and owned generating. /// This is the full set that should be advertised during state sync to prevent /// peers from pruning artifacts that are still actively in use. @@ -383,7 +391,7 @@ impl ProtocolStorage { let state = self.reserved.read().await; let not_found: Vec<_> = not_found .into_iter() - .filter(|id| !state.contains_generating(id) && !state.using.contains_key(id)) + .filter(|id| !state.contains_reserved(id)) .collect(); Ok(RemoveOutdatedResult::new(outdated, not_found)) } From e5691e3ab6872aa897167b6d9e2c50cc4fa94984 Mon Sep 17 00:00:00 2001 From: Serhii Volovyk Date: Wed, 8 Apr 2026 12:16:44 +0300 Subject: [PATCH 21/21] address comments --- .../node/src/storage/protocol_storage.rs | 55 ++++++++++++++----- integration-tests/tests/cases/state_sync.rs | 2 +- 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/chain-signatures/node/src/storage/protocol_storage.rs b/chain-signatures/node/src/storage/protocol_storage.rs index 3c18b183f..b54dac002 100644 --- a/chain-signatures/node/src/storage/protocol_storage.rs +++ b/chain-signatures/node/src/storage/protocol_storage.rs @@ -84,13 +84,8 @@ impl ArtifactSlot { impl Drop for ArtifactSlot { fn drop(&mut self) { - let storage = self.storage.clone(); - let id = self.id; - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - storage.reserved.write().await.generating.remove(&id); - }); - } + self.storage + .remove_reserved(self.id, ReservedKind::Generating); } } @@ -107,12 +102,7 @@ pub struct ArtifactTakenDropper { impl Drop for ArtifactTakenDropper { fn drop(&mut self) { if let Some(storage) = self.dropper.take() { - let id = self.id; - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - storage.reserved.write().await.using.remove(&id); - }); - } + storage.remove_reserved(self.id, ReservedKind::Using); } } } @@ -133,6 +123,12 @@ impl ArtifactTaken { } } +#[derive(Debug, Clone, Copy)] +enum ReservedKind { + Generating, + Using, +} + /// Tracks artifact IDs that are in-flight but not yet in Redis. /// Protected by a single `RwLock` to avoid multi-lock ordering issues. #[derive(Debug)] @@ -212,7 +208,38 @@ impl ProtocolStorage { } fn me(&self) -> Result { - self.me.get().copied().ok_or(StorageError::NotInitialized) + self.me.get().copied().ok_or_else(|| { + tracing::error!("ProtocolStorage::set_me() was not called"); + StorageError::NotInitialized + }) + } + + /// Remove an ID from the reserved state, trying synchronous lock first. + /// Falls back to spawning an async task if the lock is contended. + fn remove_reserved(&self, id: A::Id, kind: ReservedKind) { + let reserved = self.reserved.clone(); + if let Ok(mut state) = reserved.try_write() { + match kind { + ReservedKind::Generating => state.generating.remove(&id), + ReservedKind::Using => state.using.remove(&id), + }; + return; + } + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let mut state = reserved.write().await; + match kind { + ReservedKind::Generating => state.generating.remove(&id), + ReservedKind::Using => state.using.remove(&id), + }; + }); + } else { + tracing::warn!( + id, + ?kind, + "dropped with contended lock outside tokio runtime; id may remain reserved" + ); + } } async fn connect(&self) -> Option { diff --git a/integration-tests/tests/cases/state_sync.rs b/integration-tests/tests/cases/state_sync.rs index 8175be786..94a9ec69b 100644 --- a/integration-tests/tests/cases/state_sync.rs +++ b/integration-tests/tests/cases/state_sync.rs @@ -435,7 +435,7 @@ async fn test_sync_matrix() { .sync( caller.me, caller_update, - responder.owned_presignatures().await, + vec![], // this matrix only tests triples ) .await;