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

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;