Merge pull request #1054 from openmina/fix/ledger/mask_leak

Fix: ledger stores more staged ledgers than it needs

Author: binier

cli/src/commands/replay/replay_state_with_input_actions.rs

@@ -10,6 +10,9 @@ pub struct ReplayStateWithInputActions {
     #[arg(long, default_value = "./target/release/libreplay_dynamic_effects.so")]
     pub dynamic_effects_lib: String,
 
+    #[arg(long)]
+    pub ignore_mismatch: bool,
+
     /// Verbosity level
     #[arg(long, short, default_value = "info")]
     pub verbosity: tracing::Level,
@@ -30,13 +33,22 @@ impl ReplayStateWithInputActions {
             }
         };
 
-        replay_state_with_input_actions(&dir, dynamic_effects_lib, check_build_env)?;
+        replay_state_with_input_actions(
+            &dir,
+            dynamic_effects_lib,
+            self.ignore_mismatch,
+            check_build_env,
+        )?;
 
         Ok(())
     }
 }
 
-pub fn check_build_env(record_env: &BuildEnv, replay_env: &BuildEnv) -> anyhow::Result<()> {
+pub fn check_build_env(
+    record_env: &BuildEnv,
+    replay_env: &BuildEnv,
+    ignore_mismatch: bool,
+) -> anyhow::Result<()> {
     let is_git_same = record_env.git.commit_hash == replay_env.git.commit_hash;
     let is_cargo_same = record_env.cargo == replay_env.cargo;
     let is_rustc_same = record_env.rustc == replay_env.rustc;
@@ -47,7 +59,8 @@ pub fn check_build_env(record_env: &BuildEnv, replay_env: &BuildEnv) -> anyhow::
             record_env.git, replay_env.git
         );
         let msg = format!("git build env mismatch!\n{diff}");
-        if console::user_attended() {
+        if ignore_mismatch {
+        } else if console::user_attended() {
             use dialoguer::Confirm;
 
             let prompt = format!("{msg}\nDo you want to continue?");

ledger/src/mask/mod.rs

@@ -41,6 +41,9 @@ pub fn alive_len() -> usize {
     exec(|list| list.len())
 }
 
-pub fn alive_collect() -> Vec<Uuid> {
+pub fn alive_collect<B>() -> B
+where
+    B: FromIterator<Uuid>,
+{
     exec(|list| list.iter().cloned().collect())
 }

ledger/src/staged_ledger/staged_ledger.rs

@@ -116,7 +116,6 @@ pub struct StagedLedger {
 }
 
 impl StagedLedger {
-    #[cfg(feature = "fuzzing")]
     pub fn ledger_ref(&self) -> &Mask {
         &self.ledger
     }

mina-p2p-messages/src/v2/hashing.rs

@@ -115,7 +115,7 @@ impl Serialize for TransactionHash {
         if serializer.is_human_readable() {
             serializer.serialize_str(&self.to_string())
         } else {
-            self.0.serialize(serializer)
+            serde_bytes::serialize(&*self.0, serializer)
         }
     }
 }
@@ -129,9 +129,7 @@ impl<'de> serde::Deserialize<'de> for TransactionHash {
             let b58: String = Deserialize::deserialize(deserializer)?;
             Ok(b58.parse().map_err(|err| serde::de::Error::custom(err))?)
         } else {
-            let v = Vec::deserialize(deserializer)?;
-            v.try_into()
-                .map_err(|_| serde::de::Error::custom("transaction hash wrong size"))
+            serde_bytes::deserialize(deserializer)
                 .map(Arc::new)
                 .map(Self)
         }

node/native/src/replay.rs

@@ -1,16 +1,19 @@
 use std::cell::RefCell;
 
 use node::{
-    core::thread, recorder::StateWithInputActionsReader, snark::BlockVerifier, ActionWithMeta,
-    BuildEnv, Store,
+    core::thread,
+    recorder::StateWithInputActionsReader,
+    snark::{BlockVerifier, TransactionVerifier},
+    ActionWithMeta, BuildEnv, Store,
 };
 
 use crate::NodeService;
 
 pub fn replay_state_with_input_actions(
     dir: &str,
     dynamic_effects_lib: Option<String>,
-    mut check_build_env: impl FnMut(&BuildEnv, &BuildEnv) -> anyhow::Result<()>,
+    ignore_mismatch: bool,
+    mut check_build_env: impl FnMut(&BuildEnv, &BuildEnv, bool) -> anyhow::Result<()>,
 ) -> anyhow::Result<crate::Node> {
     eprintln!("replaying node based on initial state and actions from the dir: {dir}");
     let reader = StateWithInputActionsReader::new(dir);
@@ -31,6 +34,8 @@ pub fn replay_state_with_input_actions(
         // index/srs doesn't match deserialized one.
         state.snark.block_verify.verifier_index = BlockVerifier::make();
         state.snark.block_verify.verifier_srs = node::snark::get_srs();
+        state.snark.user_command_verify.verifier_index = TransactionVerifier::make();
+        state.snark.user_command_verify.verifier_srs = node::snark::get_srs();
         state
     };
 
@@ -46,7 +51,7 @@ pub fn replay_state_with_input_actions(
     let store = node.store_mut();
 
     let replay_env = BuildEnv::get();
-    check_build_env(&store.state().config.build, &replay_env)?;
+    check_build_env(&store.state().config.build, &replay_env, ignore_mismatch)?;
 
     eprintln!("reading actions from dir: {dir}");
 

node/src/ledger/ledger_reducer.rs

@@ -17,6 +17,16 @@ impl LedgerState {
                 } = action
                 {
                     if let Ok(state) = state_context.get_substate_mut() {
+                        if result.alive_masks > 294 {
+                            // TODO(binier): should be a bug condition, but can't be
+                            // because we get false positive during testing, since
+                            // multiple nodes/ledger run in the same process.
+                            openmina_core::log::warn!(
+                                meta.time();
+                                "ledger mask leak: more than 294 ledger masks ({}) detected!",
+                                result.alive_masks
+                            );
+                        }
                         state.alive_masks = result.alive_masks;
                     }
                 }

node/src/ledger/ledger_service.rs

@@ -1,10 +1,7 @@
 use super::{
     ledger_empty_hash_at_depth,
-    read::LedgerReadResponse,
-    read::{LedgerReadId, LedgerReadRequest},
-    write::CommitResult,
-    write::LedgerWriteRequest,
-    write::LedgerWriteResponse,
+    read::{LedgerReadId, LedgerReadRequest, LedgerReadResponse},
+    write::{CommitResult, LedgerWriteRequest, LedgerWriteResponse, LedgersToKeep},
     LedgerAddress, LedgerEvent, LEDGER_DEPTH,
 };
 use crate::{
@@ -65,8 +62,8 @@ use mina_p2p_messages::{
 };
 use mina_signer::CompressedPubKey;
 use openmina_core::{
-    block::AppliedBlock,
-    block::ArcBlockWithHash,
+    block::{AppliedBlock, ArcBlockWithHash},
+    bug_condition,
     constants::constraint_constants,
     snark::{Snark, SnarkJobId},
     thread,
@@ -137,18 +134,18 @@ impl StagedLedgersStorage {
 
     fn retain<F>(&mut self, fun: F)
     where
-        F: Fn(&LedgerHash, &[Arc<MinaBaseStagedLedgerHashStableV1>]) -> bool,
+        F: Fn(&MinaBaseStagedLedgerHashStableV1) -> bool,
     {
-        self.by_merkle_root_hash
-            .retain(|merkle_root_hash, staged_ledger_hashes| {
-                let retain = fun(merkle_root_hash, staged_ledger_hashes.as_slice());
-                if !retain {
-                    for staged_ledger_hash in staged_ledger_hashes {
-                        self.staged_ledgers.remove(staged_ledger_hash);
-                    }
+        self.by_merkle_root_hash.retain(|_, staged_ledger_hashes| {
+            staged_ledger_hashes.retain(|hash| {
+                if fun(hash) {
+                    return true;
                 }
-                retain
+                self.staged_ledgers.remove(hash);
+                false
             });
+            !staged_ledger_hashes.is_empty()
+        });
     }
 
     fn extend<I>(&mut self, iterator: I)
@@ -857,7 +854,7 @@ impl LedgerCtx {
 
     pub fn commit(
         &mut self,
-        ledgers_to_keep: BTreeSet<LedgerHash>,
+        ledgers_to_keep: LedgersToKeep,
         root_snarked_ledger_updates: TransitionFrontierRootSnarkedLedgerUpdates,
         needed_protocol_states: BTreeMap<StateHash, MinaStateProtocolStateValueStableV2>,
         new_root: &ArcBlockWithHash,
@@ -902,13 +899,13 @@ impl LedgerCtx {
         );
 
         self.staged_ledgers
-            .retain(|hash, _| ledgers_to_keep.contains(hash));
+            .retain(|hash| ledgers_to_keep.contains(hash));
         self.staged_ledgers.extend(
             self.sync
                 .staged_ledgers
                 .take()
                 .into_iter()
-                .filter(|(hash, _)| ledgers_to_keep.contains(&hash.non_snark.ledger_hash)),
+                .filter(|(hash, _)| ledgers_to_keep.contains(&**hash)),
         );
 
         for ledger_hash in [
@@ -973,13 +970,45 @@ impl LedgerCtx {
                 .unwrap_or_default(),
         );
 
+        // self.check_alive_masks();
+
         CommitResult {
             alive_masks: ::ledger::mask::alive_len(),
             available_jobs,
             needed_protocol_states,
         }
     }
 
+    #[allow(dead_code)]
+    fn check_alive_masks(&mut self) {
+        let mut alive: BTreeSet<_> = ::ledger::mask::alive_collect();
+        let staged_ledgers = self
+            .staged_ledgers
+            .staged_ledgers
+            .iter()
+            .map(|(hash, ledger)| (&hash.non_snark.ledger_hash, ledger.ledger_ref()));
+
+        let alive_ledgers = self
+            .snarked_ledgers
+            .iter()
+            .chain(staged_ledgers)
+            .map(|(hash, mask)| {
+                let uuid = mask.get_uuid();
+                if !alive.remove(&uuid) {
+                    bug_condition!("mask not found among alive masks! uuid: {uuid}, hash: {hash}");
+                }
+                (uuid, hash)
+            })
+            .collect::<Vec<_>>();
+        openmina_core::debug!(redux::Timestamp::global_now(); "alive_ledgers_after_commit: {alive_ledgers:#?}");
+
+        if !alive.is_empty() {
+            bug_condition!(
+                "masks alive which are no longer part of the ledger service: {alive:#?}"
+            );
+        }
+    }
+
     pub fn get_num_accounts(
         &mut self,
         ledger_hash: v2::LedgerHash,

node/src/ledger/write/mod.rs

@@ -54,7 +54,7 @@ pub enum LedgerWriteRequest {
         skip_verification: bool,
     },
     Commit {
-        ledgers_to_keep: BTreeSet<v2::LedgerHash>,
+        ledgers_to_keep: LedgersToKeep,
         root_snarked_ledger_updates: TransitionFrontierRootSnarkedLedgerUpdates,
         needed_protocol_states: BTreeMap<v2::StateHash, v2::MinaStateProtocolStateValueStableV2>,
         new_root: AppliedBlock,
@@ -162,6 +162,64 @@ impl TryFrom<&BlockApplyResult> for v2::ArchiveTransitionFronntierDiff {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug, Ord, PartialOrd, Eq, PartialEq, Default, Clone)]
+pub struct LedgersToKeep {
+    snarked: BTreeSet<v2::LedgerHash>,
+    staged: BTreeSet<Arc<v2::MinaBaseStagedLedgerHashStableV1>>,
+}
+
+impl LedgersToKeep {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn contains<'a, T>(&self, key: T) -> bool
+    where
+        T: 'a + Into<LedgerToKeep<'a>>,
+    {
+        match key.into() {
+            LedgerToKeep::Snarked(hash) => self.snarked.contains(hash),
+            LedgerToKeep::Staged(hash) => self.staged.contains(hash),
+        }
+    }
+
+    pub fn add_snarked(&mut self, hash: v2::LedgerHash) -> bool {
+        self.snarked.insert(hash)
+    }
+
+    pub fn add_staged(&mut self, hash: Arc<v2::MinaBaseStagedLedgerHashStableV1>) -> bool {
+        self.staged.insert(hash)
+    }
+}
+
+impl<'a> FromIterator<&'a ArcBlockWithHash> for LedgersToKeep {
+    fn from_iter<T: IntoIterator<Item = &'a ArcBlockWithHash>>(iter: T) -> Self {
+        let mut res = Self::new();
+        let best_tip = iter.into_iter().fold(None, |best_tip, block| {
+            res.add_snarked(block.snarked_ledger_hash().clone());
+            res.add_staged(Arc::new(block.staged_ledger_hashes().clone()));
+            match best_tip {
+                None => Some(block),
+                Some(tip) if tip.height() < block.height() => Some(block),
+                old => old,
+            }
+        });
+
+        if let Some(best_tip) = best_tip {
+            res.add_snarked(best_tip.staking_epoch_ledger_hash().clone());
+            res.add_snarked(best_tip.next_epoch_ledger_hash().clone());
+        }
+
+        res
+    }
+}
+
+#[derive(derive_more::From)]
+pub enum LedgerToKeep<'a> {
+    Snarked(&'a v2::LedgerHash),
+    Staged(&'a v2::MinaBaseStagedLedgerHashStableV1),
+}
+
 #[derive(Serialize, Deserialize, Debug, Default, Clone)]
 pub struct CommitResult {
     pub alive_masks: usize,

node/src/recorder/recorder.rs

@@ -109,6 +109,8 @@ impl Recorder {
                 let mut writer = BufWriter::new(file);
 
                 let encoded = data.encode().unwrap();
+                // RecordedActionWithMeta::decode(&encoded)
+                //     .expect(&format!("failed to decode encoded message: {:?}", data));
                 writer
                     .write_all(&(encoded.len() as u64).to_be_bytes())
                     .unwrap();

node/src/transition_frontier/sync/transition_frontier_sync_effects.rs

@@ -4,7 +4,7 @@ use p2p::channels::rpc::{P2pChannelsRpcAction, P2pRpcId};
 use p2p::{P2pNetworkPubsubAction, PeerId};
 use redux::ActionMeta;
 
-use crate::ledger::write::{LedgerWriteAction, LedgerWriteRequest};
+use crate::ledger::write::{LedgerWriteAction, LedgerWriteRequest, LedgersToKeep};
 use crate::p2p::channels::rpc::P2pRpcRequest;
 use crate::service::TransitionFrontierSyncLedgerSnarkedService;
 use crate::{p2p_ready, Service, Store, TransitionFrontierAction};
@@ -354,16 +354,8 @@ impl TransitionFrontierSyncAction {
                 };
                 let ledgers_to_keep = chain
                     .iter()
-                    .flat_map(|b| {
-                        [
-                            b.snarked_ledger_hash(),
-                            b.merkle_root_hash(),
-                            b.staking_epoch_ledger_hash(),
-                            b.next_epoch_ledger_hash(),
-                        ]
-                    })
-                    .cloned()
-                    .collect();
+                    .map(|block| &block.block)
+                    .collect::<LedgersToKeep>();
                 let mut root_snarked_ledger_updates = root_snarked_ledger_updates.clone();
                 if transition_frontier
                     .best_chain