use crate::{blobservice::BlobService, chunkservice::ChunkService, Error}; use data_encoding::BASE64; use tokio::{sync::mpsc::channel, task}; use tokio_stream::wrappers::ReceiverStream; use tonic::{async_trait, Request, Response, Status, Streaming}; use tracing::{debug, instrument}; pub struct GRPCBlobServiceWrapper { blob_service: BS, chunk_service: CS, } impl GRPCBlobServiceWrapper { pub fn new(blob_service: BS, chunk_service: CS) -> Self { Self { blob_service, chunk_service, } } // upload the chunk to the chunk service, and return its digest (or an error) when done. #[instrument(skip(chunk_service))] fn upload_chunk(chunk_service: CS, chunk_data: Vec) -> Result, Error> { let mut hasher = blake3::Hasher::new(); if chunk_data.len() >= 128 * 1024 { hasher.update_rayon(&chunk_data); } else { hasher.update(&chunk_data); } let digest = hasher.finalize(); if chunk_service.has(digest.as_bytes())? { debug!("already has chunk, skipping"); } let digest_resp = chunk_service.put(chunk_data)?; assert_eq!(digest_resp, digest.as_bytes()); Ok(digest.as_bytes().to_vec()) } } #[async_trait] impl< BS: BlobService + Send + Sync + Clone + 'static, CS: ChunkService + Send + Sync + Clone + 'static, > super::blob_service_server::BlobService for GRPCBlobServiceWrapper { type ReadStream = ReceiverStream>; #[instrument(skip(self))] async fn stat( &self, request: Request, ) -> Result, Status> { let rq = request.into_inner(); match self.blob_service.stat(&rq) { Ok(None) => Err(Status::not_found(format!( "blob {} not found", BASE64.encode(&rq.digest) ))), Ok(Some(blob_meta)) => Ok(Response::new(blob_meta)), Err(e) => Err(e.into()), } } #[instrument(skip(self))] async fn read( &self, request: Request, ) -> Result, Status> { let req = request.into_inner(); let (tx, rx) = channel(5); // query the chunk service for more detailed blob info let stat_resp = self.blob_service.stat(&super::StatBlobRequest { digest: req.digest.to_vec(), include_chunks: true, ..Default::default() })?; match stat_resp { None => { // If the stat didn't return any blobmeta, the client might // still have asked for a single chunk to be read. // Check the chunkstore. if let Some(data) = self.chunk_service.get(&req.digest)? { // We already know the hash matches, and contrary to // iterating over a blobmeta, we can't know the size, // so send the contents of that chunk over, // as the first (and only) element of the stream. task::spawn(async move { let res = Ok(super::BlobChunk { data }); // send the result to the client. If the client already left, that's also fine. if (tx.send(res).await).is_err() { debug!("receiver dropped"); } }); } else { return Err(Status::not_found(format!( "blob {} not found", BASE64.encode(&req.digest), ))); } } Some(blobmeta) => { let chunk_client = self.chunk_service.clone(); // TODO: use BlobReader? // But then we might not be able to send compressed chunks as-is. // Might require implementing https://docs.rs/futures-core/latest/futures_core/stream/trait.Stream.html for it // first, so we can .next().await in here. task::spawn(async move { for chunkmeta in blobmeta.chunks { // request chunk. // We don't need to validate the digest again, as // that's required for all implementations of ChunkService. let res = match chunk_client.get(&chunkmeta.digest) { Err(e) => Err(e.into()), // TODO: make this a separate error type Ok(None) => Err(Error::StorageError(format!( "consistency error: chunk {} for blob {} not found", BASE64.encode(&chunkmeta.digest), BASE64.encode(&req.digest), )) .into()), Ok(Some(data)) => { // We already know the hash matches, but also // check the size matches what chunkmeta said. if data.len() as u32 != chunkmeta.size { Err(Error::StorageError(format!( "consistency error: chunk {} for blob {} has wrong size, expected {}, got {}", BASE64.encode(&chunkmeta.digest), BASE64.encode(&req.digest), chunkmeta.size, data.len(), )).into()) } else { // send out the current chunk // TODO: we might want to break this up further if too big? Ok(super::BlobChunk { data }) } } }; // send the result to the client if (tx.send(res).await).is_err() { debug!("receiver dropped"); break; } } }); } } let receiver_stream = ReceiverStream::new(rx); Ok(Response::new(receiver_stream)) } #[instrument(skip(self))] async fn put( &self, request: Request>, ) -> Result, Status> { let mut req_inner = request.into_inner(); // initialize a blake3 hasher calculating the hash of the whole blob. let mut blob_hasher = blake3::Hasher::new(); // start a BlobMeta, which we'll fill while looping over the chunks let mut blob_meta = super::BlobMeta::default(); // is filled with bytes received from the client. let mut buf: Vec = vec![]; // This reads data from the client, chunks it up using fastcdc, // uploads all chunks to the [ChunkService], and fills a // [super::BlobMeta] linking to these chunks. while let Some(blob_chunk) = req_inner.message().await? { // calculate blob hash, and use rayon if data is > 128KiB. if blob_chunk.data.len() > 128 * 1024 { blob_hasher.update_rayon(&blob_chunk.data); } else { blob_hasher.update(&blob_chunk.data); } // extend buf with the newly received data buf.append(&mut blob_chunk.data.clone()); // TODO: play with chunking sizes let chunker_avg_size = 64 * 1024; let chunker_min_size = chunker_avg_size / 4; let chunker_max_size = chunker_avg_size * 4; // initialize a chunker with the current buffer let chunker = fastcdc::v2020::FastCDC::new( &buf, chunker_min_size, chunker_avg_size, chunker_max_size, ); // ask the chunker for cutting points in the buffer. let mut start_pos = 0 as usize; buf = loop { // ask the chunker for the next cutting point. let (_fp, end_pos) = chunker.cut(start_pos, buf.len() - start_pos); // whenever the last cut point is pointing to the end of the buffer, // keep that chunk left in there. // We don't know if the chunker decided to cut here simply because it was // at the end of the buffer, or if it would also cut if there // were more data. // // Split off all previous chunks and keep this chunk data in the buffer. if end_pos == buf.len() { break buf.split_off(start_pos); } // Upload that chunk to the chunk service and record it in BlobMeta. // TODO: make upload_chunk async and upload concurrently? let chunk_data = &buf[start_pos..end_pos]; let chunk_digest = Self::upload_chunk(self.chunk_service.clone(), chunk_data.to_vec())?; blob_meta.chunks.push(super::blob_meta::ChunkMeta { digest: chunk_digest, size: chunk_data.len() as u32, }); // move start_pos over the processed chunk. start_pos = end_pos; } } // Also upload the last chunk (what's left in `buf`) to the chunk // service and record it in BlobMeta. let buf_len = buf.len() as u32; let chunk_digest = Self::upload_chunk(self.chunk_service.clone(), buf)?; blob_meta.chunks.push(super::blob_meta::ChunkMeta { digest: chunk_digest, size: buf_len, }); let blob_digest = blob_hasher.finalize().as_bytes().to_vec(); // check if we have the received blob in the [BlobService] already. let resp = self.blob_service.stat(&super::StatBlobRequest { digest: blob_digest.to_vec(), ..Default::default() })?; // if not, store. if resp.is_none() { self.blob_service.put(&blob_digest, blob_meta)?; } // return to client. Ok(Response::new(super::PutBlobResponse { digest: blob_digest, })) } }