From 35ea0b0d2ed2d72fdafbb0b64756578a389ec484 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Mon, 13 Feb 2023 19:33:34 +0100 Subject: feat(tvix/store): add tests for GRPCBlobServiceWrapper Change-Id: I0c9a2d887b5b94c43afee97a8a1e701cc9f74f23 Reviewed-on: https://cl.tvl.fyi/c/depot/+/8098 Reviewed-by: raitobezarius Tested-by: BuildkiteCI --- tvix/store/src/proto/tests/grpc_blobservice.rs | 210 +++++++++++++++++++++++++ tvix/store/src/proto/tests/mod.rs | 1 + 2 files changed, 211 insertions(+) create mode 100644 tvix/store/src/proto/tests/grpc_blobservice.rs (limited to 'tvix/store/src') diff --git a/tvix/store/src/proto/tests/grpc_blobservice.rs b/tvix/store/src/proto/tests/grpc_blobservice.rs new file mode 100644 index 0000000000..e616d141f1 --- /dev/null +++ b/tvix/store/src/proto/tests/grpc_blobservice.rs @@ -0,0 +1,210 @@ +use crate::blobservice::{BlobService, SledBlobService}; +use crate::chunkservice::{ChunkService, SledChunkService}; +use crate::proto::blob_meta::ChunkMeta; +use crate::proto::blob_service_server::BlobService as GRPCBlobService; +use crate::proto::{BlobChunk, GRPCBlobServiceWrapper, ReadBlobRequest, StatBlobRequest}; +use lazy_static::lazy_static; +use std::path::Path; +use tempfile::TempDir; + +lazy_static! { + // 2 bytes + static ref BLOB_A: Vec = vec![0x00, 0x01]; + static ref BLOB_A_DIGEST: Vec = blake3::hash(&BLOB_A).as_bytes().to_vec(); + + // 1MB + static ref BLOB_B: Vec = (0..255).collect::>().repeat(4 * 1024); + static ref BLOB_B_DIGEST: Vec = blake3::hash(&BLOB_B).as_bytes().to_vec(); +} + +fn gen_grpc_blob_service( + p: &Path, +) -> GRPCBlobServiceWrapper< + impl BlobService + Send + Sync + Clone + 'static, + impl ChunkService + Send + Sync + Clone + 'static, +> { + let blob_service = SledBlobService::new(p.join("blobs")).unwrap(); + let chunk_service = SledChunkService::new(p.join("chunks")).unwrap(); + GRPCBlobServiceWrapper::new(blob_service, chunk_service) +} + +/// Trying to read a non-existent blob should return a not found error. +#[tokio::test] +async fn not_found_read() { + let service = gen_grpc_blob_service(TempDir::new().unwrap().path()); + + let resp = service + .read(tonic::Request::new(ReadBlobRequest { + digest: BLOB_A_DIGEST.to_vec(), + })) + .await; + + let e = resp.expect_err("must_be_err"); + assert_eq!(e.code(), tonic::Code::NotFound); +} + +/// Trying to stat a non-existent blob should return a not found error. +#[tokio::test] +async fn not_found_stat() { + let service = gen_grpc_blob_service(TempDir::new().unwrap().path()); + + let resp = service + .stat(tonic::Request::new(StatBlobRequest { + digest: BLOB_A_DIGEST.to_vec(), + ..Default::default() + })) + .await + .expect_err("must fail"); + + // The resp should be a status with Code::NotFound + assert_eq!(resp.code(), tonic::Code::NotFound); +} + +/// Put a blob in the store, get it back. We send something small enough so it +/// won't get split into multiple chunks. +#[tokio::test] +async fn put_read_stat() { + let service = gen_grpc_blob_service(TempDir::new().unwrap().path()); + + // Send blob A. + let put_resp = service + .put(tonic_mock::streaming_request(vec![BlobChunk { + data: BLOB_A.clone(), + }])) + .await + .expect("must succeed") + .into_inner(); + + assert_eq!(BLOB_A_DIGEST.to_vec(), put_resp.digest); + + // Stat for the digest of A. It should return one chunk. + let resp = service + .stat(tonic::Request::new(StatBlobRequest { + digest: BLOB_A_DIGEST.to_vec(), + include_chunks: true, + ..Default::default() + })) + .await + .expect("must succeed") + .into_inner(); + + assert_eq!(1, resp.chunks.len()); + // the `chunks` field should point to the single chunk. + assert_eq!( + vec![ChunkMeta { + digest: BLOB_A_DIGEST.to_vec(), + size: BLOB_A.len() as u32, + }], + resp.chunks, + ); + + // Read the chunk. It should return the same data. + let resp = service + .read(tonic::Request::new(ReadBlobRequest { + digest: BLOB_A_DIGEST.to_vec(), + })) + .await; + + let mut rx = resp.expect("must succeed").into_inner().into_inner(); + + // the stream should contain one element, a BlobChunk with the same contents as BLOB_A. + let item = rx + .recv() + .await + .expect("must be some") + .expect("must succeed"); + + assert_eq!(BLOB_A.to_vec(), item.data); + + // … and no more elements + assert!(rx.recv().await.is_none()); +} + +/// Put a bigger blob in the store, and get it back. +/// Assert the stat request actually returns more than one chunk, and +/// we can read each chunk individually, as well as the whole blob via the +/// `read()` method. +#[tokio::test] +async fn put_read_stat_large() { + let service = gen_grpc_blob_service(TempDir::new().unwrap().path()); + + // split up BLOB_B into BlobChunks containing 1K bytes each. + let blob_b_blobchunks: Vec = BLOB_B + .chunks(1024) + .map(|x| BlobChunk { data: x.to_vec() }) + .collect(); + + assert!(blob_b_blobchunks.len() > 1); + + // Send blob B + let put_resp = service + .put(tonic_mock::streaming_request(blob_b_blobchunks)) + .await + .expect("must succeed") + .into_inner(); + + assert_eq!(BLOB_B_DIGEST.to_vec(), put_resp.digest); + + // Stat for the digest of B + let resp = service + .stat(tonic::Request::new(StatBlobRequest { + digest: BLOB_B_DIGEST.to_vec(), + include_chunks: true, + ..Default::default() + })) + .await + .expect("must succeed") + .into_inner(); + + // it should return more than one chunk. + assert_ne!(1, resp.chunks.len()); + + // The size added up should equal the size of BLOB_B. + let mut size_in_stat: u32 = 0; + for chunk in &resp.chunks { + size_in_stat += chunk.size + } + assert_eq!(BLOB_B.len() as u32, size_in_stat); + + // Reading the whole blob by its digest via the read() interface should succeed. + { + let resp = service + .read(tonic::Request::new(ReadBlobRequest { + digest: BLOB_B_DIGEST.to_vec(), + })) + .await; + + let mut rx = resp.expect("must succeed").into_inner().into_inner(); + + let mut buf: Vec = Vec::new(); + while let Some(item) = rx.recv().await { + let mut blob_chunk = item.expect("must not be err"); + buf.append(&mut blob_chunk.data); + } + + assert_eq!(BLOB_B.to_vec(), buf); + } + + // Reading the whole blob by reading individual chunks should also succeed. + { + let mut buf: Vec = Vec::new(); + for chunk in &resp.chunks { + // request this individual chunk via read + let resp = service + .read(tonic::Request::new(ReadBlobRequest { + digest: chunk.digest.clone(), + })) + .await; + + let mut rx = resp.expect("must succeed").into_inner().into_inner(); + + // append all items from the stream to the buffer + while let Some(item) = rx.recv().await { + let mut blob_chunk = item.expect("must not be err"); + buf.append(&mut blob_chunk.data); + } + } + // finished looping over all chunks, compare + assert_eq!(BLOB_B.to_vec(), buf); + } +} diff --git a/tvix/store/src/proto/tests/mod.rs b/tvix/store/src/proto/tests/mod.rs index 04631afe40..9995be9d37 100644 --- a/tvix/store/src/proto/tests/mod.rs +++ b/tvix/store/src/proto/tests/mod.rs @@ -1,3 +1,4 @@ mod directory; mod directory_nodes_iterator; +mod grpc_blobservice; mod pathinfo; -- cgit 1.4.1