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

2 files changed, 25 insertions, 6 deletions

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 {