Instrument gRPC method handlers using their method name (#35682)

GitOrigin-RevId: baaaa9a8cdd70ce127d93ad1915981e93627b332

Author: goffrie

Cargo.lock

@@ -52,6 +52,7 @@ openidconnect = { workspace = true }
 packed_value = { path = "../packed_value" }
 parking_lot = { workspace = true }
 pb = { path = "../pb" }
+pb_extras = { path = "../pb_extras" }
 pin-project = { workspace = true }
 prometheus = { workspace = true }
 prometheus-hyper = { workspace = true }

crates/common/Cargo.toml

@@ -1,14 +1,18 @@
 use std::{
     convert::Infallible,
     net::SocketAddr,
+    sync::Arc,
     task::{
         Context,
         Poll,
     },
 };
 
+use fnv::FnvHashMap;
 use futures::Future;
+use http::Request;
 use pb::error_metadata::ErrorMetadataStatusExt;
+use pb_extras::ReflectionService;
 use sentry::integrations::tower as sentry_tower;
 use tokio::net::TcpSocket;
 use tokio_metrics::Instrumented;
@@ -35,8 +39,11 @@ use crate::{
     runtime::TaskManager,
 };
 
+// maps the full route `/service.Service/Method` to just `Method`
+type KnownMethods = FnvHashMap<String, &'static str>;
 pub struct ConvexGrpcService {
     routes: Routes,
+    known_methods: KnownMethods,
     health_reporter: HealthReporter,
     service_names: Vec<&'static str>,
 }
@@ -47,6 +54,7 @@ impl ConvexGrpcService {
         let routes = Routes::new(health_service);
         Self {
             routes,
+            known_methods: FnvHashMap::default(),
             health_reporter,
             service_names: Vec::new(),
         }
@@ -58,7 +66,7 @@ impl ConvexGrpcService {
                 http::Request<tonic::body::BoxBody>,
                 Response = http::Response<tonic::body::BoxBody>,
                 Error = Infallible,
-            > + NamedService
+            > + ReflectionService
             + Clone
             + Send
             + 'static,
@@ -69,15 +77,20 @@ impl ConvexGrpcService {
         // line with all names when we start serving.
         let service_name = <S as NamedService>::NAME;
         self.service_names.push(service_name);
+        for method_name in S::METHODS {
+            self.known_methods
+                .insert(format!("/{service_name}/{method_name}"), method_name);
+        }
         self
     }
 
     pub async fn serve<F>(mut self, addr: SocketAddr, shutdown: F) -> anyhow::Result<()>
     where
         F: Future<Output = ()>,
     {
+        let known_methods = Arc::new(self.known_methods);
         let convex_layers = ServiceBuilder::new()
-            .layer_fn(TokioInstrumentationService::new)
+            .layer_fn(move |s| TokioInstrumentationService::new(known_methods.clone(), s))
             .layer(sentry_tower::NewSentryLayer::new_from_top())
             .layer(sentry_tower::SentryHttpLayer::with_transaction());
 
@@ -118,18 +131,22 @@ pub fn handle_response<T>(response: Result<Response<T>, Status>) -> anyhow::Resu
 
 #[derive(Clone)]
 struct TokioInstrumentationService<S> {
+    known_methods: Arc<KnownMethods>,
     inner: S,
 }
 
 impl<S> TokioInstrumentationService<S> {
-    fn new(inner: S) -> Self {
-        Self { inner }
+    fn new(known_methods: Arc<KnownMethods>, inner: S) -> Self {
+        Self {
+            known_methods,
+            inner,
+        }
     }
 }
 
-impl<S, Request> Service<Request> for TokioInstrumentationService<S>
+impl<S, T> Service<Request<T>> for TokioInstrumentationService<S>
 where
-    S: Service<Request>,
+    S: Service<Request<T>>,
 {
     type Error = S::Error;
     type Future = Instrumented<S::Future>;
@@ -139,7 +156,12 @@ where
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, req: Request) -> Self::Future {
-        TaskManager::instrument("grpc_handler", self.inner.call(req))
+    fn call(&mut self, req: Request<T>) -> Self::Future {
+        let name = self
+            .known_methods
+            .get(req.uri().path())
+            .copied()
+            .unwrap_or("grpc_handler");
+        TaskManager::instrument(name, self.inner.call(req))
     }
 }

crates/common/src/grpc/mod.rs

@@ -12,6 +12,7 @@ anyhow = { workspace = true }
 convex_sync_types = { path = "../convex/sync_types" }
 errors = { path = "../errors" }
 http = { workspace = true }
+pb_extras = { path = "../pb_extras" }
 prost = { workspace = true }
 prost-reflect = { workspace = true }
 prost-types = { workspace = true }

crates/pb/Cargo.toml

@@ -40,6 +40,7 @@ pub mod usage {
     include!(concat!(env!("OUT_DIR"), "/usage.rs"));
 }
 
+include!(concat!(env!("OUT_DIR"), "/_extras.rs"));
 use std::sync::LazyLock;
 
 use prost_reflect::DescriptorPool;

crates/pb/src/lib.rs

@@ -9,6 +9,7 @@ doctest = false
 
 [dependencies]
 cfg-if = { workspace = true }
+prost = { workspace = true }
 tonic-build = { workspace = true }
 
 [package.metadata.cargo-machete]

crates/pb_build/Cargo.toml

@@ -1,5 +1,6 @@
 use std::{
     ffi::OsStr,
+    fmt::Write as _,
     fs,
     io::Result,
     path::{
@@ -8,6 +9,9 @@ use std::{
     },
 };
 
+use prost::Message;
+use tonic_build::FileDescriptorSet;
+
 cfg_if::cfg_if! {
     if #[cfg(target_os = "macos")] {
         const PROTOC_BINARY_NAME: &str = "protoc-macos-universal";
@@ -76,8 +80,8 @@ pub fn pb_build(features: Vec<&'static str>, mut extra_includes: Vec<&'static st
     let mut includes = vec!["protos/"];
     includes.append(&mut extra_includes);
 
-    let mut builder =
-        tonic_build::configure().file_descriptor_set_path(out_dir.join("descriptors.bin"));
+    let descriptor_set_path = out_dir.join("descriptors.bin");
+    let mut builder = tonic_build::configure().file_descriptor_set_path(&descriptor_set_path);
     for (proto_path, rust_path) in external_paths {
         builder = builder.extern_path(proto_path, rust_path);
     }
@@ -99,6 +103,47 @@ pub fn pb_build(features: Vec<&'static str>, mut extra_includes: Vec<&'static st
     }
     mods.sort();
 
+    // Read back the file descriptor set to codegen `ReflectionService` impls.
+    let file_descriptor_set = FileDescriptorSet::decode(&std::fs::read(&descriptor_set_path)?[..])?;
+    let mut extras = String::new();
+    for file in file_descriptor_set.file {
+        let Some(package_name) = file.package else {
+            continue;
+        };
+        if !packages.contains(&package_name) {
+            continue;
+        }
+        for service in file.service {
+            let Some(service_name) = service.name else {
+                continue;
+            };
+            let lower_name = naive_snake_case(&service_name);
+            let server_mod = format!("crate::{package_name}::{lower_name}_server");
+            write!(
+                &mut extras,
+                r#"impl<T> pb_extras::ReflectionService for {server_mod}::{service_name}Server<T> {{
+    const METHODS: &[&str] = &[
+"#
+            )
+            .unwrap();
+            for method in service.method {
+                let Some(method_name) = method.name else {
+                    continue;
+                };
+                writeln!(&mut extras, "        {method_name:?},").unwrap();
+            }
+            write!(
+                &mut extras,
+                r#"    ];
+}}
+"#
+            )
+            .unwrap();
+        }
+    }
+
+    std::fs::write(out_dir.join("_extras.rs"), extras)?;
+
     // Now let's build the lib.rs file.
     let mut lib_file_contents = String::new();
     lib_file_contents.push_str("// @generated - do not modify. Modify build.rs instead.\n");
@@ -118,6 +163,7 @@ pub fn pb_build(features: Vec<&'static str>, mut extra_includes: Vec<&'static st
 
     lib_file_contents.push_str(
         r#"
+include!(concat!(env!("OUT_DIR"), "/_extras.rs"));
 use std::sync::LazyLock;
 
 use prost_reflect::DescriptorPool;
@@ -135,3 +181,20 @@ pub static DESCRIPTOR_POOL: LazyLock<DescriptorPool> =
 
     Ok(())
 }
+
+// copied from `tonic-build`
+fn naive_snake_case(name: &str) -> String {
+    let mut s = String::new();
+    let mut it = name.chars().peekable();
+
+    while let Some(x) = it.next() {
+        s.push(x.to_ascii_lowercase());
+        if let Some(y) = it.peek() {
+            if y.is_uppercase() {
+                s.push('_');
+            }
+        }
+    }
+
+    s
+}

crates/pb_build/src/lib.rs

@@ -0,0 +1,11 @@
+[package]
+name = "pb_extras"
+version = "0.1.0"
+edition = "2021"
+license = "LicenseRef-FSL-1.1-Apache-2.0"
+
+[dependencies]
+tonic = { workspace = true }
+
+[lints]
+workspace = true

crates/pb_extras/Cargo.toml

@@ -0,0 +1,6 @@
+use tonic::server::NamedService;
+
+pub trait ReflectionService: NamedService {
+    /// The list of methods supported by this service, e.g. "ExecuteVectorQuery"
+    const METHODS: &[&str];
+}