Optimization and Refractor: UTXO index and database iterators.

Cargo.toml

@@ -235,6 +235,7 @@ secp256k1 = { version = "0.29.0", features = [
     "rand-std",
     "serde",
 ] } # TODO "0.28.0"
+self_cell = "1.2.2"
 separator = "0.4.1"
 seqlock = "0.2.0"
 serde = { version = "1.0.190", features = ["derive", "rc"] }

components/addressmanager/src/stores/address_store.rs

@@ -72,6 +72,19 @@ impl From<DbAddressKey> for AddressKey {
     }
 }
 
+impl TryFrom<&[u8]> for DbAddressKey {
+    type Error = &'static str;
+
+    fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
+        if slice.len() != ADDRESS_KEY_SIZE {
+            return Err("Invalid slice length for DbAddressKey");
+        }
+        let mut bytes = [0u8; ADDRESS_KEY_SIZE];
+        bytes.copy_from_slice(slice);
+        Ok(Self(bytes))
+    }
+}
+
 #[derive(Clone)]
 pub struct DbAddressesStore {
     db: Arc<DB>,
@@ -85,15 +98,11 @@ impl DbAddressesStore {
 
     pub fn iterator(&self) -> impl Iterator<Item = Result<(AddressKey, Entry), Box<dyn Error>>> + '_ {
         self.access.iterator().map(|iter_result| match iter_result {
-            Ok((key_bytes, connection_failed_count)) => match <[u8; ADDRESS_KEY_SIZE]>::try_from(&key_bytes[..]) {
-                Ok(address_key_slice) => {
-                    let addr_key = DbAddressKey(address_key_slice);
-                    let address: AddressKey = addr_key.into();
-                    Ok((address, connection_failed_count))
-                }
-                Err(e) => Err(e.into()),
-            },
-            Err(e) => Err(e),
+            Ok((key, entry)) => {
+                let address: AddressKey = key.into();
+                Ok((address, entry))
+            }
+            Err(e) => Err(e.into()),
         })
     }
 }

components/addressmanager/src/stores/banned_address_store.rs

@@ -66,6 +66,19 @@ impl From<AddressKey> for IpAddr {
     }
 }
 
+impl TryFrom<&[u8]> for AddressKey {
+    type Error = &'static str;
+
+    fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
+        if slice.len() != ADDRESS_KEY_SIZE {
+            return Err("Invalid slice length for AddressKey");
+        }
+        let mut bytes = [0u8; ADDRESS_KEY_SIZE];
+        bytes.copy_from_slice(slice);
+        Ok(Self(bytes))
+    }
+}
+
 #[derive(Clone)]
 pub struct DbBannedAddressesStore {
     db: Arc<DB>,
@@ -79,15 +92,11 @@ impl DbBannedAddressesStore {
 
     pub fn iterator(&self) -> impl Iterator<Item = Result<(IpAddr, ConnectionBanTimestamp), Box<dyn Error>>> + '_ {
         self.access.iterator().map(|iter_result| match iter_result {
-            Ok((key_bytes, connection_ban_timestamp)) => match <[u8; ADDRESS_KEY_SIZE]>::try_from(&key_bytes[..]) {
-                Ok(address_key_slice) => {
-                    let addr_key = AddressKey(address_key_slice);
-                    let address: IpAddr = addr_key.into();
-                    Ok((address, connection_ban_timestamp))
-                }
-                Err(e) => Err(e.into()),
-            },
-            Err(e) => Err(e),
+            Ok((key, connection_ban_timestamp)) => {
+                let address: IpAddr = key.into();
+                Ok((address, connection_ban_timestamp))
+            }
+            Err(e) => Err(e.into()),
         })
     }
 }

components/consensusmanager/src/session.rs

@@ -278,15 +278,6 @@ impl ConsensusSessionOwned {
         self.clone().spawn_blocking(move |c| c.get_virtual_chain_from_block(low, chain_path_added_limit)).await
     }
 
-    pub async fn async_get_virtual_utxos(
-        &self,
-        from_outpoint: Option<TransactionOutpoint>,
-        chunk_size: usize,
-        skip_first: bool,
-    ) -> Vec<(TransactionOutpoint, UtxoEntry)> {
-        self.clone().spawn_blocking(move |c| c.get_virtual_utxos(from_outpoint, chunk_size, skip_first)).await
-    }
-
     pub async fn async_get_tips(&self) -> Vec<Hash> {
         self.clone().spawn_blocking(|c| c.get_tips()).await
     }
@@ -432,7 +423,7 @@ impl ConsensusSessionOwned {
         from_outpoint: Option<TransactionOutpoint>,
         chunk_size: usize,
         skip_first: bool,
-    ) -> ConsensusResult<Vec<(TransactionOutpoint, UtxoEntry)>> {
+    ) -> ConsensusResult<Vec<(TransactionOutpoint, Arc<UtxoEntry>)>> {
         self.clone()
             .spawn_blocking(move |c| c.get_pruning_point_utxos(expected_pruning_point, from_outpoint, chunk_size, skip_first))
             .await

consensus/core/src/api/mod.rs

@@ -214,12 +214,7 @@ pub trait ConsensusApi: Send + Sync {
         unimplemented!()
     }
 
-    fn get_virtual_utxos(
-        &self,
-        from_outpoint: Option<TransactionOutpoint>,
-        chunk_size: usize,
-        skip_first: bool,
-    ) -> Vec<(TransactionOutpoint, UtxoEntry)> {
+    fn get_virtual_utxo_iter_owned(&self) -> Box<dyn Iterator<Item = (TransactionOutpoint, Arc<UtxoEntry>)> + Send> {
         unimplemented!()
     }
 
@@ -364,7 +359,7 @@ pub trait ConsensusApi: Send + Sync {
         from_outpoint: Option<TransactionOutpoint>,
         chunk_size: usize,
         skip_first: bool,
-    ) -> ConsensusResult<Vec<(TransactionOutpoint, UtxoEntry)>> {
+    ) -> ConsensusResult<Vec<(TransactionOutpoint, Arc<UtxoEntry>)>> {
         unimplemented!()
     }
 

consensus/core/src/tx.rs

@@ -8,6 +8,11 @@
 
 mod script_public_key;
 
+use crate::mass::{ContextualMasses, NonContextualMasses};
+use crate::{
+    hashing,
+    subnets::{self, SubnetworkId},
+};
 use borsh::{BorshDeserialize, BorshSerialize};
 use kaspa_utils::hex::ToHex;
 use kaspa_utils::mem_size::MemSizeEstimator;
@@ -27,12 +32,6 @@ use std::{
 };
 use wasm_bindgen::prelude::*;
 
-use crate::mass::{ContextualMasses, NonContextualMasses};
-use crate::{
-    hashing,
-    subnets::{self, SubnetworkId},
-};
-
 /// COINBASE_TRANSACTION_INDEX is the index of the coinbase transaction in every block
 pub const COINBASE_TRANSACTION_INDEX: usize = 0;
 /// A 32-byte Kaspa transaction identifier.

consensus/src/consensus/factory.rs

@@ -20,7 +20,7 @@ use kaspa_utils::mem_size::MemSizeEstimator;
 use parking_lot::RwLock;
 use rocksdb::WriteBatch;
 use serde::{Deserialize, Serialize};
-use std::{collections::HashMap, error::Error, fs, path::PathBuf, sync::Arc};
+use std::{collections::HashMap, fs, path::PathBuf, sync::Arc};
 
 #[derive(Serialize, Deserialize, Clone)]
 pub struct ConsensusEntry {
@@ -181,14 +181,14 @@ impl MultiConsensusManagementStore {
         Ok(())
     }
 
-    fn iterator(&self) -> impl Iterator<Item = Result<ConsensusEntry, Box<dyn Error>>> + '_ {
+    fn iterator(&self) -> impl Iterator<Item = Result<ConsensusEntry, StoreError>> + '_ {
         self.entries.iterator().map(|iter_result| match iter_result {
             Ok((_, entry)) => Ok(entry),
             Err(e) => Err(e),
         })
     }
 
-    fn iterate_inactive_entries(&self) -> impl Iterator<Item = Result<ConsensusEntry, Box<dyn Error>>> + '_ {
+    fn iterate_inactive_entries(&self) -> impl Iterator<Item = Result<ConsensusEntry, StoreError>> + '_ {
         let current_consensus_key = self.metadata.read().unwrap().current_consensus_key;
         self.iterator().filter(move |entry_result| {
             if let Ok(entry) = entry_result {

consensus/src/consensus/mod.rs

@@ -1005,15 +1005,9 @@ impl ConsensusApi for Consensus {
         self.lkg_virtual_state.load().parents.len()
     }
 
-    fn get_virtual_utxos(
-        &self,
-        from_outpoint: Option<TransactionOutpoint>,
-        chunk_size: usize,
-        skip_first: bool,
-    ) -> Vec<(TransactionOutpoint, UtxoEntry)> {
+    fn get_virtual_utxo_iter_owned(&self) -> Box<dyn Iterator<Item = (TransactionOutpoint, Arc<UtxoEntry>)> + Send> {
         let virtual_stores = self.virtual_stores.read();
-        let iter = virtual_stores.utxo_set.seek_iterator(from_outpoint, chunk_size, skip_first);
-        iter.map(|item| item.unwrap()).collect()
+        Box::new(virtual_stores.utxo_set.iterator_owned().map(|r| r.unwrap()))
     }
 
     fn get_tips(&self) -> Vec<Hash> {
@@ -1030,7 +1024,7 @@ impl ConsensusApi for Consensus {
         from_outpoint: Option<TransactionOutpoint>,
         chunk_size: usize,
         skip_first: bool,
-    ) -> ConsensusResult<Vec<(TransactionOutpoint, UtxoEntry)>> {
+    ) -> ConsensusResult<Vec<(TransactionOutpoint, Arc<UtxoEntry>)>> {
         if self.pruning_point_store.read().pruning_point().unwrap() != expected_pruning_point {
             return Err(ConsensusError::UnexpectedPruningPoint);
         }

consensus/src/model/stores/mod.rs

@@ -40,6 +40,19 @@ impl AsRef<[u8]> for U64Key {
     }
 }
 
+impl TryFrom<&[u8]> for U64Key {
+    type Error = &'static str;
+
+    fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
+        if slice.len() != size_of::<u64>() {
+            return Err("Invalid slice length for U64Key");
+        }
+        let mut bytes = [0u8; size_of::<u64>()];
+        bytes.copy_from_slice(slice);
+        Ok(Self(bytes))
+    }
+}
+
 impl Display for U64Key {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", u64::from_le_bytes(self.0))

consensus/src/model/stores/reachability.rs

@@ -628,8 +628,6 @@ impl ReachabilityStoreReader for StagingReachabilityStore<'_> {
             .access
             .iterator()
             .map(|r| r.unwrap().0)
-            .map(|k| <[u8; kaspa_hashes::HASH_SIZE]>::try_from(&k[..]).unwrap())
-            .map(Hash::from_bytes)
             .chain(self.staging_writes.keys().copied())
             .collect::<BlockHashSet>()
             .difference(&self.staging_deletions)

consensus/src/model/stores/relations.rs

@@ -346,8 +346,6 @@ impl RelationsStoreReader for StagingRelationsStore<'_> {
             .parents_access
             .iterator()
             .map(|r| r.unwrap().0)
-            .map(|k| <[u8; kaspa_hashes::HASH_SIZE]>::try_from(&k[..]).unwrap())
-            .map(Hash::from_bytes)
             .chain(self.parents_overrides.keys().copied())
             .collect::<BlockHashSet>()
             .difference(&self.entry_deletions)

consensus/src/model/stores/utxo_set.rs

@@ -11,19 +11,22 @@ use kaspa_database::prelude::{BatchDbWriter, CachedDbAccess, DirectDbWriter};
 use kaspa_database::prelude::{CachePolicy, StoreError};
 use kaspa_hashes::Hash;
 use rocksdb::WriteBatch;
-use std::{error::Error, fmt::Display, sync::Arc};
+use std::{fmt::Display, sync::Arc};
 
-type UtxoCollectionIterator<'a> = Box<dyn Iterator<Item = Result<(TransactionOutpoint, UtxoEntry), Box<dyn Error>>> + 'a>;
+type UtxoCollectionIterator<'a> = Box<dyn Iterator<Item = Result<(TransactionOutpoint, Arc<UtxoEntry>), StoreError>> + 'a>;
+type OwnedUtxoCollectionIterator = Box<dyn Iterator<Item = Result<(TransactionOutpoint, Arc<UtxoEntry>), StoreError>> + Send>;
 
 pub trait UtxoSetStoreReader {
     fn get(&self, outpoint: &TransactionOutpoint) -> Result<Arc<UtxoEntry>, StoreError>;
-    fn seek_iterator(&self, from_outpoint: Option<TransactionOutpoint>, limit: usize, skip_first: bool) -> UtxoCollectionIterator<'_>;
+    fn iterator(&self) -> UtxoCollectionIterator<'_>;
+    fn iterator_owned(&self) -> OwnedUtxoCollectionIterator;
+    fn seek_iterator(&self, seek_from: Option<TransactionOutpoint>, limit: usize, skip_first: bool) -> UtxoCollectionIterator<'_>;
 }
 
 pub trait UtxoSetStore: UtxoSetStoreReader {
     /// Updates the store according to the UTXO diff -- adding and deleting entries correspondingly.
     /// Note we define `self` as `mut` in order to require write access even though the compiler does not require it.
-    /// This is because concurrent readers can interfere with cache consistency.  
+    /// This is because concurrent readers can interfere with cache consistency.
     fn write_diff(&mut self, utxo_diff: &UtxoDiff) -> Result<(), StoreError>;
     fn write_many(&mut self, utxos: &[(TransactionOutpoint, UtxoEntry)]) -> Result<(), StoreError>;
 }
@@ -50,7 +53,7 @@ impl TryFrom<&[u8]> for UtxoKey {
         if UTXO_KEY_SIZE < slice.len() {
             return Err("src slice is too large");
         }
-        if slice.len() < kaspa_hashes::HASH_SIZE + 1 {
+        if slice.len() < kaspa_hashes::HASH_SIZE {
             return Err("src slice is too short");
         }
         // If the slice is shorter than HASH len + u32 len then we pad with zeros, effectively
@@ -111,19 +114,6 @@ impl DbUtxoSetStore {
         Ok(())
     }
 
-    pub fn iterator(&self) -> impl Iterator<Item = Result<(TransactionOutpoint, Arc<UtxoEntry>), Box<dyn Error>>> + '_ {
-        self.access.iterator().map(|iter_result| match iter_result {
-            Ok((key_bytes, utxo_entry)) => match UtxoKey::try_from(key_bytes.as_ref()) {
-                Ok(utxo_key) => {
-                    let outpoint: TransactionOutpoint = utxo_key.into();
-                    Ok((outpoint, utxo_entry))
-                }
-                Err(e) => Err(e.into()),
-            },
-            Err(e) => Err(e),
-        })
-    }
-
     /// Clear the store completely in DB and cache
     pub fn clear(&mut self) -> Result<(), StoreError> {
         self.access.delete_all(DirectDbWriter::new(&self.db))
@@ -151,12 +141,27 @@ impl UtxoSetStoreReader for DbUtxoSetStore {
         self.access.read((*outpoint).into())
     }
 
-    fn seek_iterator(&self, from_outpoint: Option<TransactionOutpoint>, limit: usize, skip_first: bool) -> UtxoCollectionIterator<'_> {
-        let seek_key = from_outpoint.map(UtxoKey::from);
-        Box::new(self.access.seek_iterator(None, seek_key, limit, skip_first).map(|res| {
-            let (key, entry) = res?;
-            let outpoint: TransactionOutpoint = UtxoKey::try_from(key.as_ref()).unwrap().into();
-            Ok((outpoint, UtxoEntry::clone(&entry)))
+    fn iterator(&self) -> UtxoCollectionIterator<'_> {
+        Box::new(self.access.iterator().map(|res| {
+            let (k, v) = res?;
+            let outpoint: TransactionOutpoint = k.into();
+            Ok((outpoint, v))
+        }))
+    }
+
+    fn seek_iterator(&self, seek_from: Option<TransactionOutpoint>, limit: usize, skip_first: bool) -> UtxoCollectionIterator<'_> {
+        Box::new(self.access.seek_iterator(None, seek_from.map(UtxoKey::from), limit, skip_first).map(|res| {
+            let (k, v) = res?;
+            let outpoint: TransactionOutpoint = k.into();
+            Ok((outpoint, v))
+        }))
+    }
+
+    fn iterator_owned(&self) -> OwnedUtxoCollectionIterator {
+        Box::new(self.access.iterator_owned().map(|res| {
+            let (k, v) = res?;
+            let outpoint: TransactionOutpoint = k.into();
+            Ok((outpoint, v))
         }))
     }
 }

consensus/src/pipeline/pruning_processor/processor.rs

@@ -19,6 +19,7 @@ use crate::{
             statuses::StatusesStoreReader,
             tips::{TipsStore, TipsStoreReader},
             utxo_diffs::UtxoDiffsStoreReader,
+            utxo_set::UtxoSetStoreReader,
             virtual_state::VirtualStateStoreReader,
         },
     },

consensus/src/pipeline/virtual_processor/processor.rs

@@ -32,6 +32,7 @@ use crate::{
             tips::{DbTipsStore, TipsStoreReader},
             utxo_diffs::{DbUtxoDiffsStore, UtxoDiffsStoreReader},
             utxo_multisets::{DbUtxoMultisetsStore, UtxoMultisetsStoreReader},
+            utxo_set::UtxoSetStoreReader,
             virtual_state::{LkgVirtualState, VirtualState, VirtualStateStoreReader, VirtualStores},
             DB,
         },
@@ -691,7 +692,7 @@ impl VirtualStateProcessor {
     /// Assumes:
     ///     1. `selected_parent` is a UTXO-valid block
     ///     2. `candidates` are an antichain ordered in descending blue work order
-    ///     3. `candidates` do not contain `selected_parent` and `selected_parent.blue work > max(candidates.blue_work)`  
+    ///     3. `candidates` do not contain `selected_parent` and `selected_parent.blue work > max(candidates.blue_work)`
     pub(super) fn pick_virtual_parents(
         &self,
         selected_parent: Hash,