maurice/snix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(tvix/tracing): gRPC trace context propagation

Browse source 
This introduces optional helper function in tvix/tracing for trace
propagation and uses these helper in the `tvix-store`.

The GRPCBlobService, GRPCDirectoryService and GRPCPathInfoService now
accept a generic client, meaning the client can be generated with either
`::new` or `::with_interceptor`.

This was tested and validated by starting a `tvix-store daemon` and
`tvix-store import`.

Change-Id: I4b194483bf09266820104b4b56e4a135dca2b77a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11863
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
This commit is contained in:
Simon Hauser 2024-06-20 11:39:09 +02:00
parent 2b20d8d82d
commit 639a00e2ab
18 changed files with 399 additions and 48 deletions
Download patch file Download diff file Expand all files Collapse all files

27
tvix/castore/src/blobservice/grpc.rs
View file

@ -17,29 +17,39 @@ use tokio_util::{
io::{CopyToBytes, SinkWriter},
sync::PollSender,
};
use tonic::{async_trait, transport::Channel, Code, Status};
use tonic::{async_trait, Code, Status};
use tracing::{instrument, Instrument as _};
/// Connects to a (remote) tvix-store BlobService over gRPC.
#[derive(Clone)]
pub struct GRPCBlobService {
pub struct GRPCBlobService<T>
where
T: Clone,
{
/// The internal reference to a gRPC client.
/// Cloning it is cheap, and it internally handles concurrent requests.
grpc_client: proto::blob_service_client::BlobServiceClient<Channel>,
grpc_client: proto::blob_service_client::BlobServiceClient<T>,
}
impl GRPCBlobService {
impl<T> GRPCBlobService<T>
where
T: tonic::client::GrpcService<tonic::body::BoxBody> + Clone,
{
/// construct a [GRPCBlobService] from a [proto::blob_service_client::BlobServiceClient].
/// panics if called outside the context of a tokio runtime.
pub fn from_client(
grpc_client: proto::blob_service_client::BlobServiceClient<Channel>,
) -> Self {
pub fn from_client(grpc_client: proto::blob_service_client::BlobServiceClient<T>) -> Self {
Self { grpc_client }
}
}
#[async_trait]
impl BlobService for GRPCBlobService {
impl<T> BlobService for GRPCBlobService<T>
where
T: tonic::client::GrpcService<tonic::body::BoxBody> + Send + Sync + Clone + 'static,
T::ResponseBody: tonic::codegen::Body<Data = tonic::codegen::Bytes> + Send + 'static,
<T::ResponseBody as tonic::codegen::Body>::Error: Into<tonic::codegen::StdError> + Send,
T::Future: Send,
{
#[instrument(skip(self, digest), fields(blob.digest=%digest))]
async fn has(&self, digest: &B3Digest) -> io::Result<bool> {
let mut grpc_client = self.grpc_client.clone();
@ -337,7 +347,6 @@ mod tests {
.await
.expect("must succeed"),
);
GRPCBlobService::from_client(client)
};