From 1dcd4b28b68a4a1b5ca19b13da119abc556806a7 Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Sat, 20 Apr 2024 21:39:38 +0000 Subject: [PATCH] Create defrag-style TUI for wallet syncing --- Cargo.lock | 100 +++++++++++ Cargo.toml | 2 + src/commands/sync.rs | 3 + src/commands/sync/defrag.rs | 331 ++++++++++++++++++++++++++++++++++++ 4 files changed, 436 insertions(+) create mode 100644 src/commands/sync/defrag.rs diff --git a/Cargo.lock b/Cargo.lock index 2195682..4558189 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -65,6 +65,21 @@ version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f" +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anyhow" version = "1.0.82" @@ -397,6 +412,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "chrono" +version = "0.4.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "num-traits", + "windows-targets 0.52.5", +] + [[package]] name = "cipher" version = "0.4.4" @@ -433,6 +460,12 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f7144d30dcf0fafbce74250a3963025d8d52177934239851c917d29f1df280c2" +[[package]] +name = "core-foundation-sys" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f" + [[package]] name = "cpufeatures" version = "0.2.12" @@ -735,6 +768,15 @@ dependencies = [ "slab", ] +[[package]] +name = "fxhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" +dependencies = [ + "byteorder", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -990,6 +1032,29 @@ dependencies = [ "tokio-io-timeout", ] +[[package]] +name = "iana-time-zone" +version = "0.1.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7ffbb5a1b541ea2561f8c41c087286cc091e21e556a4f09a8f6cbf17b69b141" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "incrementalmerkletree" version = "0.5.1" @@ -1049,6 +1114,15 @@ version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +[[package]] +name = "js-sys" +version = "0.3.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29c15563dc2726973df627357ce0c9ddddbea194836909d655df6a75d2cf296d" +dependencies = [ + "wasm-bindgen", +] + [[package]] name = "jubjub" version = "0.10.0" @@ -2566,6 +2640,22 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tui-logger" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4358d7a45f901c23c4e43e0885c159f035b2ca3a90e646f4d1dbae80b45a6c79" +dependencies = [ + "chrono", + "fxhash", + "lazy_static", + "log", + "parking_lot", + "ratatui", + "tracing", + "tracing-subscriber", +] + [[package]] name = "typenum" version = "1.17.0" @@ -2760,6 +2850,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets 0.52.5", +] + [[package]] name = "windows-sys" version = "0.48.0" @@ -3155,6 +3254,7 @@ dependencies = [ "tonic", "tracing", "tracing-subscriber", + "tui-logger", "zcash_client_backend", "zcash_client_sqlite", "zcash_keys", diff --git a/Cargo.toml b/Cargo.toml index 3963c8b..379590d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ zcash_protocol = "0.1" crossterm = { version = "0.27", optional = true, features = ["event-stream"] } ratatui = { version = "0.26", optional = true } tokio-util = { version = "0.7", optional = true } +tui-logger = { version = "0.11", optional = true, features = ["tracing-support"] } [features] default = [] @@ -40,4 +41,5 @@ tui = [ "dep:crossterm", "dep:ratatui", "dep:tokio-util", + "dep:tui-logger", ] diff --git a/src/commands/sync.rs b/src/commands/sync.rs index b03f32a..70d07f9 100644 --- a/src/commands/sync.rs +++ b/src/commands/sync.rs @@ -30,6 +30,9 @@ use crate::{ remote::connect_to_lightwalletd, }; +#[cfg(feature = "tui")] +mod defrag; + const BATCH_SIZE: u32 = 10_000; // Options accepted for the `sync` command diff --git a/src/commands/sync/defrag.rs b/src/commands/sync/defrag.rs new file mode 100644 index 0000000..ea04447 --- /dev/null +++ b/src/commands/sync/defrag.rs @@ -0,0 +1,331 @@ +use std::{collections::BTreeMap, ops::Range}; + +use crossterm::event::KeyCode; +use futures_util::FutureExt; +use ratatui::{ + prelude::*, + widgets::{Block, Paragraph}, +}; +use tokio::sync::mpsc; +use tracing::{error, info}; +use tui_logger::{TuiLoggerLevelOutput, TuiLoggerSmartWidget}; +use zcash_client_backend::data_api::scanning::{ScanPriority, ScanRange}; +use zcash_protocol::consensus::BlockHeight; + +use crate::tui; + +pub(super) struct AppHandle { + action_tx: mpsc::UnboundedSender, +} + +impl AppHandle { + /// Returns `true` if the TUI exited. + pub(super) fn set_scan_ranges( + &self, + scan_ranges: &[ScanRange], + chain_tip: BlockHeight, + ) -> bool { + match self.action_tx.send(Action::UpdateScanRanges { + scan_ranges: scan_ranges.to_vec(), + chain_tip, + }) { + Ok(()) => false, + Err(e) => { + error!("Failed to send: {}", e); + true + } + } + } + + /// Returns `true` if the TUI exited. + pub(super) fn set_fetching_range(&self, fetching_range: Option>) -> bool { + match self.action_tx.send(Action::SetFetching(fetching_range)) { + Ok(()) => false, + Err(e) => { + error!("Failed to send: {}", e); + true + } + } + } + + /// Returns `true` if the TUI exited. + pub(super) fn set_scanning_range(&self, scanning_range: Option>) -> bool { + match self.action_tx.send(Action::SetScanning(scanning_range)) { + Ok(()) => false, + Err(e) => { + error!("Failed to send: {}", e); + true + } + } + } +} + +pub(super) struct App { + should_quit: bool, + wallet_birthday: BlockHeight, + scan_ranges: BTreeMap, + fetching_range: Option>, + scanning_range: Option>, + action_tx: mpsc::UnboundedSender, + action_rx: mpsc::UnboundedReceiver, + logger_state: tui_logger::TuiWidgetState, +} + +impl App { + pub(super) fn new(wallet_birthday: BlockHeight) -> Self { + let (action_tx, action_rx) = mpsc::unbounded_channel(); + Self { + should_quit: false, + wallet_birthday, + scan_ranges: BTreeMap::new(), + fetching_range: None, + scanning_range: None, + action_tx, + action_rx, + logger_state: tui_logger::TuiWidgetState::new(), + } + } + + pub(super) fn handle(&self) -> AppHandle { + AppHandle { + action_tx: self.action_tx.clone(), + } + } + + pub(super) async fn run(&mut self, mut tui: tui::Tui) -> anyhow::Result<()> { + tui.enter()?; + + loop { + let next_event = tui.next().fuse(); + let next_action = self.action_rx.recv().fuse(); + tokio::select! { + Some(event) = next_event => if let Some(action) = Action::for_event(event) { + self.action_tx.send(action)?; + }, + Some(action) = next_action => match action { + Action::Quit => { + info!("Quit requested"); + self.should_quit = true; + break; + } + Action::Tick => {} + Action::LoggerEvent(event) => self.logger_state.transition(event), + Action::UpdateScanRanges { scan_ranges, chain_tip } => { + self.update_scan_ranges(scan_ranges, chain_tip); + } + Action::SetFetching(fetching_range) => self.fetching_range = fetching_range, + Action::SetScanning(scanning_range) => self.scanning_range = scanning_range, + Action::Render => { + tui.draw(|f| self.ui(f))?; + } + } + } + + if self.should_quit { + break; + } + } + + self.action_rx.close(); + tui.exit()?; + + Ok(()) + } + + fn update_scan_ranges(&mut self, mut scan_ranges: Vec, chain_tip: BlockHeight) { + scan_ranges.sort_by_key(|range| range.block_range().start); + + self.scan_ranges = scan_ranges + .into_iter() + .flat_map(|range| { + [ + (range.block_range().start, range.priority()), + // If this range is followed by an adjacent range, this will be + // overwritten. Otherwise, this is either a gap between unscanned + // ranges (which by definition is scanned), or the "mempool height" + // which we coerce down to the chain tip height. + ( + range.block_range().end.min(chain_tip), + ScanPriority::Scanned, + ), + ] + }) + .collect(); + + // If we weren't passed a ScanRange starting at the wallet birthday, it means we + // have scanned that height. + self.scan_ranges + .entry(self.wallet_birthday) + .or_insert(ScanPriority::Scanned); + + // If we inserted the chain tip height above, mark it as such. If we didn't insert + // it above, do so here. + self.scan_ranges + .entry(chain_tip) + .and_modify(|e| *e = ScanPriority::ChainTip) + .or_insert(ScanPriority::ChainTip); + } + + fn ui(&mut self, frame: &mut Frame) { + let [upper_area, log_area] = + Layout::vertical([Constraint::Min(0), Constraint::Length(15)]).areas(frame.size()); + + let defrag_area = { + let block = Block::bordered().title("Wallet Defragmentor"); + let inner_area = block.inner(upper_area); + frame.render_widget(block, upper_area); + inner_area + }; + + if let Some(block_count) = self + .scan_ranges + .last_key_value() + .map(|(&last, _)| u32::from(last - self.wallet_birthday)) + { + // Determine the density of blocks we will be rendering. + let blocks_per_cell = block_count / u32::from(defrag_area.area()); + let blocks_per_row = blocks_per_cell * u32::from(defrag_area.width); + + // Split the area into cells. + for i in 0..defrag_area.width { + for j in 0..defrag_area.height { + // Determine the priority of the cell. + let cell_start = self.wallet_birthday + + (blocks_per_row * u32::from(j)) + + (blocks_per_cell * u32::from(i)); + let cell_end = cell_start + blocks_per_cell; + + let (cell_text, cell_color) = if self + .fetching_range + .as_ref() + .map(|range| range.contains(&cell_start) || range.contains(&(cell_end - 1))) + .unwrap_or(false) + { + ("↓", Color::Magenta) + } else if self + .scanning_range + .as_ref() + .map(|range| range.contains(&cell_start) || range.contains(&(cell_end - 1))) + .unwrap_or(false) + { + ("@", Color::Magenta) + } else { + let cell_priority = self + .scan_ranges + .range(cell_start..cell_end) + .fold(None, |acc: Option, (_, &priority)| { + if let Some(acc) = acc { + Some(acc.max(priority)) + } else { + Some(priority) + } + }) + .or_else(|| { + self.scan_ranges + .range(..=cell_start) + .next_back() + .map(|(_, &priority)| priority) + }) + .or_else(|| { + self.scan_ranges + .range((cell_end - 1)..) + .next() + .map(|(_, &priority)| priority) + }) + .unwrap_or(ScanPriority::Ignored); + + ( + " ", + match cell_priority { + ScanPriority::Ignored => Color::Black, + ScanPriority::Scanned => Color::Green, + ScanPriority::Historic => Color::Black, + ScanPriority::OpenAdjacent => Color::LightBlue, + ScanPriority::FoundNote => Color::Yellow, + ScanPriority::ChainTip => Color::Blue, + ScanPriority::Verify => Color::Red, + }, + ) + }; + + frame.render_widget( + Paragraph::new(cell_text).bg(cell_color), + Rect::new(defrag_area.x + i, defrag_area.y + j, 1, 1), + ); + } + } + } + + frame.render_widget( + TuiLoggerSmartWidget::default() + .style_error(Style::default().fg(Color::Red)) + .style_debug(Style::default().fg(Color::Green)) + .style_warn(Style::default().fg(Color::Yellow)) + .style_trace(Style::default().fg(Color::Magenta)) + .style_info(Style::default().fg(Color::Cyan)) + .output_separator(':') + .output_timestamp(Some("%H:%M:%S".to_string())) + .output_level(Some(TuiLoggerLevelOutput::Abbreviated)) + .output_target(true) + .output_file(true) + .output_line(true) + .state(&self.logger_state), + log_area, + ); + } +} + +#[derive(Clone, Debug)] +pub(super) enum Action { + Quit, + Tick, + LoggerEvent(tui_logger::TuiWidgetEvent), + UpdateScanRanges { + scan_ranges: Vec, + chain_tip: BlockHeight, + }, + SetFetching(Option>), + SetScanning(Option>), + Render, +} + +impl Action { + fn for_event(event: tui::Event) -> Option { + match event { + tui::Event::Error => None, + tui::Event::Tick => Some(Action::Tick), + tui::Event::Render => Some(Action::Render), + tui::Event::Key(key) => match key.code { + KeyCode::Char('q') => Some(Action::Quit), + KeyCode::Char(' ') => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::SpaceKey)) + } + KeyCode::Up => Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::UpKey)), + KeyCode::Down => Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::DownKey)), + KeyCode::Left => Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::LeftKey)), + KeyCode::Right => Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::RightKey)), + KeyCode::Char('+') => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::PlusKey)) + } + KeyCode::Char('-') => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::MinusKey)) + } + KeyCode::Char('h') => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::HideKey)) + } + KeyCode::Char('f') => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::FocusKey)) + } + KeyCode::PageUp => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::PrevPageKey)) + } + KeyCode::PageDown => { + Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::NextPageKey)) + } + KeyCode::Esc => Some(Action::LoggerEvent(tui_logger::TuiWidgetEvent::EscapeKey)), + _ => None, + }, + _ => None, + } + } +}