Merge pull request #174 from dashpay/feat/enhanced-storage-clear-and-balance-display

feat: enhance clear_storage and add wallet balance to status display

Author: PastaPastaPasta

dash-spv/src/client/core.rs

@@ -210,10 +210,68 @@ impl<
 
     // ============ Storage Operations ============
 
-    /// Clear all persisted storage (headers, filters, state, sync state).
+    /// Clear all persisted storage (headers, filters, state, sync state) and reset in-memory state.
     pub async fn clear_storage(&mut self) -> Result<()> {
-        let mut storage = self.storage.lock().await;
-        storage.clear().await.map_err(SpvError::Storage)
+        // Wipe on-disk persistence fully
+        {
+            let mut storage = self.storage.lock().await;
+            storage.clear().await.map_err(SpvError::Storage)?;
+        }
+
+        // Reset in-memory chain state to a clean baseline for the current network
+        {
+            let mut state = self.state.write().await;
+            *state = ChainState::new_for_network(self.config.network);
+        }
+
+        // Reset sync manager filter state (headers/filters progress trackers)
+        self.sync_manager.filter_sync_mut().clear_filter_state().await;
+
+        // Reset in-memory statistics and received filter height tracking without
+        // replacing the SharedFilterHeights Arc (to keep existing references valid)
+        let received_heights = {
+            let stats = self.stats.read().await;
+            stats.received_filter_heights.clone()
+        };
+
+        {
+            use std::time::Duration;
+            let mut stats = self.stats.write().await;
+            stats.connected_peers = 0;
+            stats.total_peers = 0;
+            stats.header_height = 0;
+            stats.filter_height = 0;
+            stats.headers_downloaded = 0;
+            stats.filter_headers_downloaded = 0;
+            stats.filters_downloaded = 0;
+            stats.filters_matched = 0;
+            stats.blocks_with_relevant_transactions = 0;
+            stats.blocks_requested = 0;
+            stats.blocks_processed = 0;
+            stats.masternode_diffs_processed = 0;
+            stats.bytes_received = 0;
+            stats.bytes_sent = 0;
+            stats.uptime = Duration::default();
+            stats.filters_requested = 0;
+            stats.filters_received = 0;
+            stats.filter_sync_start_time = None;
+            stats.last_filter_received_time = None;
+            stats.active_filter_requests = 0;
+            stats.pending_filter_requests = 0;
+            stats.filter_request_timeouts = 0;
+            stats.filter_requests_retried = 0;
+        }
+
+        received_heights.lock().await.clear();
+
+        // Reset mempool tracking (state and bloom filter)
+        {
+            let mut mempool_state = self.mempool_state.write().await;
+            *mempool_state = MempoolState::default();
+        }
+        self.mempool_filter = None;
+
+        Ok(())
     }
 
     /// Clear only the persisted sync state snapshot (keep headers/filters).
@@ -295,20 +353,28 @@ impl<
 
     /// Helper to create a StatusDisplay instance.
     #[cfg(feature = "terminal-ui")]
-    pub(super) async fn create_status_display(&self) -> StatusDisplay<'_, S> {
+    pub(super) async fn create_status_display(&self) -> StatusDisplay<'_, S, W> {
         StatusDisplay::new(
             &self.state,
             &self.stats,
             self.storage.clone(),
+            Some(&self.wallet),
             &self.terminal_ui,
             &self.config,
         )
     }
 
     /// Helper to create a StatusDisplay instance (without terminal UI).
     #[cfg(not(feature = "terminal-ui"))]
-    pub(super) async fn create_status_display(&self) -> StatusDisplay<'_, S> {
-        StatusDisplay::new(&self.state, &self.stats, self.storage.clone(), &None, &self.config)
+    pub(super) async fn create_status_display(&self) -> StatusDisplay<'_, S, W> {
+        StatusDisplay::new(
+            &self.state,
+            &self.stats,
+            self.storage.clone(),
+            Some(&self.wallet),
+            &None,
+            &self.config,
+        )
     }
 
     /// Update the status display.

dash-spv/src/client/status_display.rs

@@ -9,32 +9,38 @@ use crate::storage::StorageManager;
 #[cfg(feature = "terminal-ui")]
 use crate::terminal::TerminalUI;
 use crate::types::{ChainState, SpvStats, SyncProgress};
+use key_wallet_manager::wallet_interface::WalletInterface;
 
 /// Status display manager for updating UI and reporting sync progress.
-pub struct StatusDisplay<'a, S: StorageManager> {
+pub struct StatusDisplay<'a, S: StorageManager, W: WalletInterface> {
     state: &'a Arc<RwLock<ChainState>>,
     stats: &'a Arc<RwLock<SpvStats>>,
     storage: Arc<Mutex<S>>,
+    wallet: Option<&'a Arc<RwLock<W>>>,
     #[cfg(feature = "terminal-ui")]
     terminal_ui: &'a Option<Arc<TerminalUI>>,
     #[allow(dead_code)]
     config: &'a ClientConfig,
 }
 
-impl<'a, S: StorageManager + Send + Sync + 'static> StatusDisplay<'a, S> {
+impl<'a, S: StorageManager + Send + Sync + 'static, W: WalletInterface + Send + Sync + 'static>
+    StatusDisplay<'a, S, W>
+{
     /// Create a new status display manager.
     #[cfg(feature = "terminal-ui")]
     pub fn new(
         state: &'a Arc<RwLock<ChainState>>,
         stats: &'a Arc<RwLock<SpvStats>>,
         storage: Arc<Mutex<S>>,
+        wallet: Option<&'a Arc<RwLock<W>>>,
         terminal_ui: &'a Option<Arc<TerminalUI>>,
         config: &'a ClientConfig,
     ) -> Self {
         Self {
             state,
             stats,
             storage,
+            wallet,
             terminal_ui,
             config,
         }
@@ -46,13 +52,15 @@ impl<'a, S: StorageManager + Send + Sync + 'static> StatusDisplay<'a, S> {
         state: &'a Arc<RwLock<ChainState>>,
         stats: &'a Arc<RwLock<SpvStats>>,
         storage: Arc<Mutex<S>>,
+        wallet: Option<&'a Arc<RwLock<W>>>,
         _terminal_ui: &'a Option<()>,
         config: &'a ClientConfig,
     ) -> Self {
         Self {
             state,
             stats,
             storage,
+            wallet,
             config,
         }
     }
@@ -140,6 +148,42 @@ impl<'a, S: StorageManager + Send + Sync + 'static> StatusDisplay<'a, S> {
         state.clone()
     }
 
+    /// Helper to try to get wallet balance if W implements Any.
+    /// This is a wrapper that handles the case where W might not implement Any.
+    fn try_get_balance_if_any(wallet: &W) -> Option<u64>
+    where
+        W: 'static,
+    {
+        // Try to use Any trait for downcasting
+        // We check if W is WalletManager<ManagedWalletInfo> using TypeId
+        use key_wallet::wallet::managed_wallet_info::ManagedWalletInfo;
+        use key_wallet_manager::wallet_manager::WalletManager;
+        use std::any::TypeId;
+
+        // Check if W is WalletManager<ManagedWalletInfo>
+        let wallet_type_id = TypeId::of::<W>();
+        let wallet_manager_type_id = TypeId::of::<WalletManager<ManagedWalletInfo>>();
+
+        if wallet_type_id == wallet_manager_type_id {
+            // Unsafe downcast: we've verified the types match, so this is safe
+            unsafe {
+                let wallet_ptr = wallet as *const W as *const WalletManager<ManagedWalletInfo>;
+                let wallet_ref = &*wallet_ptr;
+                return Some(wallet_ref.get_total_balance());
+            }
+        }
+
+        None
+    }
+
+    /// Format balance in DASH with 8 decimal places.
+    fn format_balance(satoshis: u64) -> String {
+        use dashcore::Amount;
+        use dashcore::Denomination;
+        let amount = Amount::from_sat(satoshis);
+        amount.to_string_with_denomination(Denomination::Dash)
+    }
+
     /// Update the status display.
     pub async fn update_status_display(&self) {
         #[cfg(feature = "terminal-ui")]
@@ -219,23 +263,39 @@ impl<'a, S: StorageManager + Send + Sync + 'static> StatusDisplay<'a, S> {
 
             // Get filter and block processing statistics
             let stats = self.stats.read().await;
+            let filters_received = stats.filters_received;
             let filters_matched = stats.filters_matched;
             let blocks_with_relevant_transactions = stats.blocks_with_relevant_transactions;
             let blocks_processed = stats.blocks_processed;
             drop(stats);
 
+            // Get wallet balance if available
+            let balance_str = if let Some(wallet_ref) = self.wallet {
+                let wallet_guard = wallet_ref.read().await;
+                // Try to get balance if W implements Any (for WalletManager support)
+                // We use a helper that requires W: Any, so we need to handle this carefully
+                // For now, we'll attempt to get balance only if possible
+                Self::try_get_balance_if_any(&*wallet_guard)
+                    .map(|balance_sat| format!(" | Balance: {}", Self::format_balance(balance_sat)))
+                    .unwrap_or_default()
+            } else {
+                String::new()
+            };
+
             tracing::info!(
-                "📊 [SYNC STATUS] Headers: {} | Filter Headers: {} | Latest ChainLock: {} | Filters Matched: {} | Blocks w/ Relevant Txs: {} | Blocks Processed: {}",
+                "📊 [SYNC STATUS] Headers: {} | Filter Headers: {} | Filters: {} | Latest ChainLock: {} | Filters Matched: {} | Blocks w/ Relevant Txs: {} | Blocks Processed: {}{}",
                 header_height,
                 filter_height,
+                filters_received,
                 if chainlock_height > 0 {
                     format!("#{}", chainlock_height)
                 } else {
                     "None".to_string()
                 },
                 filters_matched,
                 blocks_with_relevant_transactions,
-                blocks_processed
+                blocks_processed,
+                balance_str
             );
         }
     }