refactor(tvix): move castore into tvix-castore crate
This splits the pure content-addressed layers from tvix-store into a `castore` crate, and only leaves PathInfo related things, as well as the CLI entrypoint in the tvix-store crate. Notable changes: - `fixtures` and `utils` had to be moved out of the `test` cfg, so they can be imported from tvix-store. - Some ad-hoc fixtures in the test were moved to proper fixtures in the same step. - The protos are now created by a (more static) recipe in the protos/ directory. The (now two) golang targets are commented out, as it's not possible to update them properly in the same CL. This will be done by a followup CL once this is merged (and whitby deployed) Bug: https://b.tvl.fyi/issues/301 Change-Id: I8d675d4bf1fb697eb7d479747c1b1e3635718107 Reviewed-on: https://cl.tvl.fyi/c/depot/+/9370 Reviewed-by: tazjin <tazjin@tvl.su> Reviewed-by: flokli <flokli@flokli.de> Autosubmit: flokli <flokli@flokli.de> Tested-by: BuildkiteCI Reviewed-by: Connor Brewster <cbrewster@hey.com>
This commit is contained in:
parent
d8ef0cfb4a
commit
32f41458c0
89 changed files with 2308 additions and 1829 deletions
177
tvix/castore/src/proto/grpc_blobservice_wrapper.rs
Normal file
177
tvix/castore/src/proto/grpc_blobservice_wrapper.rs
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
use crate::blobservice::BlobService;
|
||||
use core::pin::pin;
|
||||
use futures::TryFutureExt;
|
||||
use std::{
|
||||
collections::VecDeque,
|
||||
io,
|
||||
ops::{Deref, DerefMut},
|
||||
pin::Pin,
|
||||
sync::Arc,
|
||||
};
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::io::ReaderStream;
|
||||
use tonic::{async_trait, Request, Response, Status, Streaming};
|
||||
use tracing::{instrument, warn};
|
||||
|
||||
pub struct GRPCBlobServiceWrapper {
|
||||
blob_service: Arc<dyn BlobService>,
|
||||
}
|
||||
|
||||
impl From<Arc<dyn BlobService>> for GRPCBlobServiceWrapper {
|
||||
fn from(value: Arc<dyn BlobService>) -> Self {
|
||||
Self {
|
||||
blob_service: value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// This is necessary because bytes::BytesMut comes up with
|
||||
// a default 64 bytes capacity that cannot be changed
|
||||
// easily if you assume a bytes::BufMut trait implementation
|
||||
// Therefore, we override the Default implementation here
|
||||
// TODO(raitobezarius?): upstream me properly
|
||||
struct BytesMutWithDefaultCapacity<const N: usize> {
|
||||
inner: bytes::BytesMut,
|
||||
}
|
||||
|
||||
impl<const N: usize> Deref for BytesMutWithDefaultCapacity<N> {
|
||||
type Target = bytes::BytesMut;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl<const N: usize> DerefMut for BytesMutWithDefaultCapacity<N> {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl<const N: usize> Default for BytesMutWithDefaultCapacity<N> {
|
||||
fn default() -> Self {
|
||||
BytesMutWithDefaultCapacity {
|
||||
inner: bytes::BytesMut::with_capacity(N),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<const N: usize> bytes::Buf for BytesMutWithDefaultCapacity<N> {
|
||||
fn remaining(&self) -> usize {
|
||||
self.inner.remaining()
|
||||
}
|
||||
|
||||
fn chunk(&self) -> &[u8] {
|
||||
self.inner.chunk()
|
||||
}
|
||||
|
||||
fn advance(&mut self, cnt: usize) {
|
||||
self.inner.advance(cnt);
|
||||
}
|
||||
}
|
||||
|
||||
unsafe impl<const N: usize> bytes::BufMut for BytesMutWithDefaultCapacity<N> {
|
||||
fn remaining_mut(&self) -> usize {
|
||||
self.inner.remaining_mut()
|
||||
}
|
||||
|
||||
unsafe fn advance_mut(&mut self, cnt: usize) {
|
||||
self.inner.advance_mut(cnt);
|
||||
}
|
||||
|
||||
fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
|
||||
self.inner.chunk_mut()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl super::blob_service_server::BlobService for GRPCBlobServiceWrapper {
|
||||
// https://github.com/tokio-rs/tokio/issues/2723#issuecomment-1534723933
|
||||
type ReadStream =
|
||||
Pin<Box<dyn futures::Stream<Item = Result<super::BlobChunk, Status>> + Send + 'static>>;
|
||||
|
||||
#[instrument(skip(self))]
|
||||
async fn stat(
|
||||
&self,
|
||||
request: Request<super::StatBlobRequest>,
|
||||
) -> Result<Response<super::BlobMeta>, Status> {
|
||||
let rq = request.into_inner();
|
||||
let req_digest = rq
|
||||
.digest
|
||||
.try_into()
|
||||
.map_err(|_e| Status::invalid_argument("invalid digest length"))?;
|
||||
|
||||
match self.blob_service.has(&req_digest).await {
|
||||
Ok(true) => Ok(Response::new(super::BlobMeta::default())),
|
||||
Ok(false) => Err(Status::not_found(format!("blob {} not found", &req_digest))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[instrument(skip(self))]
|
||||
async fn read(
|
||||
&self,
|
||||
request: Request<super::ReadBlobRequest>,
|
||||
) -> Result<Response<Self::ReadStream>, Status> {
|
||||
let rq = request.into_inner();
|
||||
|
||||
let req_digest = rq
|
||||
.digest
|
||||
.try_into()
|
||||
.map_err(|_e| Status::invalid_argument("invalid digest length"))?;
|
||||
|
||||
match self.blob_service.open_read(&req_digest).await {
|
||||
Ok(Some(reader)) => {
|
||||
fn stream_mapper(
|
||||
x: Result<bytes::Bytes, io::Error>,
|
||||
) -> Result<super::BlobChunk, Status> {
|
||||
match x {
|
||||
Ok(bytes) => Ok(super::BlobChunk { data: bytes }),
|
||||
Err(e) => Err(Status::from(e)),
|
||||
}
|
||||
}
|
||||
|
||||
let chunks_stream = ReaderStream::new(reader).map(stream_mapper);
|
||||
Ok(Response::new(Box::pin(chunks_stream)))
|
||||
}
|
||||
Ok(None) => Err(Status::not_found(format!("blob {} not found", &req_digest))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[instrument(skip(self))]
|
||||
async fn put(
|
||||
&self,
|
||||
request: Request<Streaming<super::BlobChunk>>,
|
||||
) -> Result<Response<super::PutBlobResponse>, Status> {
|
||||
let req_inner = request.into_inner();
|
||||
|
||||
let data_stream = req_inner.map(|x| {
|
||||
x.map(|x| VecDeque::from(x.data.to_vec()))
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
|
||||
});
|
||||
|
||||
let mut data_reader = tokio_util::io::StreamReader::new(data_stream);
|
||||
|
||||
let mut blob_writer = pin!(self.blob_service.open_write().await);
|
||||
|
||||
tokio::io::copy(&mut data_reader, &mut blob_writer)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
warn!("error copying: {}", e);
|
||||
Status::internal("error copying")
|
||||
})?;
|
||||
|
||||
let digest = blob_writer
|
||||
.close()
|
||||
.map_err(|e| {
|
||||
warn!("error closing stream: {}", e);
|
||||
Status::internal("error closing stream")
|
||||
})
|
||||
.await?
|
||||
.to_vec();
|
||||
|
||||
Ok(Response::new(super::PutBlobResponse {
|
||||
digest: digest.into(),
|
||||
}))
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue