Skip to main content

hitbox_fn/
extractor.rs

1//! Cache key extraction for function arguments.
2
3use std::marker::PhantomData;
4
5use async_trait::async_trait;
6use hitbox::{Extractor, KeyPart, KeyParts};
7
8use crate::Args;
9
10/// Trait for types that can contribute to a cache key.
11///
12/// Implement this trait for your custom types to control how they appear in cache keys.
13/// Blanket implementations are provided for common primitive types.
14///
15/// # Example
16///
17/// ```ignore
18/// use hitbox_fn::KeyExtract;
19/// use hitbox_core::KeyPart;
20///
21/// struct UserId(u64);
22///
23/// impl KeyExtract for UserId {
24///     fn extract(&self) -> Vec<KeyPart> {
25///         vec![KeyPart::new("user_id", Some(self.0.to_string()))]
26///     }
27/// }
28/// ```
29///
30/// # Automatic composition
31///
32/// When all elements of a tuple implement `KeyExtract`, the tuple automatically
33/// implements it too by concatenating all key parts:
34///
35/// ```ignore
36/// // If UserId and TenantId both implement KeyExtract,
37/// // then Args<(UserId, TenantId)> automatically implements KeyExtract
38/// let args = Args((UserId(42), TenantId("acme".into())));
39/// let parts = args.extract();
40/// // parts = [KeyPart("user_id", "42"), KeyPart("tenant", "acme")]
41/// ```
42pub trait KeyExtract {
43    /// Extract key parts from this value.
44    fn extract(&self) -> Vec<KeyPart>;
45}
46
47// KeyExtract implementations for primitive types
48
49macro_rules! impl_key_extract_for_scalar {
50    ($($ty:ty => $name:expr),* $(,)?) => {
51        $(
52            impl KeyExtract for $ty {
53                fn extract(&self) -> Vec<KeyPart> {
54                    vec![KeyPart::new($name, Some(self.to_string()))]
55                }
56            }
57        )*
58    };
59}
60
61impl_key_extract_for_scalar! {
62    u8 => "u8", u16 => "u16", u32 => "u32", u64 => "u64", u128 => "u128", usize => "usize",
63    i8 => "i8", i16 => "i16", i32 => "i32", i64 => "i64", i128 => "i128", isize => "isize",
64    bool => "bool",
65}
66
67impl KeyExtract for String {
68    fn extract(&self) -> Vec<KeyPart> {
69        vec![KeyPart::new("str", Some(self.clone()))]
70    }
71}
72
73impl KeyExtract for &str {
74    fn extract(&self) -> Vec<KeyPart> {
75        vec![KeyPart::new("str", Some(self.to_string()))]
76    }
77}
78
79impl<T: KeyExtract> KeyExtract for Option<T> {
80    fn extract(&self) -> Vec<KeyPart> {
81        match self {
82            Some(v) => v.extract(),
83            None => vec![KeyPart::new("none", None::<&str>)],
84        }
85    }
86}
87
88impl<T: KeyExtract + ?Sized> KeyExtract for &T {
89    fn extract(&self) -> Vec<KeyPart> {
90        (*self).extract()
91    }
92}
93
94// KeyExtract implementations for Args<tuples>
95
96macro_rules! impl_key_extract_for_args {
97    () => {
98        impl KeyExtract for Args<()> {
99            fn extract(&self) -> Vec<KeyPart> {
100                vec![]
101            }
102        }
103    };
104    ($first:tt : $T0:ident $(, $idx:tt : $T:ident)*) => {
105        impl<$T0: KeyExtract $(, $T: KeyExtract)*> KeyExtract for Args<($T0, $($T,)*)> {
106            fn extract(&self) -> Vec<KeyPart> {
107                [self.0.$first.extract() $(, self.0.$idx.extract())*]
108                    .into_iter()
109                    .flatten()
110                    .collect()
111            }
112        }
113    };
114}
115
116impl_key_extract_for_args!();
117impl_key_extract_for_args!(0: T0);
118impl_key_extract_for_args!(0: T0, 1: T1);
119impl_key_extract_for_args!(0: T0, 1: T1, 2: T2);
120impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3);
121impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4);
122impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5);
123impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6);
124impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6, 7: T7);
125impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6, 7: T7, 8: T8);
126impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6, 7: T7, 8: T8, 9: T9);
127impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6, 7: T7, 8: T8, 9: T9, 10: T10);
128impl_key_extract_for_args!(0: T0, 1: T1, 2: T2, 3: T3, 4: T4, 5: T5, 6: T6, 7: T7, 8: T8, 9: T9, 10: T10, 11: T11);
129
130// FnExtractor - bridges KeyExtract to Extractor
131
132/// Generic extractor that uses [`KeyExtract`] trait to generate cache keys.
133///
134/// This extractor bridges the simple `KeyExtract` trait to hitbox's `Extractor` trait,
135/// adding the function path as a prefix to ensure different functions have different cache keys.
136///
137/// # Example
138///
139/// ```ignore
140/// use hitbox_fn::{Args, FnExtractor};
141///
142/// let extractor = FnExtractor::<Args<(UserId, TenantId)>>::new("my_module::fetch_user");
143/// ```
144pub struct FnExtractor<T> {
145    fn_path: &'static str,
146    _marker: PhantomData<T>,
147}
148
149impl<T> FnExtractor<T> {
150    /// Create a new function extractor with the given fully qualified function path.
151    ///
152    /// The function path is used as a prefix in the cache key to ensure different
153    /// functions produce different keys even with the same arguments.
154    pub fn new(fn_path: &'static str) -> Self {
155        Self {
156            fn_path,
157            _marker: PhantomData,
158        }
159    }
160}
161
162#[async_trait]
163impl<T> Extractor for FnExtractor<T>
164where
165    T: KeyExtract + Send + Sync,
166{
167    type Subject = T;
168
169    async fn get(&self, subject: Self::Subject) -> KeyParts<Self::Subject> {
170        let extracted = subject.extract();
171        let mut parts = KeyParts::new(subject);
172        parts.push(KeyPart::new("fn", Some(self.fn_path)));
173        for part in extracted {
174            parts.push(part);
175        }
176        parts
177    }
178}