Skip to main content

kameo_persistence_macros/
lib.rs

1use heck::ToShoutySnakeCase;
2use proc_macro::TokenStream;
3use quote::quote;
4use syn::{parse_macro_input, DeriveInput};
5
6#[proc_macro_derive(PersistentActor, attributes(snapshot))]
7pub fn derive_persistent_actor(input: TokenStream) -> TokenStream {
8    let input = parse_macro_input!(input as DeriveInput);
9
10    let name = &input.ident;
11    let snapshot_type = find_snapshot_type(&input);
12
13    let regiestry_ident = syn::Ident::new(
14        &format!("{}_REGISTRY", name.to_string().to_shouty_snake_case()),
15        name.span(),
16    );
17
18    let expanded = quote! {
19
20        static #regiestry_ident: ::std::sync::LazyLock<::std::sync::RwLock<::kameo_persistence::BiHashMap<::url::Url, ::kameo::prelude::WeakActorRef<#name>>>> =
21            ::std::sync::LazyLock::new(|| ::std::sync::RwLock::new(::kameo_persistence::BiHashMap::new()));
22
23
24        impl ::kameo_persistence::PersistentActor for #name {
25            type Snapshot = #snapshot_type;
26
27
28            fn register_persistent(persistence_key: ::url::Url, actor_ref: &::kameo::prelude::ActorRef<Self>) -> ::anyhow::Result<()> {
29                let Ok(mut registry) = #regiestry_ident.write() else {
30                    ::anyhow::bail!("Failed to acquire write lock on registry");
31                };
32                if let Some(old_pair) = registry.insert(persistence_key, actor_ref.downgrade()) {
33                    #[cfg(feature = "tracing")]
34                    ::tracing::warn!("Existing persistent actor reference for {old_pair:?} is replaced");
35                }
36                Ok(())
37            }
38
39            fn persistence_key(actor_ref: &::kameo::prelude::ActorRef<Self>) -> Option<::url::Url> {
40                let registry = #regiestry_ident.read().unwrap();
41                registry.get_left(&actor_ref.downgrade()).cloned()
42            }
43
44            fn lookup_persistent(persistence_key: &::url::Url) -> Option<::kameo::prelude::ActorRef<Self>> {
45                let registry = #regiestry_ident.read().unwrap();
46                registry
47                    .get_right(persistence_key)
48                    .and_then(|weak_ref| weak_ref.upgrade())
49            }
50        }
51    };
52
53    TokenStream::from(expanded)
54}
55
56fn find_snapshot_type(input: &DeriveInput) -> syn::Type {
57    // Look for #[snapshot(Type)] attribute
58    for attr in &input.attrs {
59        if attr.path().is_ident("snapshot") {
60            if let Ok(snapshot_type) = attr.parse_args::<syn::Type>() {
61                return snapshot_type;
62            }
63        }
64    }
65
66    syn::parse_quote! { <Self as ::kameo::prelude::Actor>::Args }
67}