mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(tvix/castore/grpc): don't use explicit channel

Browse source 
We can just use the `BoxStream` directly, or a `once` with the single
`Directory`.

In the recursive case, we also did not properly close the channel after
the first error.

Change-Id: Ifad56d307fc7861107b6d3cffd28d35631d526e6
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11635
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Autosubmit: flokli <flokli@flokli.de>
This commit is contained in:
Florian Klink 2024-05-12 14:52:09 +03:00 committed by clbot
parent b26569028c
commit 30995a0990
1 changed files with 35 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
tvix/castore/src/proto/grpc_directoryservice_wrapper.rs
View file

@ -1,12 +1,12 @@
use crate::directoryservice::ClosureValidator; use crate::directoryservice::ClosureValidator;
use crate::proto; use crate::proto;
use crate::{directoryservice::DirectoryService, B3Digest}; use crate::{directoryservice::DirectoryService, B3Digest};
use futures::StreamExt; use futures::stream::BoxStream;
use futures::TryStreamExt;
use std::ops::Deref; use std::ops::Deref;
use tokio::sync::mpsc::channel; use tokio_stream::once;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{async_trait, Request, Response, Status, Streaming}; use tonic::{async_trait, Request, Response, Status, Streaming};
use tracing::{debug, instrument, warn}; use tracing::{instrument, warn};
pub struct GRPCDirectoryServiceWrapper<T> { pub struct GRPCDirectoryServiceWrapper<T> {
directory_service: T, directory_service: T,
@ -23,63 +23,52 @@ impl<T> proto::directory_service_server::DirectoryService for GRPCDirectoryServi
where where
T: Deref<Target = dyn DirectoryService> + Send + Sync + 'static, T: Deref<Target = dyn DirectoryService> + Send + Sync + 'static,
{ {
type GetStream = ReceiverStream<tonic::Result<proto::Directory, Status>>; type GetStream = BoxStream<'static, tonic::Result<proto::Directory, Status>>;
#[instrument(skip_all)] #[instrument(skip_all)]
async fn get( async fn get<'a>(
&self, &'a self,
request: Request<proto::GetDirectoryRequest>, request: Request<proto::GetDirectoryRequest>,
) -> Result<Response<Self::GetStream>, Status> { ) -> Result<Response<Self::GetStream>, Status> {
let (tx, rx) = channel(5);
let req_inner = request.into_inner(); let req_inner = request.into_inner();
// look at the digest in the request and put it in the top of the queue. let by_what = &req_inner
match &req_inner.by_what { .by_what
None => return Err(Status::invalid_argument("by_what needs to be specified")), .ok_or_else(|| Status::invalid_argument("invalid by_what"))?;
Some(proto::get_directory_request::ByWhat::Digest(ref digest)) => {
match by_what {
proto::get_directory_request::ByWhat::Digest(ref digest) => {
let digest: B3Digest = digest let digest: B3Digest = digest
.clone() .clone()
.try_into() .try_into()
.map_err(|_e| Status::invalid_argument("invalid digest length"))?; .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
if !req_inner.recursive { Ok(tonic::Response::new({
let e: Result<proto::Directory, Status> = match self if !req_inner.recursive {
.directory_service let directory = self
.get(&digest) .directory_service
.await .get(&digest)
{ .await
Ok(Some(directory)) => Ok(directory), .map_err(|e| {
Ok(None) => { warn!(err = %e, directory.digest=%digest, "failed to get directory");
Err(Status::not_found(format!("directory {} not found", digest))) tonic::Status::new(tonic::Code::Internal, e.to_string())
} })?
Err(e) => { .ok_or_else(|| {
warn!(err = %e, directory.digest=%digest, "failed to get directory"); Status::not_found(format!("directory {} not found", digest))
Err(e.into()) })?;
}
};
if tx.send(e).await.is_err() { Box::pin(once(Ok(directory)))
debug!("receiver dropped"); } else {
// If recursive was requested, traverse via get_recursive.
Box::pin(
self.directory_service.get_recursive(&digest).map_err(|e| {
tonic::Status::new(tonic::Code::Internal, e.to_string())
}),
)
} }
} else { }))
// If recursive was requested, traverse via get_recursive.
let mut directories_it = self.directory_service.get_recursive(&digest);
while let Some(e) = directories_it.next().await {
// map err in res from Error to Status
let res = e.map_err(|e| Status::internal(e.to_string()));
if tx.send(res).await.is_err() {
debug!("receiver dropped");
break;
}
}
}
} }
} }
let receiver_stream = ReceiverStream::new(rx);
Ok(Response::new(receiver_stream))
} }
#[instrument(skip_all)] #[instrument(skip_all)]