Merge pull request #81 from Steake/copilot/implement-state-persistence-rocksdb

Steake · web-flow · commit 81091ccbd4b4 · 2025-12-11T05:12:23.000+07:00
Implement RocksDB-backed state persistence for accounts and bonds
diff --git a/crates/bitcell-node/Cargo.toml b/crates/bitcell-node/Cargo.toml
@@ -39,3 +39,4 @@ base64 = "0.21"
 
 [dev-dependencies]
 proptest.workspace = true
+tempfile = "3.23.0"
diff --git a/crates/bitcell-node/src/blockchain.rs b/crates/bitcell-node/src/blockchain.rs
@@ -5,11 +5,10 @@
 ///! - Block validation including signature, VRF, and transaction verification
 ///! - Transaction indexing for efficient lookups
 ///! - State management with Merkle tree root computation
-
 use crate::{Result, MetricsRegistry};
 use bitcell_consensus::{Block, BlockHeader, Transaction, BattleProof};
 use bitcell_crypto::{Hash256, PublicKey, SecretKey};
-use bitcell_economics::{COIN, INITIAL_BLOCK_REWARD, HALVING_INTERVAL, MAX_HALVINGS};
+use bitcell_economics::{INITIAL_BLOCK_REWARD, HALVING_INTERVAL, MAX_HALVINGS};
 use bitcell_state::StateManager;
 use std::sync::{Arc, RwLock};
 use std::collections::HashMap;
@@ -78,6 +77,54 @@ impl Blockchain {
         blockchain
     }
     
+    /// Create new blockchain with persistent storage
+    /// 
+    /// This method initializes the blockchain with RocksDB-backed state storage.
+    /// State will be persisted to disk and restored across node restarts.
+    /// 
+    /// # Arguments
+    /// * `secret_key` - Node's secret key for signing
+    /// * `metrics` - Metrics registry
+    /// * `data_path` - Path to the data directory for persistent storage
+    pub fn with_storage(
+        secret_key: Arc<SecretKey>, 
+        metrics: MetricsRegistry,
+        data_path: &std::path::Path,
+    ) -> std::result::Result<Self, String> {
+        // Create storage manager
+        let storage_path = data_path.join("state");
+        let storage = Arc::new(
+            bitcell_state::StorageManager::new(&storage_path)
+                .map_err(|e| format!("Failed to create storage: {}", e))?
+        );
+        
+        // Create state manager with storage
+        let state = StateManager::with_storage(storage)
+            .map_err(|e| format!("Failed to initialize state: {:?}", e))?;
+        
+        let genesis = Self::create_genesis_block(&secret_key);
+        let genesis_hash = genesis.hash();
+        
+        let mut blocks = HashMap::new();
+        blocks.insert(GENESIS_HEIGHT, genesis);
+        
+        let blockchain = Self {
+            height: Arc::new(RwLock::new(GENESIS_HEIGHT)),
+            latest_hash: Arc::new(RwLock::new(genesis_hash)),
+            blocks: Arc::new(RwLock::new(blocks)),
+            tx_index: Arc::new(RwLock::new(HashMap::new())),
+            state: Arc::new(RwLock::new(state)),
+            metrics: metrics.clone(),
+            secret_key,
+        };
+        
+        // Initialize metrics
+        blockchain.metrics.set_chain_height(GENESIS_HEIGHT);
+        blockchain.metrics.set_sync_progress(100);
+        
+        Ok(blockchain)
+    }
+    
     /// Create genesis block
     fn create_genesis_block(secret_key: &SecretKey) -> Block {
         let header = BlockHeader {
@@ -518,6 +565,40 @@ mod tests {
         // Test reward becomes 0 after 64 halvings
         assert_eq!(Blockchain::calculate_block_reward(HALVING_INTERVAL * 64), 0);
     }
+
+    #[test]
+    fn test_blockchain_with_persistent_storage() {
+        use tempfile::TempDir;
+        
+        let temp_dir = TempDir::new().unwrap();
+        let data_path = temp_dir.path();
+        let sk = Arc::new(SecretKey::generate());
+        let pubkey = [1u8; 33];
+        
+        // Create blockchain with storage and modify state
+        {
+            let metrics = MetricsRegistry::new();
+            let blockchain = Blockchain::with_storage(sk.clone(), metrics, data_path).unwrap();
+            
+            // Add an account to state
+            let mut state = blockchain.state.write().unwrap();
+            state.update_account(pubkey, bitcell_state::Account {
+                balance: 1000,
+                nonce: 5,
+            });
+        }
+        
+        // Recreate blockchain from same storage and verify persistence
+        {
+            let metrics = MetricsRegistry::new();
+            let blockchain = Blockchain::with_storage(sk, metrics, data_path).unwrap();
+            
+            let state = blockchain.state.read().unwrap();
+            let account = state.get_account_owned(&pubkey).expect("Account should persist");
+            assert_eq!(account.balance, 1000);
+            assert_eq!(account.nonce, 5);
+        }
+    }
     
     #[test]
     fn test_vrf_block_production_and_validation() {
diff --git a/crates/bitcell-node/src/config.rs b/crates/bitcell-node/src/config.rs
@@ -14,6 +14,8 @@ pub struct NodeConfig {
     /// Block production interval in seconds.
     /// Defaults to 10 seconds for testing. Use 600 (10 minutes) for production.
     pub block_time_secs: u64,
+    /// Data directory for persistent storage. If None, uses in-memory storage only.
+    pub data_dir: Option<std::path::PathBuf>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -33,6 +35,7 @@ impl Default for NodeConfig {
             bootstrap_nodes: vec![],
             key_seed: None,
             block_time_secs: 10, // Default to 10 seconds for testing
+            data_dir: None, // Default to in-memory storage for testing
         }
     }
 }
diff --git a/crates/bitcell-node/src/main.rs b/crates/bitcell-node/src/main.rs
@@ -1,7 +1,6 @@
 //! BitCell node binary
 
 use bitcell_node::{NodeConfig, ValidatorNode, MinerNode};
-use bitcell_crypto::SecretKey;
 use clap::{Parser, Subcommand};
 use std::path::PathBuf;
 
@@ -81,18 +80,18 @@ async fn main() {
     let cli = Cli::parse();
 
     match cli.command {
-        Commands::Validator { port, rpc_port, data_dir: _, enable_dht, bootstrap, key_seed, key_file, private_key } => {
+        Commands::Validator { port, rpc_port, data_dir, enable_dht, bootstrap, key_seed, key_file, private_key } => {
             println!("🌌 BitCell Validator Node");
             println!("=========================");
             
             let mut config = NodeConfig::default();
             config.network_port = port;
             config.enable_dht = enable_dht;
             config.key_seed = key_seed.clone();
+            config.data_dir = data_dir;
             if let Some(bootstrap_node) = bootstrap {
                 config.bootstrap_nodes.push(bootstrap_node);
             }
-            // TODO: Use data_dir
             
             // Resolve secret key
             let secret_key = match bitcell_node::keys::resolve_secret_key(
@@ -122,7 +121,13 @@ async fn main() {
             // Or we can modify NodeConfig to hold the secret key? No, NodeConfig is serializable.
             
             // Let's update ValidatorNode::new to take the secret key as an argument.
-            let mut node = ValidatorNode::with_key(config, secret_key.clone());
+            let mut node = match ValidatorNode::with_key(config, secret_key.clone()) {
+                Ok(node) => node,
+                Err(e) => {
+                    eprintln!("Error initializing validator node: {}", e);
+                    std::process::exit(1);
+                }
+            };
             
             // Start metrics server on port + 2 to avoid conflict with P2P port (30333) and RPC port (30334)
             let metrics_port = port + 2;
@@ -162,14 +167,15 @@ async fn main() {
             tokio::signal::ctrl_c().await.expect("Failed to listen for Ctrl+C");
             println!("\nShutting down...");
         }
-        Commands::Miner { port, rpc_port, data_dir: _, enable_dht, bootstrap, key_seed, key_file, private_key } => {
+        Commands::Miner { port, rpc_port, data_dir, enable_dht, bootstrap, key_seed, key_file, private_key } => {
             println!("⛏️  BitCell Miner Node");
             println!("======================");
             
             let mut config = NodeConfig::default();
             config.network_port = port;
             config.enable_dht = enable_dht;
             config.key_seed = key_seed.clone();
+            config.data_dir = data_dir;
             if let Some(bootstrap_node) = bootstrap {
                 config.bootstrap_nodes.push(bootstrap_node);
             }
@@ -190,7 +196,13 @@ async fn main() {
             
             println!("Miner Public Key: {:?}", secret_key.public_key());
             
-            let mut node = MinerNode::with_key(config, secret_key.clone());
+            let mut node = match MinerNode::with_key(config, secret_key.clone()) {
+                Ok(node) => node,
+                Err(e) => {
+                    eprintln!("Error initializing miner node: {}", e);
+                    std::process::exit(1);
+                }
+            };
             
             let metrics_port = port + 2;
 
@@ -228,14 +240,15 @@ async fn main() {
             tokio::signal::ctrl_c().await.expect("Failed to listen for Ctrl+C");
             println!("\nShutting down...");
         }
-        Commands::FullNode { port, rpc_port, data_dir: _, enable_dht, bootstrap, key_seed, key_file, private_key } => {
+        Commands::FullNode { port, rpc_port, data_dir, enable_dht, bootstrap, key_seed, key_file, private_key } => {
             println!("🌍 BitCell Full Node");
             println!("====================");
             
             let mut config = NodeConfig::default();
             config.network_port = port;
             config.enable_dht = enable_dht;
             config.key_seed = key_seed.clone();
+            config.data_dir = data_dir;
             if let Some(bootstrap_node) = bootstrap {
                 config.bootstrap_nodes.push(bootstrap_node);
             }
@@ -257,7 +270,13 @@ async fn main() {
             println!("Full Node Public Key: {:?}", secret_key.public_key());
 
             // Reuse ValidatorNode for now as FullNode logic is similar (just no voting)
-            let mut node = ValidatorNode::with_key(config, secret_key.clone());
+            let mut node = match ValidatorNode::with_key(config, secret_key.clone()) {
+                Ok(node) => node,
+                Err(e) => {
+                    eprintln!("Error initializing full node: {}", e);
+                    std::process::exit(1);
+                }
+            };
             
             let metrics_port = port + 2;
 
diff --git a/crates/bitcell-node/src/miner.rs b/crates/bitcell-node/src/miner.rs
@@ -1,17 +1,14 @@
-///! Miner node implementation
-
+//! Miner node implementation
 use crate::{NodeConfig, Result, MetricsRegistry, Blockchain, TransactionPool, NetworkManager};
 use bitcell_crypto::SecretKey;
 use bitcell_ca::{Glider, GliderPattern};
-use bitcell_state::StateManager;
 use std::sync::Arc;
 use bitcell_consensus::Transaction;
 
 /// Miner node
 pub struct MinerNode {
     pub config: NodeConfig,
     pub secret_key: Arc<SecretKey>,
-    pub state: StateManager,
     pub glider_strategy: GliderPattern,
     pub metrics: MetricsRegistry,
     pub blockchain: Blockchain,
@@ -20,25 +17,38 @@ pub struct MinerNode {
 }
 
 impl MinerNode {
-    pub fn new(config: NodeConfig, secret_key: SecretKey) -> Self {
+    pub fn new(config: NodeConfig, secret_key: SecretKey) -> crate::Result<Self> {
         Self::with_key(config, Arc::new(secret_key))
     }
 
-    pub fn with_key(config: NodeConfig, secret_key: Arc<SecretKey>) -> Self {
+    pub fn with_key(config: NodeConfig, secret_key: Arc<SecretKey>) -> crate::Result<Self> {
         let metrics = MetricsRegistry::new();
-        let blockchain = Blockchain::new(secret_key.clone(), metrics.clone());
+        
+        // Create blockchain with or without persistent storage based on config
+        let blockchain = if let Some(ref data_path) = config.data_dir {
+            // Ensure data directory exists
+            std::fs::create_dir_all(data_path)
+                .map_err(|e| crate::Error::Config(format!("Failed to create data directory: {}", e)))?;
+            
+            println!("📦 Using persistent storage at: {}", data_path.display());
+            Blockchain::with_storage(secret_key.clone(), metrics.clone(), data_path)
+                .map_err(|e| crate::Error::Config(format!("Failed to initialize blockchain with storage: {}", e)))?
+        } else {
+            println!("⚠️  Using in-memory storage (data will not persist)");
+            Blockchain::new(secret_key.clone(), metrics.clone())
+        };
+        
         let network = Arc::new(NetworkManager::new(secret_key.public_key(), metrics.clone()));
         
-        Self {
+        Ok(Self {
             config,
             secret_key,
-            state: StateManager::new(),
             glider_strategy: GliderPattern::Standard,
             metrics,
             blockchain,
             tx_pool: TransactionPool::default(),
             network,
-        }
+        })
     }
 
     pub async fn start(&mut self) -> Result<()> {
@@ -162,15 +172,15 @@ mod tests {
     fn test_miner_creation() {
         let config = NodeConfig::default();
         let sk = SecretKey::generate();
-        let miner = MinerNode::new(config, sk);
+        let miner = MinerNode::new(config, sk).unwrap();
         assert_eq!(miner.glider_strategy, GliderPattern::Standard);
     }
 
     #[test]
     fn test_glider_generation() {
         let config = NodeConfig::default();
         let sk = SecretKey::generate();
-        let miner = MinerNode::new(config, sk);
+        let miner = MinerNode::new(config, sk).unwrap();
         let glider = miner.generate_glider();
         assert_eq!(glider.pattern, GliderPattern::Standard);
     }
diff --git a/crates/bitcell-node/src/validator.rs b/crates/bitcell-node/src/validator.rs
@@ -2,7 +2,6 @@
 
 use crate::{NodeConfig, Result, MetricsRegistry, Blockchain, TransactionPool};
 use bitcell_consensus::Block;
-use bitcell_state::StateManager;
 use bitcell_network::PeerManager;
 use bitcell_crypto::SecretKey;
 use std::sync::Arc;
@@ -15,7 +14,6 @@ const MAX_TXS_PER_BLOCK: usize = 1000;
 /// Validator node
 pub struct ValidatorNode {
     pub config: NodeConfig,
-    pub state: StateManager,
     pub peers: PeerManager,
     pub metrics: MetricsRegistry,
     pub blockchain: Blockchain,
@@ -26,7 +24,7 @@ pub struct ValidatorNode {
 }
 
 impl ValidatorNode {
-    pub fn new(config: NodeConfig) -> Self {
+    pub fn new(config: NodeConfig) -> crate::Result<Self> {
         let secret_key = if let Some(seed) = &config.key_seed {
             println!("Generating validator key from seed: {}", seed);
             let hash = bitcell_crypto::Hash256::hash(seed.as_bytes());
@@ -37,23 +35,36 @@ impl ValidatorNode {
         Self::with_key(config, secret_key)
     }
 
-    pub fn with_key(config: NodeConfig, secret_key: Arc<SecretKey>) -> Self {
+    pub fn with_key(config: NodeConfig, secret_key: Arc<SecretKey>) -> crate::Result<Self> {
         let metrics = MetricsRegistry::new();
-        let blockchain = Blockchain::new(secret_key.clone(), metrics.clone());
+        
+        // Create blockchain with or without persistent storage based on config
+        let blockchain = if let Some(ref data_path) = config.data_dir {
+            // Ensure data directory exists
+            std::fs::create_dir_all(data_path)
+                .map_err(|e| crate::Error::Config(format!("Failed to create data directory: {}", e)))?;
+            
+            println!("📦 Using persistent storage at: {}", data_path.display());
+            Blockchain::with_storage(secret_key.clone(), metrics.clone(), data_path)
+                .map_err(|e| crate::Error::Config(format!("Failed to initialize blockchain with storage: {}", e)))?
+        } else {
+            println!("⚠️  Using in-memory storage (data will not persist)");
+            Blockchain::new(secret_key.clone(), metrics.clone())
+        };
+        
         let tournament_manager = Arc::new(crate::tournament::TournamentManager::new(metrics.clone()));
         let network = Arc::new(crate::network::NetworkManager::new(secret_key.public_key(), metrics.clone()));
         
-        Self {
+        Ok(Self {
             config,
-            state: StateManager::new(),
             peers: PeerManager::new(),
             metrics,
             blockchain,
             tx_pool: TransactionPool::default(),
             secret_key,
             tournament_manager,
             network,
-        }
+        })
     }
 
     pub async fn start(&mut self) -> Result<()> {
@@ -276,7 +287,9 @@ mod tests {
     #[test]
     fn test_validator_creation() {
         let config = NodeConfig::default();
-        let node = ValidatorNode::new(config);
-        assert_eq!(node.state.accounts.len(), 0);
+        let node = ValidatorNode::new(config).unwrap();
+        let state = node.blockchain.state();
+        let state_guard = state.read().unwrap();
+        assert_eq!(state_guard.accounts.len(), 0);
     }
 }
diff --git a/crates/bitcell-state/src/lib.rs b/crates/bitcell-state/src/lib.rs
diff --git a/crates/bitcell-state/src/storage.rs b/crates/bitcell-state/src/storage.rs

Original file line number	Diff line number	Diff line change
`@@ -39,3 +39,4 @@ base64 = "0.21"`
`39`	`39`
`40`	`40`	`[dev-dependencies]`
`41`	`41`	`proptest.workspace = true`
	`42`	`+tempfile = "3.23.0"`
Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,8 @@ pub struct NodeConfig {`
`14`	`14`	`/// Block production interval in seconds.`
`15`	`15`	`/// Defaults to 10 seconds for testing. Use 600 (10 minutes) for production.`
`16`	`16`	`pub block_time_secs: u64,`
	`17`	`+ /// Data directory for persistent storage. If None, uses in-memory storage only.`
	`18`	`+ pub data_dir: Option<std::path::PathBuf>,`
`17`	`19`	`}`
`18`	`20`
`19`	`21`	`#[derive(Debug, Clone, Serialize, Deserialize)]`
`@@ -33,6 +35,7 @@ impl Default for NodeConfig {`
`33`	`35`	`bootstrap_nodes: vec![],`
`34`	`36`	`key_seed: None,`
`35`	`37`	`block_time_secs: 10, // Default to 10 seconds for testing`
	`38`	`+ data_dir: None, // Default to in-memory storage for testing`
`36`	`39`	`}`
`37`	`40`	`}`
`38`	`41`	`}`