Rust异步函数指标埋点问题：同步/异步结果统计异常

要解决异步函数的类型不匹配问题，核心是区分同步与异步函数的返回值结构：同步函数直接返回Result<T, E>，而异步函数（经#[tonic::async_trait]处理后）返回的是包裹着Result<T, E>的Pin<Box<dyn Future<Output = Result<T, E>>>>。我们需要针对两种情况分别生成结果拦截逻辑。

实现步骤 #

• 检测函数类型：在过程宏中解析函数时，判断是否带有async关键字，区分同步/异步场景。
• 同步函数逻辑：直接包裹原函数调用，匹配返回的Result更新计数器，逻辑和你现有代码一致。
• 异步函数逻辑：对原异步函数返回的Future进行包装，在await获取结果后匹配Result更新计数器，再将原结果返回，确保最终返回的Future类型符合要求。

示例代码 #

假设使用syn和quote编写过程宏，核心实现如下：

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn};

#[proc_macro_attribute]
pub fn collect_metrics(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    let fn_name = &input.sig.ident;
    let fn_body = &input.block;
    let asyncness = &input.sig.asyncness;
    let inputs = &input.sig.inputs;
    let output = &input.sig.output;

    // 根据是否异步生成不同逻辑
    let expanded = if asyncness.is_some() {
        // 异步函数处理：包装Future，await后处理结果
        quote! {
            #input

            #asyncness fn #fn_name(#(#inputs),*) #output {
                let inner_future = async move {
                    #fn_body
                };
                async move {
                    let result = inner_future.await;
                    // 替换为你的实际指标统计逻辑
                    match &result {
                        Ok(_) => {
                            metrics::inc_success(stringify!(#fn_name));
                        }
                        Err(_) => {
                            metrics::inc_failure(stringify!(#fn_name));
                        }
                    }
                    result
                }.boxed()
            }
        }
    } else {
        // 同步函数处理：直接调用后匹配结果
        quote! {
            #input

            fn #fn_name(#(#inputs),*) #output {
                let result = {
                    #fn_body
                };
                // 替换为你的实际指标统计逻辑
                match &result {
                    Ok(_) => {
                        metrics::inc_success(stringify!(#fn_name));
                    }
                    Err(_) => {
                        metrics::inc_failure(stringify!(#fn_name));
                    }
                }
                result
            }
        }
    };

    TokenStream::from(expanded)
}

关键说明 #

• 异步函数的类型适配：通过外层async块包裹原Future，await后处理结果，最后调用.boxed()（需引入futures::future::FutureExt）转换成Pin<Box<dyn Future<...>>>，完全匹配tonic::async_trait要求的返回类型。
• 指标统计逻辑：示例中的metrics::inc_success/metrics::inc_failure需替换为你实际使用的指标库调用（比如Prometheus客户端的计数器递增方法）。
• 签名兼容性：生成的代码会完整保留原函数的签名、参数和返回值类型，可直接用于#[tonic::async_trait]装饰的trait实现方法。

依赖配置 #

需要在Cargo.toml中添加以下依赖：

[build-dependencies]
proc-macro2 = "1.0"
quote = "1.0"
syn = { version = "2.0", features = ["full"] }

[dependencies]
futures = "0.3" # 用于Future的boxed方法
# 你的指标库依赖，比如 prometheus = "0.13"

内容的提问来源于stack exchange，提问作者Julian Duru