feat(tvix/store): implement reflection

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

tvix/Cargo.lock

@@ -1987,6 +1987,21 @@ dependencies = [
  "syn 1.0.103",
 ]
 
+[[package]]
+name = "tonic-reflection"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0455f730d540a1484bffc3c55c94100b18a662597b982c2e9073f2c55c602616"
+dependencies = [
+ "bytes",
+ "prost",
+ "prost-types",
+ "tokio",
+ "tokio-stream",
+ "tonic",
+ "tonic-build",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
@@ -2156,6 +2171,7 @@ dependencies = [
  "tokio-stream",
  "tonic",
  "tonic-build",
+ "tonic-reflection",
 ]
 
 [[package]]

tvix/Cargo.nix

@@ -3792,6 +3792,7 @@ rec {
           "default" = [ "std" ];
           "std" = [ "prost/std" ];
         };
+        resolvedDefaultFeatures = [ "default" "std" ];
       };
       "quick-error" = rec {
         crateName = "quick-error";
@@ -5501,7 +5502,7 @@ rec {
           "time" = [ "tokio/time" ];
           "tokio-util" = [ "dep:tokio-util" ];
         };
-        resolvedDefaultFeatures = [ "default" "time" ];
+        resolvedDefaultFeatures = [ "default" "net" "time" ];
       };
       "tokio-util" = rec {
         crateName = "tokio-util";
@@ -5770,6 +5771,53 @@ rec {
         };
         resolvedDefaultFeatures = [ "default" "prost" "prost-build" "transport" ];
       };
+      "tonic-reflection" = rec {
+        crateName = "tonic-reflection";
+        version = "0.5.0";
+        edition = "2018";
+        sha256 = "05i6c1fcbwkkj0p2r63vb5iac60b22a5rif3zx5li8a0slqgfm84";
+        authors = [
+          "James Nugent <james@jen20.com>"
+          "Samani G. Gikandi <samani@gojulas.com>"
+        ];
+        dependencies = [
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "prost";
+            packageId = "prost";
+          }
+          {
+            name = "prost-types";
+            packageId = "prost-types";
+          }
+          {
+            name = "tokio";
+            packageId = "tokio";
+            features = [ "sync" ];
+          }
+          {
+            name = "tokio-stream";
+            packageId = "tokio-stream";
+            features = [ "net" ];
+          }
+          {
+            name = "tonic";
+            packageId = "tonic";
+            features = [ "codegen" "prost" ];
+          }
+        ];
+        buildDependencies = [
+          {
+            name = "tonic-build";
+            packageId = "tonic-build";
+            features = [ "transport" "prost" ];
+          }
+        ];
+
+      };
       "tower" = rec {
         crateName = "tower";
         version = "0.4.13";
@@ -6429,6 +6477,11 @@ rec {
             name = "tonic";
             packageId = "tonic";
           }
+          {
+            name = "tonic-reflection";
+            packageId = "tonic-reflection";
+            optional = true;
+          }
         ];
         buildDependencies = [
           {
@@ -6446,7 +6499,12 @@ rec {
             packageId = "test-case";
           }
         ];
-
+        features = {
+          "default" = [ "reflection" ];
+          "reflection" = [ "tonic-reflection" ];
+          "tonic-reflection" = [ "dep:tonic-reflection" ];
+        };
+        resolvedDefaultFeatures = [ "default" "reflection" "tonic-reflection" ];
       };
       "typenum" = rec {
         crateName = "typenum";

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"
+]

tvix/store/build.rs

@@ -1,24 +1,32 @@
 use std::io::Result;
 
 fn main() -> Result<()> {
-    tonic_build::configure()
+    #[allow(unused_mut)]
-        .build_server(true)
+    let mut builder = tonic_build::configure();
-        .build_client(true)
+
-        .compile(
+    #[cfg(feature = "reflection")]
-            &[
+    {
-                "tvix/store/protos/castore.proto",
+        let out_dir = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
-                "tvix/store/protos/pathinfo.proto",
+        let descriptor_path = out_dir.join("tvix.store.v1.bin");
-                "tvix/store/protos/rpc_blobstore.proto",
+
-                "tvix/store/protos/rpc_directory.proto",
+        builder = builder.file_descriptor_set_path(&descriptor_path);
-                "tvix/store/protos/rpc_pathinfo.proto",
+    };
-            ],
+
-            // If we are in running `cargo build` manually, using `../..` works fine,
+    builder.build_server(true).build_client(true).compile(
-            // but in case we run inside a nix build, we need to instead point PROTO_ROOT
+        &[
-            // to a sparseTree containing that structure.
+            "tvix/store/protos/castore.proto",
-            &[match std::env::var_os("PROTO_ROOT") {
+            "tvix/store/protos/pathinfo.proto",
-                Some(proto_root) => proto_root.to_str().unwrap().to_owned(),
+            "tvix/store/protos/rpc_blobstore.proto",
-                None => "../..".to_string(),
+            "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(())
 }

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;

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);
+    let mut router = server
-
-    Server::builder()
         .add_service(BlobServiceServer::new(blob_service))
         .add_service(DirectoryServiceServer::new(directory_service))
-        .add_service(PathInfoServiceServer::new(path_info_service))
+        .add_service(PathInfoServiceServer::new(path_info_service));
-        .serve(listen_address)
+
-        .await?;
+    #[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(())
 }

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 {