feat(tvix/store): implement reflection r/5557

This implements grpc.reflection.v1alpha.ServerReflection, and will make tools
like evans automatically discover available services, without having to
specify the path to the .proto files client-side.

It's behind a reflection feature flag, which is enabled by default.

Change-Id: Icbcb5eb05ceede5b9952e38a2ba72eaa6fa8a437
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7435
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

5 files changed, 68 insertions, 25 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 5cae4e0dfcd7..90672b6d037f 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -15,9 +15,21 @@ tokio = { version = "1.23.0", features = ["rt-multi-thread"] }
 tokio-stream = "0.1.11"
 tonic = "0.8.2"
 
+[dependencies.tonic-reflection]
+optional = true
+version = "0.5.0"
+
 [build-dependencies]
 prost-build = "0.11.2"
 tonic-build = "0.8.2"
 
 [dev-dependencies]
 test-case = "2.2.2"
+
+[features]
+default = [
+  "reflection"
+]
+reflection = [
+  "tonic-reflection"
+]
diff --git a/tvix/store/build.rs b/tvix/store/build.rs
index 3fdc7418c1c5..ef1674cfbd0b 100644
--- a/tvix/store/build.rs
+++ b/tvix/store/build.rs
@@ -1,24 +1,32 @@
 use std::io::Result;
 
 fn main() -> Result<()> {
-    tonic_build::configure()
-        .build_server(true)
-        .build_client(true)
-        .compile(
-            &[
-                "tvix/store/protos/castore.proto",
-                "tvix/store/protos/pathinfo.proto",
-                "tvix/store/protos/rpc_blobstore.proto",
-                "tvix/store/protos/rpc_directory.proto",
-                "tvix/store/protos/rpc_pathinfo.proto",
-            ],
-            // If we are in running `cargo build` manually, using `../..` works fine,
-            // but in case we run inside a nix build, we need to instead point PROTO_ROOT
-            // to a sparseTree containing that structure.
-            &[match std::env::var_os("PROTO_ROOT") {
-                Some(proto_root) => proto_root.to_str().unwrap().to_owned(),
-                None => "../..".to_string(),
-            }],
-        )?;
+    #[allow(unused_mut)]
+    let mut builder = tonic_build::configure();
+
+    #[cfg(feature = "reflection")]
+    {
+        let out_dir = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
+        let descriptor_path = out_dir.join("tvix.store.v1.bin");
+
+        builder = builder.file_descriptor_set_path(&descriptor_path);
+    };
+
+    builder.build_server(true).build_client(true).compile(
+        &[
+            "tvix/store/protos/castore.proto",
+            "tvix/store/protos/pathinfo.proto",
+            "tvix/store/protos/rpc_blobstore.proto",
+            "tvix/store/protos/rpc_directory.proto",
+            "tvix/store/protos/rpc_pathinfo.proto",
+        ],
+        // If we are in running `cargo build` manually, using `../..` works fine,
+        // but in case we run inside a nix build, we need to instead point PROTO_ROOT
+        // to a sparseTree containing that structure.
+        &[match std::env::var_os("PROTO_ROOT") {
+            Some(proto_root) => proto_root.to_str().unwrap().to_owned(),
+            None => "../..".to_string(),
+        }],
+    )?;
     Ok(())
 }
diff --git a/tvix/store/default.nix b/tvix/store/default.nix
index 1dbf01923545..68b529e723b1 100644
--- a/tvix/store/default.nix
+++ b/tvix/store/default.nix
@@ -20,6 +20,10 @@ depot.tvix.crates.workspaceMembers.tvix-store.build.override {
       nativeBuildInputs = protobufDep prev;
     };
 
+    tonic-reflection = prev: {
+      nativeBuildInputs = protobufDep prev;
+    };
+
     tvix-store = prev: {
       PROTO_ROOT = protoRoot;
       nativeBuildInputs = protobufDep prev;
diff --git a/tvix/store/src/main.rs b/tvix/store/src/main.rs
index cc9002438bce..b28c3fdeefe8 100644
--- a/tvix/store/src/main.rs
+++ b/tvix/store/src/main.rs
@@ -2,6 +2,9 @@ use crate::proto::blob_service_server::BlobServiceServer;
 use crate::proto::directory_service_server::DirectoryServiceServer;
 use crate::proto::path_info_service_server::PathInfoServiceServer;
 
+#[cfg(feature = "reflection")]
+use crate::proto::FILE_DESCRIPTOR_SET;
+
 use clap::Parser;
 use tonic::{transport::Server, Result};
 
@@ -31,18 +34,28 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .parse()
         .unwrap();
 
+    let mut server = Server::builder();
+
     let blob_service = dummy_blob_service::DummyBlobService {};
     let directory_service = dummy_directory_service::DummyDirectoryService {};
     let path_info_service = dummy_path_info_service::DummyPathInfoService {};
 
-    println!("tvix-store listening on {}", listen_address);
-
-    Server::builder()
+    let mut router = server
         .add_service(BlobServiceServer::new(blob_service))
         .add_service(DirectoryServiceServer::new(directory_service))
-        .add_service(PathInfoServiceServer::new(path_info_service))
-        .serve(listen_address)
-        .await?;
+        .add_service(PathInfoServiceServer::new(path_info_service));
+
+    #[cfg(feature = "reflection")]
+    {
+        let reflection_svc = tonic_reflection::server::Builder::configure()
+            .register_encoded_file_descriptor_set(FILE_DESCRIPTOR_SET)
+            .build()?;
+        router = router.add_service(reflection_svc);
+    }
+
+    println!("tvix-store listening on {}", listen_address);
+
+    router.serve(listen_address).await?;
 
     Ok(())
 }
diff --git a/tvix/store/src/proto.rs b/tvix/store/src/proto.rs
index eac37820658e..75e5c34ab137 100644
--- a/tvix/store/src/proto.rs
+++ b/tvix/store/src/proto.rs
@@ -7,6 +7,12 @@ use prost::Message;
 
 tonic::include_proto!("tvix.store.v1");
 
+#[cfg(feature = "reflection")]
+/// Compiled file descriptors for implementing [gRPC
+/// reflection](https://github.com/grpc/grpc/blob/master/doc/server-reflection.md) with e.g.
+/// [`tonic_reflection`](https://docs.rs/tonic-reflection).
+pub const FILE_DESCRIPTOR_SET: &[u8] = tonic::include_file_descriptor_set!("tvix.store.v1");
+
 /// Errors that can occur during the validation of Directory messages.
 #[derive(Debug, PartialEq, Eq, Error)]
 pub enum ValidateDirectoryError {