From f04ae4d2ae51d55843699e4115e089eda354fbe9 Mon Sep 17 00:00:00 2001
From: Archisman Mridha <archismanmridha12345@gmail.com>
Date: Mon, 12 Feb 2024 17:22:29 +0530
Subject: [PATCH] Allow specifying `struct_name` when annotating a method
 (#171)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Specifying the struct name is useful for cases where the autometrics annotation isn’t used on a complete impl block

---------

Co-authored-by: Gerry Agbobada <gerry@fiberplane.com>
Fixes: #139
---
 autometrics-macros/src/lib.rs         |  2 +-
 autometrics-macros/src/parse.rs       | 11 ++++++++++-
 autometrics/tests/integration_test.rs | 23 +++++++++++++++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/autometrics-macros/src/lib.rs b/autometrics-macros/src/lib.rs
index 984d716..3978f2c 100644
--- a/autometrics-macros/src/lib.rs
+++ b/autometrics-macros/src/lib.rs
@@ -29,7 +29,7 @@ pub fn autometrics(
     let item = parse_macro_input!(item as Item);
 
     let result = match item {
-        Item::Function(item) => instrument_function(&args, item, None),
+        Item::Function(item) => instrument_function(&args, item, args.struct_name.as_deref()),
         Item::Impl(item) => instrument_impl_block(&args, item, &async_trait),
     };
 
diff --git a/autometrics-macros/src/parse.rs b/autometrics-macros/src/parse.rs
index e7c015d..cf7f7b8 100644
--- a/autometrics-macros/src/parse.rs
+++ b/autometrics-macros/src/parse.rs
@@ -1,5 +1,5 @@
 use syn::parse::{Parse, ParseStream};
-use syn::{Expr, ItemFn, ItemImpl, Result, Token};
+use syn::{Expr, ItemFn, ItemImpl, LitStr, Result, Token};
 
 mod kw {
     syn::custom_keyword!(track_concurrency);
@@ -8,6 +8,7 @@ mod kw {
     syn::custom_keyword!(latency);
     syn::custom_keyword!(ok_if);
     syn::custom_keyword!(error_if);
+    syn::custom_keyword!(struct_name);
 }
 
 /// Autometrics can be applied to individual functions or to
@@ -32,6 +33,9 @@ pub(crate) struct AutometricsArgs {
     pub ok_if: Option<Expr>,
     pub error_if: Option<Expr>,
     pub objective: Option<Expr>,
+
+    // Fix for https://github.com/autometrics-dev/autometrics-rs/issues/139.
+    pub struct_name: Option<String>,
 }
 
 impl Parse for AutometricsArgs {
@@ -67,6 +71,11 @@ impl Parse for AutometricsArgs {
                     return Err(input.error("expected only a single `objective` argument"));
                 }
                 args.objective = Some(input.parse()?);
+            } else if lookahead.peek(kw::struct_name) {
+                let _ = input.parse::<kw::struct_name>()?;
+                let _ = input.parse::<Token![=]>()?;
+                let struct_name = input.parse::<LitStr>()?.value();
+                args.struct_name = Some(struct_name);
             } else if lookahead.peek(Token![,]) {
                 let _ = input.parse::<Token![,]>()?;
             } else {
diff --git a/autometrics/tests/integration_test.rs b/autometrics/tests/integration_test.rs
index bedaaf3..2cb0983 100644
--- a/autometrics/tests/integration_test.rs
+++ b/autometrics/tests/integration_test.rs
@@ -71,6 +71,29 @@ fn impl_block() {
         && line.ends_with("} 1")));
 }
 
+#[test]
+fn struct_name_autometrics_macro_attribute() {
+    prometheus_exporter::try_init().ok();
+
+    struct Bar;
+
+    impl Bar {
+        #[autometrics(struct_name = "Bar")]
+        fn test_fn() -> &'static str {
+            "Hello world!"
+        }
+    }
+
+    Bar::test_fn();
+
+    let metrics = prometheus_exporter::encode_to_string().unwrap();
+    assert!(metrics.lines().any(|line| {
+        line.starts_with("function_calls_total{")
+            && line.contains(r#"function="Bar::test_fn""#)
+            && line.ends_with("} 1")
+    }));
+}
+
 #[test]
 fn result() {
     prometheus_exporter::try_init().ok();