use std::{cmp::min, path::PathBuf};
use abscissa_core::{config, Command, FrameworkError};
use color_eyre::eyre::{eyre, Report};
use tokio::time::Instant;
use tower::{Service, ServiceExt};
use zebra_chain::{block::Height, parameters::Network};
use zebra_state as old_zs;
use zebra_state as new_zs;
use crate::{
application::ZebradApp,
components::tokio::{RuntimeRun, TokioComponent},
config::ZebradConfig,
prelude::*,
BoxError,
};
const PROGRESS_HEIGHT_INTERVAL: u32 = 5_000;
#[derive(Command, Debug, clap::Parser)]
pub struct CopyStateCmd {
#[clap(long, short, help = "stop copying at this source height")]
max_source_height: Option<u32>,
#[clap(
long,
short,
help = "config file path for the target state (default: ephemeral), \
the source state uses the main zebrad config"
)]
target_config_path: Option<PathBuf>,
#[clap(help = "tracing filters which override the zebrad.toml config")]
filters: Vec<String>,
}
impl CopyStateCmd {
async fn start(&self) -> Result<(), Report> {
let base_config = APPLICATION.config();
let source_config = base_config.state.clone();
let target_config = self
.target_config_path
.as_ref()
.map(|path| ZebradApp::default().load_config(path))
.transpose()?
.map(|app_config| app_config.state)
.unwrap_or_else(new_zs::Config::ephemeral);
info!(?base_config, "state copy base config");
self.copy(&base_config.network.network, source_config, target_config)
.await
.map_err(|e| eyre!(e))
}
async fn copy(
&self,
network: &Network,
source_config: old_zs::Config,
target_config: new_zs::Config,
) -> Result<(), BoxError> {
info!(
?source_config,
"initializing source state service (old format)"
);
let source_start_time = Instant::now();
let (
mut source_state,
_source_read_only_state_service,
_source_latest_chain_tip,
_source_chain_tip_change,
) = old_zs::spawn_init(source_config.clone(), network, Height::MAX, 0).await?;
let elapsed = source_start_time.elapsed();
info!(?elapsed, "finished initializing source state service");
info!(
?target_config, target_config_path = ?self.target_config_path,
"initializing target state service (new format)"
);
let target_start_time = Instant::now();
let (
mut target_state,
_target_read_only_state_service,
_target_latest_chain_tip,
_target_chain_tip_change,
) = new_zs::spawn_init(target_config.clone(), network, Height::MAX, 0).await?;
let elapsed = target_start_time.elapsed();
info!(?elapsed, "finished initializing target state service");
info!("fetching source and target tip heights");
let source_tip = source_state
.ready()
.await?
.call(old_zs::Request::Tip)
.await?;
let source_tip = match source_tip {
old_zs::Response::Tip(Some(source_tip)) => source_tip,
old_zs::Response::Tip(None) => Err("empty source state: no blocks to copy")?,
response => Err(format!("unexpected response to Tip request: {response:?}",))?,
};
let source_tip_height = source_tip.0 .0;
let initial_target_tip = target_state
.ready()
.await?
.call(new_zs::Request::Tip)
.await?;
let initial_target_tip = match initial_target_tip {
new_zs::Response::Tip(target_tip) => target_tip,
response => Err(format!("unexpected response to Tip request: {response:?}",))?,
};
let min_target_height = initial_target_tip
.map(|target_tip| target_tip.0 .0 + 1)
.unwrap_or(0);
let max_copy_height = self
.max_source_height
.map(|max_source_height| min(source_tip_height, max_source_height))
.unwrap_or(source_tip_height);
if min_target_height >= max_copy_height {
info!(
?min_target_height,
?max_copy_height,
max_source_height = ?self.max_source_height,
?source_tip,
?initial_target_tip,
"target is already at or after max copy height"
);
return Ok(());
}
info!(
?min_target_height,
?max_copy_height,
max_source_height = ?self.max_source_height,
?source_tip,
?initial_target_tip,
"starting copy from source to target"
);
let copy_start_time = Instant::now();
for height in min_target_height..=max_copy_height {
let source_block = source_state
.ready()
.await?
.call(old_zs::Request::Block(Height(height).into()))
.await?;
let source_block = match source_block {
old_zs::Response::Block(Some(source_block)) => {
trace!(?height, %source_block, "read source block");
source_block
}
old_zs::Response::Block(None) => {
Err(format!("unexpected missing source block, height: {height}",))?
}
response => Err(format!(
"unexpected response to Block request, height: {height}, \n \
response: {response:?}",
))?,
};
let source_block_hash = source_block.hash();
let target_block_commit_hash = target_state
.ready()
.await?
.call(new_zs::Request::CommitCheckpointVerifiedBlock(
source_block.clone().into(),
))
.await?;
let target_block_commit_hash = match target_block_commit_hash {
new_zs::Response::Committed(target_block_commit_hash) => {
trace!(?target_block_commit_hash, "wrote target block");
target_block_commit_hash
}
response => Err(format!(
"unexpected response to CommitCheckpointVerifiedBlock request, height: {height}\n \
response: {response:?}",
))?,
};
let target_block = target_state
.ready()
.await?
.call(new_zs::Request::Block(Height(height).into()))
.await?;
let target_block = match target_block {
new_zs::Response::Block(Some(target_block)) => {
trace!(?height, %target_block, "read target block");
target_block
}
new_zs::Response::Block(None) => {
Err(format!("unexpected missing target block, height: {height}",))?
}
response => Err(format!(
"unexpected response to Block request, height: {height},\n \
response: {response:?}",
))?,
};
let target_block_data_hash = target_block.hash();
if source_block_hash != target_block_commit_hash
|| source_block_hash != target_block_data_hash
|| source_block != target_block
{
Err(format!(
"unexpected mismatch between source and target blocks,\n \
max copy height: {max_copy_height:?},\n \
source hash: {source_block_hash:?},\n \
target commit hash: {target_block_commit_hash:?},\n \
target data hash: {target_block_data_hash:?},\n \
source block: {source_block:?},\n \
target block: {target_block:?}",
))?;
}
if height % PROGRESS_HEIGHT_INTERVAL == 0 {
let elapsed = copy_start_time.elapsed();
info!(
?height,
?max_copy_height,
?elapsed,
"copied block from source to target"
);
}
}
let elapsed = copy_start_time.elapsed();
info!(?max_copy_height, ?elapsed, "finished copying blocks");
info!(?max_copy_height, "fetching final target tip");
let final_target_tip = target_state
.ready()
.await?
.call(new_zs::Request::Tip)
.await?;
let final_target_tip = match final_target_tip {
new_zs::Response::Tip(Some(target_tip)) => target_tip,
new_zs::Response::Tip(None) => Err("empty target state: expected written blocks")?,
response => Err(format!("unexpected response to Tip request: {response:?}",))?,
};
let final_target_tip_height = final_target_tip.0 .0;
let final_target_tip_hash = final_target_tip.1;
let target_tip_source_depth = source_state
.ready()
.await?
.call(old_zs::Request::Depth(final_target_tip_hash))
.await?;
let target_tip_source_depth = match target_tip_source_depth {
old_zs::Response::Depth(source_depth) => source_depth,
response => Err(format!(
"unexpected response to Depth request: {response:?}",
))?,
};
if max_copy_height == source_tip_height {
let expected_target_depth = Some(0);
if source_tip != final_target_tip || target_tip_source_depth != expected_target_depth {
Err(format!(
"unexpected mismatch between source and target tips,\n \
max copy height: {max_copy_height:?},\n \
source tip: {source_tip:?},\n \
target tip: {final_target_tip:?},\n \
actual target tip depth in source: {target_tip_source_depth:?},\n \
expect target tip depth in source: {expected_target_depth:?}",
))?;
} else {
info!(
?max_copy_height,
?source_tip,
?final_target_tip,
?target_tip_source_depth,
"source and target states contain the same blocks"
);
}
} else {
let expected_target_depth = source_tip_height.checked_sub(final_target_tip_height);
if target_tip_source_depth != expected_target_depth {
Err(format!(
"unexpected mismatch between source and target tips,\n \
max copy height: {max_copy_height:?},\n \
source tip: {source_tip:?},\n \
target tip: {final_target_tip:?},\n \
actual target tip depth in source: {target_tip_source_depth:?},\n \
expect target tip depth in source: {expected_target_depth:?}",
))?;
} else {
info!(
?max_copy_height,
?source_tip,
?final_target_tip,
?target_tip_source_depth,
"target state reached the max copy height"
);
}
}
Ok(())
}
}
impl Runnable for CopyStateCmd {
fn run(&self) {
info!(
max_source_height = ?self.max_source_height,
target_config_path = ?self.target_config_path,
"starting cached chain state copy"
);
let rt = APPLICATION
.state()
.components_mut()
.get_downcast_mut::<TokioComponent>()
.expect("TokioComponent should be available")
.rt
.take();
rt.expect("runtime should not already be taken")
.run(self.start());
info!("finished cached chain state copy");
}
}
impl config::Override<ZebradConfig> for CopyStateCmd {
fn override_config(&self, mut config: ZebradConfig) -> Result<ZebradConfig, FrameworkError> {
if !self.filters.is_empty() {
config.tracing.filter = Some(self.filters.join(","));
}
Ok(config)
}
}