Skip to content

Commit

Permalink
split db download and formal download snapshot commands, cleanup inputs
Browse files Browse the repository at this point in the history
  • Loading branch information
@johnjmartin
johnjmartin committed Jan 10, 2024
1 parent ac1b2a2 commit d96cc9d
Show file tree
Hide file tree
Showing 2 changed files with 196 additions and 150 deletions.
254 changes: 177 additions & 77 deletions crates/sui-tool/src/commands.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -258,13 +258,8 @@ pub enum ToolCommand {
DownloadDBSnapshot {
#[clap(long = "epoch")]
epoch: u64,
#[clap(long = "genesis")]
genesis: PathBuf,
#[clap(long = "path", default_value = "/tmp")]
path: PathBuf,
/// skip downloading checkpoints dir. Overridden to `true` if `--formal` flag specified
#[clap(long = "skip-checkpoints")]
skip_checkpoints: bool,
/// skip downloading indexes dir. Overridden to `true` if `--formal` flag specified,
/// as index staging is not yet supported for formal snapshots.
#[clap(long = "skip-indexes")]
Expand All @@ -273,11 +268,44 @@ pub enum ToolCommand {
/// value based on number of available logical cores.
#[clap(long = "num-parallel-downloads")]
num_parallel_downloads: Option<usize>,
/// If true, restore from formal (slim, DB agnostic) snapshot. Note that this is only supported
/// for protocol versions supporting `commit_root_state_digest`. For mainnet, this is
/// epoch 20+, and for testnet this is epoch 12+
#[clap(long = "formal")]
formal: bool,
/// Network to download snapshot for. Defaults to "mainnet".
/// If `--snapshot-bucket` or `--archive-bucket` is not specified,
/// the value of this flag is used to construct default bucket names.
#[clap(long = "network", default_value = "mainnet")]
network: Chain,
/// Snapshot bucket name. If not specified, defaults are
/// based on value of `--network` and `--formal` flags.
#[clap(long = "snapshot-bucket")]
snapshot_bucket: Option<String>,
/// Snapshot bucket type. Defaults to "gcs" if `--formal`
/// flag specified, otherwise "s3".
#[clap(long = "snapshot-bucket-type")]
snapshot_bucket_type: Option<ObjectStoreType>,
/// Path to snapshot directory on local filesystem.
/// Only applicable if `--snapshot-bucket-type` is "file".
#[clap(long = "snapshot-path")]
snapshot_path: Option<PathBuf>,
/// If true, no authentication is needed for snapshot restores
#[clap(long = "no-sign-request")]
no_sign_request: bool,
/// If false (default), log level will be overridden to "off",
/// and output will be reduced to necessary status information.
#[clap(long = "verbose")]
verbose: bool,
},

#[clap(name = "download-formal-snapshot")]
DownloadFormalSnapshot {
#[clap(long = "epoch")]
epoch: u64,
#[clap(long = "genesis")]
genesis: PathBuf,
#[clap(long = "path", default_value = "/tmp")]
path: PathBuf,
/// Number of parallel downloads to perform. Defaults to a reasonable
/// value based on number of available logical cores.
#[clap(long = "num-parallel-downloads")]
num_parallel_downloads: Option<usize>,
/// If true, perform snapshot and checkpoint summary verification. Only
/// applicable if `--formal` flag is specified. Defaults to true.
#[clap(long = "verify")]
Expand Down Expand Up @@ -513,14 +541,11 @@ impl ToolCommand {
let config = sui_config::NodeConfig::load(config_path)?;
restore_from_db_checkpoint(&config, &db_checkpoint_path).await?;
}
ToolCommand::DownloadDBSnapshot {
ToolCommand::DownloadFormalSnapshot {
epoch,
genesis,
path,
skip_checkpoints,
skip_indexes,
num_parallel_downloads,
formal,
verify,
network,
snapshot_bucket,
Expand All @@ -542,55 +567,24 @@ impl ToolCommand {
.expect("Failed to get number of CPUs")
});
let snapshot_bucket =
snapshot_bucket.or_else(|| match (formal, network, no_sign_request) {
(true, Chain::Mainnet, false) => Some(
snapshot_bucket.or_else(|| match (network, no_sign_request) {
(Chain::Mainnet, false) => Some(
env::var("MAINNET_FORMAL_SIGNED_BUCKET")
.unwrap_or("mysten-mainnet-formal".to_string()),
),
(false, Chain::Mainnet, false) => Some(
env::var("MAINNET_DB_SIGNED_BUCKET")
.unwrap_or("mysten-mainnet-snapshots".to_string()),
),
(true, Chain::Mainnet, true) => {
env::var("MAINNET_FORMAL_UNSIGNED_BUCKET").ok()
}
(false, Chain::Mainnet, true) => {
env::var("MAINNET_DB_UNSIGNED_BUCKET").ok()
}
(true, Chain::Testnet, true) => {
env::var("TESTNET_FORMAL_UNSIGNED_BUCKET").ok()
}
(false, Chain::Testnet, true) => {
env::var("TESTNET_DB_UNSIGNED_BUCKET").ok()
}
(true, Chain::Testnet, _) => Some(
(Chain::Mainnet, true) => env::var("MAINNET_FORMAL_UNSIGNED_BUCKET").ok(),
(Chain::Testnet, true) => env::var("TESTNET_FORMAL_UNSIGNED_BUCKET").ok(),
(Chain::Testnet, _) => Some(
env::var("TESTNET_FORMAL_SIGNED_BUCKET")
.unwrap_or("mysten-testnet-formal".to_string()),
),
(false, Chain::Testnet, _) => Some(
env::var("TESTNET_DB_SIGNED_BUCKET")
.unwrap_or("mysten-testnet-snapshots".to_string()),
),
(_, Chain::Unknown, _) => {
(Chain::Unknown, _) => {
panic!("Cannot generate default snapshot bucket for unknown network");
}
});

let snapshot_bucket_type = snapshot_bucket_type.unwrap_or({
if formal {
ObjectStoreType::GCS
} else {
ObjectStoreType::S3
}
});

// index staging is not yet supported for formal snapshots
let skip_indexes = skip_indexes || formal;
// Checkpoint db does not exist in formal snapshots and
// is not reconstructed during formal snapshot restore
let skip_checkpoints = skip_checkpoints || formal;
let aws_endpoint = env::var("AWS_SNAPSHOT_ENDPOINT").ok().or_else(|| {
if formal && no_sign_request {
if no_sign_request {
if network == Chain::Mainnet {
Some("https://formal-snapshot.mainnet.sui.io".to_string())
} else if network == Chain::Testnet {
Expand All @@ -602,6 +596,13 @@ impl ToolCommand {
None
}
});

let snapshot_bucket_type = if no_sign_request {
ObjectStoreType::S3
} else {
snapshot_bucket_type
.expect("--snapshot-bucket-type must be set if not using --no-sign-request")
};
let snapshot_store_config = match snapshot_bucket_type {
ObjectStoreType::S3 => ObjectStoreConfig {
object_store: Some(ObjectStoreType::S3),
Expand All @@ -615,7 +616,7 @@ impl ToolCommand {
)
.ok()
.and_then(|b| b.parse().ok())
.unwrap_or(formal && no_sign_request),
.unwrap_or(no_sign_request),
object_store_connection_limit: 200,
no_sign_request,
..Default::default()
Expand Down Expand Up @@ -716,32 +717,131 @@ impl ToolCommand {
e
);
}
if formal {
let verify = verify.unwrap_or(true);
download_formal_snapshot(
&path,
epoch,
&genesis,
snapshot_store_config,
archive_store_config,
num_parallel_downloads,
network,
verify,
)
.await?;
let verify = verify.unwrap_or(true);
download_formal_snapshot(
&path,
epoch,
&genesis,
snapshot_store_config,
archive_store_config,
num_parallel_downloads,
network,
verify,
)
.await?;
}
ToolCommand::DownloadDBSnapshot {
epoch,
path,
skip_indexes,
num_parallel_downloads,
network,
snapshot_bucket,
snapshot_bucket_type,
snapshot_path,
no_sign_request,
verbose,
} => {
if !verbose {
tracing_handle
.update_log("off")
.expect("Failed to update log level");
}
let num_parallel_downloads = num_parallel_downloads.unwrap_or_else(|| {
num_cpus::get()
.checked_sub(1)
.expect("Failed to get number of CPUs")
});
let snapshot_bucket =
snapshot_bucket.or_else(|| match (network, no_sign_request) {
(Chain::Mainnet, false) => Some(
env::var("MAINNET_DB_SIGNED_BUCKET")
.unwrap_or("mysten-mainnet-snapshots".to_string()),
),
(Chain::Mainnet, true) => env::var("MAINNET_DB_UNSIGNED_BUCKET").ok(),
(Chain::Testnet, true) => env::var("TESTNET_DB_UNSIGNED_BUCKET").ok(),
(Chain::Testnet, _) => Some(
env::var("TESTNET_DB_SIGNED_BUCKET")
.unwrap_or("mysten-testnet-snapshots".to_string()),
),
(Chain::Unknown, _) => {
panic!("Cannot generate default snapshot bucket for unknown network");
}
});

let aws_endpoint = env::var("AWS_SNAPSHOT_ENDPOINT").ok();
let snapshot_bucket_type = if no_sign_request {
ObjectStoreType::S3
} else {
download_db_snapshot(
&path,
epoch,
&genesis,
snapshot_store_config,
archive_store_config,
skip_checkpoints,
skip_indexes,
num_parallel_downloads,
)
.await?;
snapshot_bucket_type
.expect("--snapshot-bucket-type must be set if not using --no-sign-request")
};
let snapshot_store_config = match snapshot_bucket_type {
ObjectStoreType::S3 => ObjectStoreConfig {
object_store: Some(ObjectStoreType::S3),
bucket: snapshot_bucket.filter(|s| !s.is_empty()),
aws_access_key_id: env::var("AWS_SNAPSHOT_ACCESS_KEY_ID").ok(),
aws_secret_access_key: env::var("AWS_SNAPSHOT_SECRET_ACCESS_KEY").ok(),
aws_region: env::var("AWS_SNAPSHOT_REGION").ok(),
aws_endpoint: aws_endpoint.filter(|s| !s.is_empty()),
aws_virtual_hosted_style_request: env::var(
"AWS_SNAPSHOT_VIRTUAL_HOSTED_REQUESTS",
)
.ok()
.and_then(|b| b.parse().ok())
.unwrap_or(no_sign_request),
object_store_connection_limit: 200,
no_sign_request,
..Default::default()
},
ObjectStoreType::GCS => ObjectStoreConfig {
object_store: Some(ObjectStoreType::GCS),
bucket: snapshot_bucket,
google_service_account: env::var("GCS_SNAPSHOT_SERVICE_ACCOUNT_FILE_PATH")
.ok(),
object_store_connection_limit: 200,
no_sign_request,
..Default::default()
},
ObjectStoreType::Azure => ObjectStoreConfig {
object_store: Some(ObjectStoreType::Azure),
bucket: snapshot_bucket,
azure_storage_account: env::var("AZURE_SNAPSHOT_STORAGE_ACCOUNT").ok(),
azure_storage_access_key: env::var("AZURE_SNAPSHOT_STORAGE_ACCESS_KEY")
.ok(),
object_store_connection_limit: 200,
no_sign_request,
..Default::default()
},
ObjectStoreType::File => {
if snapshot_path.is_some() {
ObjectStoreConfig {
object_store: Some(ObjectStoreType::File),
directory: snapshot_path,
..Default::default()
}
} else {
panic!(
"--snapshot-path must be specified for --snapshot-bucket-type=file"
);
}
}
};

if let Err(e) = check_completed_snapshot(&snapshot_store_config, epoch).await {
panic!(
"Aborting snapshot restore: {}, snapshot may not be uploaded yet",
e
);
}
download_db_snapshot(
&path,
epoch,
snapshot_store_config,
skip_indexes,
num_parallel_downloads,
)
.await?;
}
ToolCommand::Replay {
rpc_url,
Expand Down
Loading

0 comments on commit d96cc9d

Please sign in to comment.