Stream RPC snapshot downloads

This commit is contained in:
Michael Vines 2020-12-19 10:34:41 -08:00
parent 0090106f60
commit b3dc988564
3 changed files with 33 additions and 12 deletions

2
Cargo.lock generated
View File

@ -3991,6 +3991,7 @@ dependencies = [
"bs58",
"bv",
"byteorder",
"bytes 0.4.12",
"chrono",
"core_affinity",
"crossbeam-channel",
@ -4056,6 +4057,7 @@ dependencies = [
"thiserror",
"tokio 0.1.22",
"tokio 0.2.22",
"tokio-codec",
"tokio-fs",
"tokio-io",
"trees",

View File

@ -78,8 +78,10 @@ tempfile = "3.1.0"
thiserror = "1.0"
tokio = { version = "0.2", features = ["full"] }
tokio_01 = { version = "0.1", package = "tokio" }
tokio_01_bytes = { version = "0.4.7", package = "bytes" }
tokio_fs_01 = { version = "0.1", package = "tokio-fs" }
tokio_io_01 = { version = "0.1", package = "tokio-io" }
tokio_codec_01 = { version = "0.1", package = "tokio-codec" }
solana-rayon-threadlimit = { path = "../rayon-threadlimit", version = "1.6.0" }
trees = "0.2.1"

View File

@ -17,6 +17,7 @@ use jsonrpc_http_server::{
};
use regex::Regex;
use solana_ledger::blockstore::Blockstore;
use solana_metrics::inc_new_counter_info;
use solana_runtime::{
bank_forks::{BankForks, SnapshotConfig},
commitment::BlockCommitmentCache,
@ -61,7 +62,7 @@ impl RpcRequestMiddleware {
Self {
ledger_path,
snapshot_archive_path_regex: Regex::new(
r"/snapshot-\d+-[[:alnum:]]+\.tar\.(bz2|zst|gz)$",
r"/snapshot-\d+-[[:alnum:]]+\.(tar|tar\.bz2|tar\.zst|tar\.gz)$",
)
.unwrap(),
snapshot_config,
@ -85,6 +86,7 @@ impl RpcRequestMiddleware {
.unwrap()
}
#[allow(dead_code)]
fn internal_server_error() -> hyper::Response<hyper::Body> {
hyper::Response::builder()
.status(hyper::StatusCode::INTERNAL_SERVER_ERROR)
@ -112,13 +114,18 @@ impl RpcRequestMiddleware {
let stem = path.split_at(1).1; // Drop leading '/' from path
let filename = {
match path {
"/genesis.tar.bz2" => self.ledger_path.join(stem),
_ => self
.snapshot_config
.as_ref()
.unwrap()
.snapshot_package_output_path
.join(stem),
"/genesis.tar.bz2" => {
inc_new_counter_info!("rpc-get_genesis", 1);
self.ledger_path.join(stem)
}
_ => {
inc_new_counter_info!("rpc-get_snapshot", 1);
self.snapshot_config
.as_ref()
.unwrap()
.snapshot_package_output_path
.join(stem)
}
}
};
@ -129,10 +136,13 @@ impl RpcRequestMiddleware {
response: Box::new(
tokio_fs_01::file::File::open(filename)
.and_then(|file| {
let buf: Vec<u8> = Vec::new();
tokio_io_01::io::read_to_end(file, buf)
.and_then(|item| Ok(hyper::Response::new(item.1.into())))
.or_else(|_| Ok(RpcRequestMiddleware::internal_server_error()))
use tokio_codec_01::{BytesCodec, FramedRead};
let stream = FramedRead::new(file, BytesCodec::new())
.map(tokio_01_bytes::BytesMut::freeze);
let body = hyper::Body::wrap_stream(stream);
Ok(hyper::Response::new(body))
})
.or_else(|_| Ok(RpcRequestMiddleware::not_found())),
),
@ -530,6 +540,13 @@ mod tests {
assert!(rrm_with_snapshot_config.is_file_get_path(
"/snapshot-100-AvFf9oS8A8U78HdjT9YG2sTTThLHJZmhaMn2g8vkWYnr.tar.bz2"
));
assert!(rrm_with_snapshot_config.is_file_get_path(
"/snapshot-100-AvFf9oS8A8U78HdjT9YG2sTTThLHJZmhaMn2g8vkWYnr.tar.zst"
));
assert!(rrm_with_snapshot_config
.is_file_get_path("/snapshot-100-AvFf9oS8A8U78HdjT9YG2sTTThLHJZmhaMn2g8vkWYnr.tar.gz"));
assert!(rrm_with_snapshot_config
.is_file_get_path("/snapshot-100-AvFf9oS8A8U78HdjT9YG2sTTThLHJZmhaMn2g8vkWYnr.tar"));
assert!(!rrm.is_file_get_path(
"/snapshot-notaslotnumber-AvFf9oS8A8U78HdjT9YG2sTTThLHJZmhaMn2g8vkWYnr.tar.bz2"