Create defrag-style TUI for wallet syncing

This commit is contained in:
Jack Grigg 2024-04-20 21:39:38 +00:00
parent 94e1337b46
commit 1dcd4b28b6
4 changed files with 436 additions and 0 deletions

100
Cargo.lock generated
View File

@ -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",

View File

@ -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",
]

View File

@ -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

331
src/commands/sync/defrag.rs Normal file
View File

@ -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<Action>,
}
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<Range<BlockHeight>>) -> 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<Range<BlockHeight>>) -> 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<BlockHeight, ScanPriority>,
fetching_range: Option<Range<BlockHeight>>,
scanning_range: Option<Range<BlockHeight>>,
action_tx: mpsc::UnboundedSender<Action>,
action_rx: mpsc::UnboundedReceiver<Action>,
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<ScanRange>, 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<ScanPriority>, (_, &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<ScanRange>,
chain_tip: BlockHeight,
},
SetFetching(Option<Range<BlockHeight>>),
SetScanning(Option<Range<BlockHeight>>),
Render,
}
impl Action {
fn for_event(event: tui::Event) -> Option<Self> {
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,
}
}
}