diff options
author | Florian Klink <flokli@flokli.de> | 2023-05-22T16·40+0300 |
---|---|---|
committer | clbot <clbot@tvl.fyi> | 2023-05-24T09·01+0000 |
commit | 457fd4c325e85d07f3034d82f96f897bc63b8462 (patch) | |
tree | b9ff64b5db67a944f90f2dd4d4723dcdce96fe3c | |
parent | 01e8fa65294b5ff8438e1abc0e70e9fe7e2da05c (diff) |
feat(tvix/store/nar): add GRPCNARCalculationService r/6184
This asks a remote tvix-store for the nar size and digest of a given root node. Change-Id: If9f916d9bfc5f8dc3166e2c6c1671c0f0124d1c1 Reviewed-on: https://cl.tvl.fyi/c/depot/+/8611 Autosubmit: flokli <flokli@flokli.de> Reviewed-by: tazjin <tazjin@tvl.su> Tested-by: BuildkiteCI
-rw-r--r-- | tvix/store/src/nar/grpc_nar_calculation_service.rs | 69 | ||||
-rw-r--r-- | tvix/store/src/nar/mod.rs | 2 |
2 files changed, 71 insertions, 0 deletions
diff --git a/tvix/store/src/nar/grpc_nar_calculation_service.rs b/tvix/store/src/nar/grpc_nar_calculation_service.rs new file mode 100644 index 000000000000..429593743914 --- /dev/null +++ b/tvix/store/src/nar/grpc_nar_calculation_service.rs @@ -0,0 +1,69 @@ +use super::NARCalculationService; +use crate::proto; +use tonic::transport::Channel; +use tonic::Status; + +/// A NAR calculation service which asks a remote tvix-store for NAR calculation +/// (via the gRPC PathInfoService). +#[derive(Clone)] +pub struct GRPCNARCalculationService { + /// A handle into the active tokio runtime. Necessary to spawn tasks. + tokio_handle: tokio::runtime::Handle, + + /// The internal reference to a gRPC client. + /// Cloning it is cheap, and it internally handles concurrent requests. + grpc_client: proto::path_info_service_client::PathInfoServiceClient<Channel>, +} + +impl GRPCNARCalculationService { + /// construct a new [GRPCNARCalculationService], by passing a handle to the + /// tokio runtime, and a gRPC client. + pub fn new( + tokio_handle: tokio::runtime::Handle, + grpc_client: proto::path_info_service_client::PathInfoServiceClient<Channel>, + ) -> Self { + Self { + tokio_handle, + grpc_client, + } + } + + /// construct a [GRPCNARCalculationService], from a [proto::path_info_service_client::PathInfoServiceClient<Channel>]. + /// panics if called outside the context of a tokio runtime. + pub fn from_client( + grpc_client: proto::path_info_service_client::PathInfoServiceClient<Channel>, + ) -> Self { + Self { + tokio_handle: tokio::runtime::Handle::current(), + grpc_client, + } + } +} + +impl NARCalculationService for GRPCNARCalculationService { + fn calculate_nar( + &self, + root_node: &proto::node::Node, + ) -> Result<(u64, [u8; 32]), super::RenderError> { + // Get a new handle to the gRPC client, and copy the root node. + let mut grpc_client = self.grpc_client.clone(); + let root_node = root_node.clone(); + + let task: tokio::task::JoinHandle<Result<_, Status>> = + self.tokio_handle.spawn(async move { + Ok(grpc_client + .calculate_nar(proto::Node { + node: Some(root_node), + }) + .await? + .into_inner()) + }); + + match self.tokio_handle.block_on(task).unwrap() { + Ok(resp) => Ok((resp.nar_size, resp.nar_sha256.to_vec().try_into().unwrap())), + Err(e) => Err(super::RenderError::StoreError(crate::Error::StorageError( + e.to_string(), + ))), + } + } +} diff --git a/tvix/store/src/nar/mod.rs b/tvix/store/src/nar/mod.rs index 2bfb541733be..a29cc5451bae 100644 --- a/tvix/store/src/nar/mod.rs +++ b/tvix/store/src/nar/mod.rs @@ -2,9 +2,11 @@ use crate::{proto, B3Digest}; use data_encoding::BASE64; use thiserror::Error; +mod grpc_nar_calculation_service; mod non_caching_calculation_service; mod renderer; +pub use grpc_nar_calculation_service::GRPCNARCalculationService; pub use non_caching_calculation_service::NonCachingNARCalculationService; pub use renderer::NARRenderer; |