feat(cubestore): EXPLAIN ANALYZE DETAILED per-query trace

rust/cubestore/cuberpc/src/lib.rs

@@ -7,15 +7,47 @@ use syn::spanned::Spanned;
 use syn::{parse_macro_input, FnArg, Item, Pat, ReturnType, TraitItem};
 
 #[proc_macro_attribute]
-pub fn service(_attr: TokenStream, input: TokenStream) -> proc_macro::TokenStream {
-    let svc = parse_macro_input!(input as RpcService);
+pub fn service(attr: TokenStream, input: TokenStream) -> proc_macro::TokenStream {
+    let args = parse_macro_input!(attr as ServiceArgs);
+    let mut svc = parse_macro_input!(input as RpcService);
+    svc.trace_guard = args.trace_guard;
 
     proc_macro::TokenStream::from(svc.into_token_stream())
 }
 
+/// Optional arguments to `#[cuberpc::service]`.
+/// `trace_guard = <path>`: generate a `Traced<Service>` decorator that wraps an
+/// `Arc<dyn Service>` and, for every async method, holds the guard returned by
+/// `<path>(method_name)` across the call. Keeps cuberpc decoupled from the host
+/// crate's tracing module.
+struct ServiceArgs {
+    trace_guard: Option<syn::Path>,
+}
+
+impl Parse for ServiceArgs {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let mut trace_guard = None;
+        if !input.is_empty() {
+            let key: Ident = input.parse()?;
+            input.parse::<syn::Token![=]>()?;
+            let path: syn::Path = input.parse()?;
+            if key == "trace_guard" {
+                trace_guard = Some(path);
+            } else {
+                return Err(syn::Error::new(
+                    key.span(),
+                    "unknown cuberpc::service argument",
+                ));
+            }
+        }
+        Ok(ServiceArgs { trace_guard })
+    }
+}
+
 struct RpcService {
     ident: Ident,
     methods: Vec<RpcMethod>,
+    trace_guard: Option<syn::Path>,
 }
 
 struct RpcMethod {
@@ -46,6 +78,7 @@ impl Parse for RpcService {
                 RpcService {
                     ident: trait_item.ident.clone(),
                     methods,
+                    trace_guard: None,
                 }
             }
             x => {
@@ -106,6 +139,62 @@ impl RpcService {
         }
     }
 
+    fn method_call_variant_name_impl(&self) -> proc_macro2::TokenStream {
+        let method_call = self.method_call_ident();
+        let arms = self
+            .methods
+            .iter()
+            .map(|m| {
+                let variant = m.variant_ident();
+                let has_args = m.args.iter().any(|a| matches!(a, FnArg::Typed(_)));
+                if has_args {
+                    quote! { #method_call::#variant(..) => stringify!(#variant) }
+                } else {
+                    quote! { #method_call::#variant => stringify!(#variant) }
+                }
+            })
+            .collect::<Vec<_>>();
+        quote! {
+            impl #method_call {
+                /// The trait method name behind this call, for tracing/metrics labels.
+                pub fn variant_name(&self) -> &'static str {
+                    match self {
+                        #( #arms ),*
+                    }
+                }
+            }
+        }
+    }
+
+    fn traced_decorator(&self) -> proc_macro2::TokenStream {
+        let Some(guard_path) = self.trace_guard.as_ref() else {
+            return quote! {};
+        };
+        let service = &self.ident;
+        let traced = format_ident!("Traced{}", self.ident);
+        let methods = self
+            .methods
+            .iter()
+            .map(|m| m.traced_method(guard_path))
+            .collect::<Vec<_>>();
+        quote! {
+            pub struct #traced {
+                inner: std::sync::Arc<dyn #service>,
+            }
+
+            impl #traced {
+                pub fn new(inner: std::sync::Arc<dyn #service>) -> std::sync::Arc<Self> {
+                    std::sync::Arc::new(Self { inner })
+                }
+            }
+
+            #[async_trait]
+            impl #service for #traced {
+                #( #methods )*
+            }
+        }
+    }
+
     fn client_transport_trait(&self) -> proc_macro2::TokenStream {
         let method_call = self.method_call_ident();
         let method_result = self.method_result_ident();
@@ -350,17 +439,53 @@ impl RpcMethod {
     fn variant_ident(&self) -> Ident {
         format_ident!("{}", self.ident.to_string().to_camel_case())
     }
+
+    fn traced_method(&self, guard_path: &syn::Path) -> proc_macro2::TokenStream {
+        let &Self {
+            ident,
+            asyncness,
+            args,
+            output,
+        } = &self;
+        let arg_names = args
+            .iter()
+            .filter_map(|a| match a {
+                FnArg::Typed(ty) => match ty.pat.as_ref() {
+                    Pat::Ident(id) => Some(id.ident.clone()),
+                    x => panic!("Unexpected pattern: {:?}", x),
+                },
+                FnArg::Receiver(_) => None,
+            })
+            .collect::<Vec<_>>();
+        let variant = self.variant_ident();
+        if *asyncness {
+            quote! {
+                async fn #ident(#( #args ),*) #output {
+                    let _g = #guard_path(stringify!(#variant));
+                    self.inner.#ident(#( #arg_names ),*).await
+                }
+            }
+        } else {
+            quote! {
+                fn #ident(#( #args ),*) #output {
+                    self.inner.#ident(#( #arg_names ),*)
+                }
+            }
+        }
+    }
 }
 
 impl ToTokens for RpcService {
     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
         tokens.extend(vec![
             self.original_trait(),
             self.method_call_enum(),
+            self.method_call_variant_name_impl(),
             self.method_result_enum(),
             self.client_transport_trait(),
             self.client_impl(),
             self.server_impl(),
+            self.traced_decorator(),
         ]);
     }
 }

rust/cubestore/cubestore/src/cluster/message.rs

@@ -1,6 +1,7 @@
 use crate::metastore::{MetaStoreRpcMethodCall, MetaStoreRpcMethodResult};
 use crate::queryplanner::query_executor::SerializedRecordBatchStream;
 use crate::queryplanner::serialized_plan::SerializedPlan;
+use crate::trace::{MainTrace, WorkerTrace};
 use crate::CubeError;
 use datafusion::arrow::datatypes::SchemaRef;
 use serde::{Deserialize, Serialize};
@@ -20,10 +21,21 @@ pub enum NetworkMessage {
     SelectResult(Result<(SchemaRef, Vec<SerializedRecordBatchStream>), CubeError>),
 
     //Perform explain analyze of worker query part and return it pretty printed physical plan
-    /// The boolean flag is whether to execute the plan to collect runtime metrics.
-    ExplainAnalyze(SerializedPlan, WorkerPlanningParams, bool),
+    ExplainAnalyze(SerializedPlan, WorkerPlanningParams),
     ExplainAnalyzeResult(Result<String, CubeError>),
 
+    /// Detailed-trace mirror of [RouterSelect]: the entry node asks a main worker to
+    /// run the full router plan for real and return the assembled `MainTrace`.
+    RouterSelectDetailed(SerializedPlan),
+    RouterSelectDetailedResult(Result<MainTrace, CubeError>),
+
+    /// Detailed-trace mirror of [Select]: a worker runs its part for real and returns
+    /// both the result rows (for the main to merge) and its `WorkerTrace`.
+    SelectDetailed(SerializedPlan, WorkerPlanningParams),
+    SelectDetailedResult(
+        Result<(SchemaRef, Vec<SerializedRecordBatchStream>, WorkerTrace), CubeError>,
+    ),
+
     /// Select that sends results in batches. The immediate response is [SelectResultSchema],
     /// followed by a stream of [SelectResultBatch].
     SelectStart(SerializedPlan, WorkerPlanningParams),