Merge pull request 'rewrite read_and_put_block as a series of steps with channels' (#734) from refactor-put into main

Reviewed-on: Deuxfleurs/garage#734
This commit is contained in:
Alex 2024-02-26 17:52:45 +00:00
commit 911a83ea7d
4 changed files with 169 additions and 83 deletions

View file

@ -387,7 +387,10 @@ pub async fn handle_upload_part_copy(
// we need to insert that data as a new block. // we need to insert that data as a new block.
async move { async move {
if must_upload { if must_upload {
garage2.block_manager.rpc_put_block(final_hash, data).await garage2
.block_manager
.rpc_put_block(final_hash, data, None)
.await
} else { } else {
Ok(()) Ok(())
} }

View file

@ -6,7 +6,6 @@ use hyper::{Request, Response};
use md5::{Digest as Md5Digest, Md5}; use md5::{Digest as Md5Digest, Md5};
use garage_table::*; use garage_table::*;
use garage_util::async_hash::*;
use garage_util::data::*; use garage_util::data::*;
use garage_model::bucket_table::Bucket; use garage_model::bucket_table::Bucket;
@ -135,17 +134,8 @@ pub async fn handle_put_part(
garage.version_table.insert(&version).await?; garage.version_table.insert(&version).await?;
// Copy data to version // Copy data to version
let first_block_hash = async_blake2sum(first_block.clone()).await; let (total_size, data_md5sum, data_sha256sum, _) =
read_and_put_blocks(&garage, &version, part_number, first_block, &mut chunker).await?;
let (total_size, data_md5sum, data_sha256sum) = read_and_put_blocks(
&garage,
&version,
part_number,
first_block,
first_block_hash,
&mut chunker,
)
.await?;
// Verify that checksums map // Verify that checksums map
ensure_checksum_matches( ensure_checksum_matches(

View file

@ -3,10 +3,13 @@ use std::sync::Arc;
use base64::prelude::*; use base64::prelude::*;
use futures::prelude::*; use futures::prelude::*;
use futures::stream::FuturesOrdered;
use futures::try_join; use futures::try_join;
use md5::{digest::generic_array::*, Digest as Md5Digest, Md5}; use md5::{digest::generic_array::*, Digest as Md5Digest, Md5};
use sha2::Sha256; use sha2::Sha256;
use tokio::sync::mpsc;
use hyper::body::Bytes; use hyper::body::Bytes;
use hyper::header::{HeaderMap, HeaderValue}; use hyper::header::{HeaderMap, HeaderValue};
use hyper::{Request, Response}; use hyper::{Request, Response};
@ -17,6 +20,7 @@ use opentelemetry::{
}; };
use garage_net::bytes_buf::BytesBuf; use garage_net::bytes_buf::BytesBuf;
use garage_rpc::rpc_helper::OrderTag;
use garage_table::*; use garage_table::*;
use garage_util::async_hash::*; use garage_util::async_hash::*;
use garage_util::data::*; use garage_util::data::*;
@ -35,6 +39,8 @@ use crate::helpers::*;
use crate::s3::api_server::{ReqBody, ResBody}; use crate::s3::api_server::{ReqBody, ResBody};
use crate::s3::error::*; use crate::s3::error::*;
const PUT_BLOCKS_MAX_PARALLEL: usize = 3;
pub async fn handle_put( pub async fn handle_put(
garage: Arc<Garage>, garage: Arc<Garage>,
req: Request<ReqBody>, req: Request<ReqBody>,
@ -168,17 +174,8 @@ pub(crate) async fn save_stream<S: Stream<Item = Result<Bytes, Error>> + Unpin>(
garage.version_table.insert(&version).await?; garage.version_table.insert(&version).await?;
// Transfer data and verify checksum // Transfer data and verify checksum
let first_block_hash = async_blake2sum(first_block.clone()).await; let (total_size, data_md5sum, data_sha256sum, first_block_hash) =
read_and_put_blocks(&garage, &version, 1, first_block, &mut chunker).await?;
let (total_size, data_md5sum, data_sha256sum) = read_and_put_blocks(
&garage,
&version,
1,
first_block,
first_block_hash,
&mut chunker,
)
.await?;
ensure_checksum_matches( ensure_checksum_matches(
data_md5sum.as_slice(), data_md5sum.as_slice(),
@ -299,84 +296,164 @@ pub(crate) async fn read_and_put_blocks<S: Stream<Item = Result<Bytes, Error>> +
version: &Version, version: &Version,
part_number: u64, part_number: u64,
first_block: Bytes, first_block: Bytes,
first_block_hash: Hash,
chunker: &mut StreamChunker<S>, chunker: &mut StreamChunker<S>,
) -> Result<(u64, GenericArray<u8, typenum::U16>, Hash), Error> { ) -> Result<(u64, GenericArray<u8, typenum::U16>, Hash, Hash), Error> {
let tracer = opentelemetry::global::tracer("garage"); let tracer = opentelemetry::global::tracer("garage");
let md5hasher = AsyncHasher::<Md5>::new(); let (block_tx, mut block_rx) = mpsc::channel::<Result<Bytes, Error>>(2);
let sha256hasher = AsyncHasher::<Sha256>::new(); let read_blocks = async {
block_tx.send(Ok(first_block)).await?;
loop {
let res = chunker
.next()
.with_context(Context::current_with_span(
tracer.start("Read block from client"),
))
.await;
match res {
Ok(Some(block)) => block_tx.send(Ok(block)).await?,
Ok(None) => break,
Err(e) => {
block_tx.send(Err(e)).await?;
break;
}
}
}
drop(block_tx);
Ok::<_, mpsc::error::SendError<_>>(())
};
futures::future::join( let (block_tx2, mut block_rx2) = mpsc::channel::<Result<Bytes, Error>>(1);
md5hasher.update(first_block.clone()), let hash_stream = async {
sha256hasher.update(first_block.clone()), let md5hasher = AsyncHasher::<Md5>::new();
) let sha256hasher = AsyncHasher::<Sha256>::new();
.with_context(Context::current_with_span( while let Some(next) = block_rx.recv().await {
tracer.start("Hash first block (md5, sha256)"), match next {
)) Ok(block) => {
.await; block_tx2.send(Ok(block.clone())).await?;
futures::future::join(
md5hasher.update(block.clone()),
sha256hasher.update(block.clone()),
)
.with_context(Context::current_with_span(
tracer.start("Hash block (md5, sha256)"),
))
.await;
}
Err(e) => {
block_tx2.send(Err(e)).await?;
break;
}
}
}
drop(block_tx2);
Ok::<_, mpsc::error::SendError<_>>(futures::join!(
md5hasher.finalize(),
sha256hasher.finalize()
))
};
let mut next_offset = first_block.len(); let (block_tx3, mut block_rx3) = mpsc::channel::<Result<(Bytes, Hash), Error>>(1);
let mut put_curr_version_block = put_block_meta( let hash_blocks = async {
garage, let mut first_block_hash = None;
version, while let Some(next) = block_rx2.recv().await {
part_number, match next {
0, Ok(block) => {
first_block_hash, let hash = async_blake2sum(block.clone())
first_block.len() as u64, .with_context(Context::current_with_span(
); tracer.start("Hash block (blake2)"),
let mut put_curr_block = garage ))
.block_manager .await;
.rpc_put_block(first_block_hash, first_block); if first_block_hash.is_none() {
first_block_hash = Some(hash);
}
block_tx3.send(Ok((block, hash))).await?;
}
Err(e) => {
block_tx3.send(Err(e)).await?;
break;
}
}
}
drop(block_tx3);
Ok::<_, mpsc::error::SendError<_>>(first_block_hash.unwrap())
};
loop { let put_blocks = async {
let (_, _, next_block) = futures::try_join!( // Structure for handling several concurrent writes to storage nodes
put_curr_block.map_err(Error::from), let order_stream = OrderTag::stream();
put_curr_version_block.map_err(Error::from), let mut write_futs = FuturesOrdered::new();
chunker.next(), let mut written_bytes = 0u64;
)?; loop {
if let Some(block) = next_block { // Simultaneously write blocks to storage nodes & await for next block to be written
let (_, _, block_hash) = futures::future::join3( let currently_running = write_futs.len();
md5hasher.update(block.clone()), let write_futs_next = async {
sha256hasher.update(block.clone()), if write_futs.is_empty() {
async_blake2sum(block.clone()), futures::future::pending().await
) } else {
.with_context(Context::current_with_span( write_futs.next().await.unwrap()
tracer.start("Hash block (md5, sha256, blake2)"), }
)) };
.await; let recv_next = async {
let block_len = block.len(); // If more than a maximum number of writes are in progress, don't add more for now
put_curr_version_block = put_block_meta( if currently_running >= PUT_BLOCKS_MAX_PARALLEL {
futures::future::pending().await
} else {
block_rx3.recv().await
}
};
let (block, hash) = tokio::select! {
result = write_futs_next => {
result?;
continue;
},
recv = recv_next => match recv {
Some(next) => next?,
None => break,
},
};
// For next block to be written: count its size and spawn future to write it
let offset = written_bytes;
written_bytes += block.len() as u64;
write_futs.push_back(put_block_and_meta(
garage, garage,
version, version,
part_number, part_number,
next_offset as u64, offset,
block_hash, hash,
block_len as u64, block,
); order_stream.order(written_bytes),
put_curr_block = garage.block_manager.rpc_put_block(block_hash, block); ));
next_offset += block_len;
} else {
break;
} }
} while let Some(res) = write_futs.next().await {
res?;
}
Ok::<_, Error>(written_bytes)
};
let total_size = next_offset as u64; let (_, stream_hash_result, block_hash_result, final_result) =
let data_md5sum = md5hasher.finalize().await; futures::join!(read_blocks, hash_stream, hash_blocks, put_blocks);
let total_size = final_result?;
// unwrap here is ok, because if hasher failed, it is because something failed
// later in the pipeline which already caused a return at the ? on previous line
let (data_md5sum, data_sha256sum) = stream_hash_result.unwrap();
let first_block_hash = block_hash_result.unwrap();
let data_sha256sum = sha256hasher.finalize().await;
let data_sha256sum = Hash::try_from(&data_sha256sum[..]).unwrap(); let data_sha256sum = Hash::try_from(&data_sha256sum[..]).unwrap();
Ok((total_size, data_md5sum, data_sha256sum)) Ok((total_size, data_md5sum, data_sha256sum, first_block_hash))
} }
async fn put_block_meta( async fn put_block_and_meta(
garage: &Garage, garage: &Garage,
version: &Version, version: &Version,
part_number: u64, part_number: u64,
offset: u64, offset: u64,
hash: Hash, hash: Hash,
size: u64, block: Bytes,
order_tag: OrderTag,
) -> Result<(), GarageError> { ) -> Result<(), GarageError> {
let mut version = version.clone(); let mut version = version.clone();
version.blocks.put( version.blocks.put(
@ -384,7 +461,10 @@ async fn put_block_meta(
part_number, part_number,
offset, offset,
}, },
VersionBlock { hash, size }, VersionBlock {
hash,
size: block.len() as u64,
},
); );
let block_ref = BlockRef { let block_ref = BlockRef {
@ -394,6 +474,9 @@ async fn put_block_meta(
}; };
futures::try_join!( futures::try_join!(
garage
.block_manager
.rpc_put_block(hash, block, Some(order_tag)),
garage.version_table.insert(&version), garage.version_table.insert(&version),
garage.block_ref_table.insert(&block_ref), garage.block_ref_table.insert(&block_ref),
)?; )?;

View file

@ -346,7 +346,12 @@ impl BlockManager {
} }
/// Send block to nodes that should have it /// Send block to nodes that should have it
pub async fn rpc_put_block(&self, hash: Hash, data: Bytes) -> Result<(), Error> { pub async fn rpc_put_block(
&self,
hash: Hash,
data: Bytes,
order_tag: Option<OrderTag>,
) -> Result<(), Error> {
let who = self.replication.write_nodes(&hash); let who = self.replication.write_nodes(&hash);
let (header, bytes) = DataBlock::from_buffer(data, self.compression_level) let (header, bytes) = DataBlock::from_buffer(data, self.compression_level)
@ -354,6 +359,11 @@ impl BlockManager {
.into_parts(); .into_parts();
let put_block_rpc = let put_block_rpc =
Req::new(BlockRpc::PutBlock { hash, header })?.with_stream_from_buffer(bytes); Req::new(BlockRpc::PutBlock { hash, header })?.with_stream_from_buffer(bytes);
let put_block_rpc = if let Some(tag) = order_tag {
put_block_rpc.with_order_tag(tag)
} else {
put_block_rpc
};
self.system self.system
.rpc .rpc