hyperactor_macros/
lib.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Defines macros used by the [`hyperactor`] crate.
10
11#![feature(proc_macro_def_site)]
12#![deny(missing_docs)]
13
14extern crate proc_macro;
15
16use convert_case::Case;
17use convert_case::Casing;
18use indoc::indoc;
19use proc_macro::TokenStream;
20use proc_macro2::Span;
21use quote::format_ident;
22use quote::quote;
23use syn::Attribute;
24use syn::Data;
25use syn::DeriveInput;
26use syn::Expr;
27use syn::Field;
28use syn::Ident;
29use syn::ItemFn;
30use syn::ItemImpl;
31use syn::Token;
32use syn::Type;
33use syn::WherePredicate;
34use syn::bracketed;
35use syn::parse::Parse;
36use syn::parse::ParseStream;
37use syn::parse_macro_input;
38use syn::punctuated::Punctuated;
39use syn::spanned::Spanned;
40
41const REPLY_VARIANT_ERROR: &str = indoc! {r#"
42`call` message expects a typed port ref (`OncePortRef` or `PortRef`) or handle (`OncePortHandle` or `PortHandle`) argument in the last position
43
44= help: use `MyCall(Arg1Type, Arg2Type, .., OncePortRef<ReplyType>)`
45= help: use `MyCall(Arg1Type, Arg2Type, .., OncePortHandle<ReplyType>)`
46"#};
47
48const REPLY_USAGE_ERROR: &str = indoc! {r#"
49`call` message expects at most one `reply` argument
50
51= help: use `MyCall(Arg1Type, Arg2Type, .., #[reply] OncePortRef<ReplyType>)`
52= help: use `MyCall(Arg1Type, Arg2Type, .., #[reply] OncePortHandle<ReplyType>)`
53"#};
54
55enum FieldFlag {
56    None,
57    Reply,
58}
59
60/// Represents a variant of an enum.
61#[allow(dead_code)]
62enum Variant {
63    /// A named variant (i.e., `MyVariant { .. }`).
64    Named {
65        enum_name: Ident,
66        name: Ident,
67        field_names: Vec<Ident>,
68        field_types: Vec<Type>,
69        field_flags: Vec<FieldFlag>,
70        is_struct: bool,
71        generics: syn::Generics,
72    },
73    /// An anonymous variant (i.e., `MyVariant(..)`).
74    Anon {
75        enum_name: Ident,
76        name: Ident,
77        field_types: Vec<Type>,
78        field_flags: Vec<FieldFlag>,
79        is_struct: bool,
80        generics: syn::Generics,
81    },
82}
83
84impl Variant {
85    /// The number of fields in the variant.
86    fn len(&self) -> usize {
87        self.field_types().len()
88    }
89
90    /// Returns whether this variant was defined as a struct.
91    fn is_struct(&self) -> bool {
92        match self {
93            Variant::Named { is_struct, .. } => *is_struct,
94            Variant::Anon { is_struct, .. } => *is_struct,
95        }
96    }
97
98    /// The name of the enum containing the variant.
99    fn enum_name(&self) -> &Ident {
100        match self {
101            Variant::Named { enum_name, .. } => enum_name,
102            Variant::Anon { enum_name, .. } => enum_name,
103        }
104    }
105
106    /// The name of the variant itself.
107    fn name(&self) -> &Ident {
108        match self {
109            Variant::Named { name, .. } => name,
110            Variant::Anon { name, .. } => name,
111        }
112    }
113
114    /// The generics of the variant itself.
115    #[allow(dead_code)]
116    fn generics(&self) -> &syn::Generics {
117        match self {
118            Variant::Named { generics, .. } => generics,
119            Variant::Anon { generics, .. } => generics,
120        }
121    }
122
123    /// The snake_name of the variant itself.
124    fn snake_name(&self) -> Ident {
125        Ident::new(
126            &self.name().to_string().to_case(Case::Snake),
127            self.name().span(),
128        )
129    }
130
131    /// The variant's qualified name.
132    fn qualified_name(&self) -> proc_macro2::TokenStream {
133        let enum_name = self.enum_name();
134        let name = self.name();
135
136        if self.is_struct() {
137            quote! { #enum_name }
138        } else {
139            quote! { #enum_name::#name }
140        }
141    }
142
143    /// Names of the fields in the variant. Anonymous variants are named
144    /// according to their position in the argument list.
145    fn field_names(&self) -> Vec<Ident> {
146        match self {
147            Variant::Named { field_names, .. } => field_names.clone(),
148            Variant::Anon { field_types, .. } => (0usize..field_types.len())
149                .map(|idx| format_ident!("arg{}", idx))
150                .collect(),
151        }
152    }
153
154    /// The types of the fields int the variant.
155    fn field_types(&self) -> &Vec<Type> {
156        match self {
157            Variant::Named { field_types, .. } => field_types,
158            Variant::Anon { field_types, .. } => field_types,
159        }
160    }
161
162    /// Return the field flags for this variant.
163    fn field_flags(&self) -> &Vec<FieldFlag> {
164        match self {
165            Variant::Named { field_flags, .. } => field_flags,
166            Variant::Anon { field_flags, .. } => field_flags,
167        }
168    }
169
170    /// The constructor for the variant, using the field names directly.
171    fn constructor(&self) -> proc_macro2::TokenStream {
172        let qualified_name = self.qualified_name();
173        let field_names = self.field_names();
174        match self {
175            Variant::Named { .. } => quote! { #qualified_name { #(#field_names),* } },
176            Variant::Anon { .. } => quote! { #qualified_name(#(#field_names),*) },
177        }
178    }
179}
180
181struct ReplyPort {
182    is_handle: bool,
183    is_once: bool,
184}
185
186impl ReplyPort {
187    fn from_last_segment(last_segment: &proc_macro2::Ident) -> ReplyPort {
188        ReplyPort {
189            is_handle: last_segment == "PortHandle" || last_segment == "OncePortHandle",
190            is_once: last_segment == "OncePortHandle" || last_segment == "OncePortRef",
191        }
192    }
193
194    fn open_op(&self) -> proc_macro2::TokenStream {
195        if self.is_once {
196            quote! { hyperactor::mailbox::open_once_port }
197        } else {
198            quote! { hyperactor::mailbox::open_port }
199        }
200    }
201
202    fn rx_modifier(&self) -> proc_macro2::TokenStream {
203        if self.is_once {
204            quote! {}
205        } else {
206            quote! { mut }
207        }
208    }
209}
210
211/// Represents a message that can be sent to a handler, each message is associated with
212/// a variant.
213#[allow(clippy::large_enum_variant)]
214enum Message {
215    /// A call message is a request-response message, the last argument is
216    /// a [`hyperactor::OncePortRef`] or [`hyperactor::OncePortHandle`].
217    Call {
218        variant: Variant,
219        /// Tells whether the reply argument is a handle.
220        reply_port: ReplyPort,
221        /// The underlying return type (i.e., the type of the reply port).
222        return_type: Type,
223        /// the log level for generated instrumentation for handlers of this message.
224        log_level: Option<Ident>,
225    },
226    OneWay {
227        variant: Variant,
228        /// the log level for generated instrumentation for handlers of this message.
229        log_level: Option<Ident>,
230    },
231}
232
233impl Message {
234    fn new(span: Span, variant: Variant, log_level: Option<Ident>) -> Result<Self, syn::Error> {
235        match &variant
236            .field_flags()
237            .iter()
238            .zip(variant.field_types())
239            .filter_map(|(flag, ty)| match flag {
240                FieldFlag::Reply => Some(ty),
241                FieldFlag::None => None,
242            })
243            .collect::<Vec<&Type>>()[..]
244        {
245            [] => Ok(Self::OneWay { variant, log_level }),
246            [reply_port_ty] => {
247                let syn::Type::Path(type_path) = reply_port_ty else {
248                    return Err(syn::Error::new(span, REPLY_VARIANT_ERROR));
249                };
250                let Some(last_segment) = type_path.path.segments.last() else {
251                    return Err(syn::Error::new(span, REPLY_VARIANT_ERROR));
252                };
253                if last_segment.ident != "OncePortRef"
254                    && last_segment.ident != "OncePortHandle"
255                    && last_segment.ident != "PortRef"
256                    && last_segment.ident != "PortHandle"
257                {
258                    return Err(syn::Error::new_spanned(last_segment, REPLY_VARIANT_ERROR));
259                }
260                let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments else {
261                    return Err(syn::Error::new_spanned(last_segment, REPLY_VARIANT_ERROR));
262                };
263                let Some(syn::GenericArgument::Type(return_ty)) = args.args.first() else {
264                    return Err(syn::Error::new_spanned(&args.args, REPLY_VARIANT_ERROR));
265                };
266                let reply_port = ReplyPort::from_last_segment(&last_segment.ident);
267                let return_type = return_ty.clone();
268                Ok(Self::Call {
269                    variant,
270                    reply_port,
271                    return_type,
272                    log_level,
273                })
274            }
275            _ => Err(syn::Error::new(span, REPLY_USAGE_ERROR)),
276        }
277    }
278
279    /// The arguments of this message.
280    fn args(&self) -> Vec<(Ident, Type)> {
281        match self {
282            Message::Call { variant, .. } => variant
283                .field_names()
284                .into_iter()
285                .zip(variant.field_types().clone())
286                .take(variant.len() - 1)
287                .collect(),
288            Message::OneWay { variant, .. } => variant
289                .field_names()
290                .into_iter()
291                .zip(variant.field_types().clone())
292                .collect(),
293        }
294    }
295
296    fn variant(&self) -> &Variant {
297        match self {
298            Message::Call { variant, .. } => variant,
299            Message::OneWay { variant, .. } => variant,
300        }
301    }
302
303    fn reply_port_position(&self) -> Option<usize> {
304        self.variant()
305            .field_flags()
306            .iter()
307            .position(|flag| matches!(flag, FieldFlag::Reply))
308    }
309
310    /// The reply port argument of this message.
311    fn reply_port_arg(&self) -> Option<(Ident, Type)> {
312        match self {
313            Message::Call { variant, .. } => {
314                let pos = self.reply_port_position()?;
315                Some((
316                    variant.field_names()[pos].clone(),
317                    variant.field_types()[pos].clone(),
318                ))
319            }
320            Message::OneWay { .. } => None,
321        }
322    }
323}
324
325fn parse_log_level(attrs: &[Attribute]) -> Result<Option<Ident>, syn::Error> {
326    let level: Option<String> = match attrs.iter().find(|attr| attr.path().is_ident("log_level")) {
327        Some(attr) => {
328            let Ok(meta) = attr.meta.require_list() else {
329                return Err(syn::Error::new(
330                    Span::call_site(),
331                    indoc! {"
332                            `log_level` attribute must specify level. Supported levels = error, warn, info, debug, trace
333
334                            = help use `#[log_level(info)]` or `#[log_level(error)]`
335                        "},
336                ));
337            };
338            let parsed = meta.parse_args_with(Punctuated::<Ident, Token![,]>::parse_terminated)?;
339            if parsed.len() != 1 {
340                return Err(syn::Error::new(
341                    Span::call_site(),
342                    indoc! {"
343                            `log_level` attribute must specify exactly one level
344
345                            = help use `#[log_level(warn)]` or `#[log_level(info)]`
346                        "},
347                ));
348            };
349            Some(parsed.first().unwrap().to_string())
350        }
351        None => None,
352    };
353
354    if level.is_none() {
355        return Ok(None);
356    }
357    let level = level.unwrap();
358
359    match level.as_str() {
360        "error" | "warn" | "info" | "debug" | "trace" => {}
361        _ => {
362            return Err(syn::Error::new(
363                Span::call_site(),
364                indoc! {"
365                            `log_level` attribute must be one of 'error, warn, info, debug, trace'
366
367                            = help use `#[log_level(warn)]` or `#[log_level(info)]`
368                        "},
369            ));
370        }
371    }
372
373    Ok(Some(Ident::new(
374        level.to_ascii_uppercase().as_str(),
375        Span::call_site(),
376    )))
377}
378
379fn parse_field_flag(field: &Field) -> FieldFlag {
380    for attr in field.attrs.iter() {
381        match &attr.meta {
382            syn::Meta::Path(path) if path.is_ident("reply") => return FieldFlag::Reply,
383            _ => {}
384        }
385    }
386    FieldFlag::None
387}
388
389/// Parse a message enum or struct into its constituent messages.
390fn parse_messages(input: DeriveInput) -> Result<Vec<Message>, syn::Error> {
391    match &input.data {
392        Data::Enum(data_enum) => {
393            let mut messages = Vec::new();
394
395            for variant in &data_enum.variants {
396                let name = variant.ident.clone();
397                let attrs = &variant.attrs;
398
399                let message_variant = match &variant.fields {
400                    syn::Fields::Unnamed(fields_) => Variant::Anon {
401                        enum_name: input.ident.clone(),
402                        name,
403                        field_types: fields_
404                            .unnamed
405                            .iter()
406                            .map(|field| field.ty.clone())
407                            .collect(),
408                        field_flags: fields_.unnamed.iter().map(parse_field_flag).collect(),
409                        is_struct: false,
410                        generics: input.generics.clone(),
411                    },
412                    syn::Fields::Named(fields_) => Variant::Named {
413                        enum_name: input.ident.clone(),
414                        name,
415                        field_names: fields_
416                            .named
417                            .iter()
418                            .map(|field| field.ident.clone().unwrap())
419                            .collect(),
420                        field_types: fields_.named.iter().map(|field| field.ty.clone()).collect(),
421                        field_flags: fields_.named.iter().map(parse_field_flag).collect(),
422                        is_struct: false,
423                        generics: input.generics.clone(),
424                    },
425                    _ => {
426                        return Err(syn::Error::new_spanned(
427                            variant,
428                            indoc! {r#"
429                                `Handler` currently only supports named or tuple struct variants
430
431                                = help use `MyCall(Arg1Type, Arg2Type, ..)`,
432                                = help use `MyCall { arg1: Arg1Type, arg2: Arg2Type, .. }`,
433                                = help use `MyCall(Arg1Type, Arg2Type, .., #[reply] OncePortRef<ReplyType>)`
434                                = help use `MyCall { arg1: Arg1Type, arg2: Arg2Type, .., reply: #[reply] OncePortRef<ReplyType>}`
435                                = help use `MyCall(Arg1Type, Arg2Type, .., #[reply] OncePortHandle<ReplyType>)`
436                                = help use `MyCall { arg1: Arg1Type, arg2: Arg2Type, .., reply: #[reply] OncePortHandle<ReplyType>}`
437                              "#},
438                        ));
439                    }
440                };
441                let log_level = parse_log_level(attrs)?;
442
443                messages.push(Message::new(
444                    variant.fields.span(),
445                    message_variant,
446                    log_level,
447                )?);
448            }
449
450            Ok(messages)
451        }
452        Data::Struct(data_struct) => {
453            let struct_name = input.ident.clone();
454            let attrs = &input.attrs;
455
456            let message_variant = match &data_struct.fields {
457                syn::Fields::Unnamed(fields_) => Variant::Anon {
458                    enum_name: struct_name.clone(),
459                    name: struct_name,
460                    field_types: fields_
461                        .unnamed
462                        .iter()
463                        .map(|field| field.ty.clone())
464                        .collect(),
465                    field_flags: fields_.unnamed.iter().map(parse_field_flag).collect(),
466                    is_struct: true,
467                    generics: input.generics.clone(),
468                },
469                syn::Fields::Named(fields_) => Variant::Named {
470                    enum_name: struct_name.clone(),
471                    name: struct_name,
472                    field_names: fields_
473                        .named
474                        .iter()
475                        .map(|field| field.ident.clone().unwrap())
476                        .collect(),
477                    field_types: fields_.named.iter().map(|field| field.ty.clone()).collect(),
478                    field_flags: fields_.named.iter().map(parse_field_flag).collect(),
479                    is_struct: true,
480                    generics: input.generics.clone(),
481                },
482                syn::Fields::Unit => Variant::Anon {
483                    enum_name: struct_name.clone(),
484                    name: struct_name,
485                    field_types: Vec::new(),
486                    field_flags: Vec::new(),
487                    is_struct: true,
488                    generics: input.generics.clone(),
489                },
490            };
491
492            let log_level = parse_log_level(attrs)?;
493            let message = Message::new(data_struct.fields.span(), message_variant, log_level)?;
494
495            Ok(vec![message])
496        }
497        _ => Err(syn::Error::new_spanned(
498            input,
499            "handlers can only be derived for enums and structs",
500        )),
501    }
502}
503
504/// Derive a custom handler trait for given an enum containing tuple
505/// structs.  The handler trait defines a method corresponding
506/// to each of the enum's variants, and a `handle` function
507/// that dispatches messages to the correct method.  The macro
508/// supports two messaging patterns: "call" and "oneway". A call is a
509/// request-response message; a [`hyperactor::mailbox::OncePortRef`] or
510/// [`hyperactor::mailbox::OncePortHandle`] in the last position is used
511/// to send the return value.
512///
513/// The macro also derives a client trait that can be automatically implemented
514/// by specifying [`HandleClient`] for `ActorHandle<Actor>` and [`RefClient`]
515/// for `ActorRef<Actor>` accordingly. We require two implementations because
516/// not `ActorRef`s require that its message type is serializable.
517///
518/// The associated [`hyperactor_macros::handler`] macro can be used to add
519/// a dispatching handler directly to an [`hyperactor::Actor`].
520///
521/// # Example
522///
523/// The following example creates a "shopping list" actor responsible for
524/// maintaining a shopping list.
525///
526/// ```
527/// use std::collections::HashSet;
528/// use std::time::Duration;
529///
530/// use async_trait::async_trait;
531/// use hyperactor::Actor;
532/// use hyperactor::HandleClient;
533/// use hyperactor::Handler;
534/// use hyperactor::Instance;
535/// use hyperactor::RefClient;
536/// use hyperactor::reference;
537/// use serde::Deserialize;
538/// use serde::Serialize;
539/// use typeuri::Named;
540///
541/// #[derive(Handler, HandleClient, RefClient, Debug, Serialize, Deserialize, Named)]
542/// enum ShoppingList {
543///     // Oneway messages dispatch messages asynchronously, with no reply.
544///     Add(String),
545///     Remove(String),
546///
547///     // Call messages dispatch a request, expecting a reply to the
548///     // provided port, which must be in the last position.
549///     Exists(String, #[reply] reference::OncePortRef<bool>),
550///
551///     List(#[reply] reference::OncePortRef<Vec<String>>),
552/// }
553///
554/// // Define an actor.
555/// #[derive(Debug)]
556/// #[hyperactor::export(ShoppingList)]
557/// #[hyperactor::spawnable]
558/// struct ShoppingListActor(HashSet<String>);
559///
560/// #[async_trait]
561/// impl Actor for ShoppingListActor {
562///     type Params = ();
563///
564///     async fn new(_params: ()) -> Result<Self, anyhow::Error> {
565///         Ok(Self(HashSet::new()))
566///     }
567/// }
568///
569/// // ShoppingListHandler is the trait generated by derive(Handler) above.
570/// // We implement the trait here for the actor, defining a handler for
571/// // each ShoppingList message.
572/// //
573/// // The `handle` attribute installs a handler that routes messages
574/// // to the `ShoppingListHandler` implementation directly. This can also
575/// // be done manually:
576/// //
577/// // ```ignore
578/// //<ShoppingListActor as ShoppingListHandler>
579/// //     ::handle(self, comm, message).await
580/// // ```
581/// #[async_trait]
582/// #[hyperactor::handle(ShoppingList)]
583/// impl ShoppingListHandler for ShoppingListActor {
584///     async fn add(&mut self, _cx: &Context<Self>, item: String) -> Result<(), anyhow::Error> {
585///         eprintln!("insert {}", item);
586///         self.0.insert(item);
587///         Ok(())
588///     }
589///
590///     async fn remove(&mut self, _cx: &Context<Self>, item: String) -> Result<(), anyhow::Error> {
591///         eprintln!("remove {}", item);
592///         self.0.remove(&item);
593///         Ok(())
594///     }
595///
596///     async fn exists(
597///         &mut self,
598///         _cx: &Context<Self>,
599///         item: String,
600///     ) -> Result<bool, anyhow::Error> {
601///         Ok(self.0.contains(&item))
602///     }
603///
604///     async fn list(&mut self, _cx: &Context<Self>) -> Result<Vec<String>, anyhow::Error> {
605///         Ok(self.0.iter().cloned().collect())
606///     }
607/// }
608///
609/// #[tokio::main]
610/// async fn main() -> Result<(), anyhow::Error> {
611///     // Spawn our actor on the current proc.
612///     let shopping_list_actor: hyperactor::ActorHandle<ShoppingListActor> =
613///         hyperactor::spawn(ShoppingListActor::default());
614///     let client = hyperactor::client("client");
615///
616///     // todo: consider making this a macro to remove the magic names
617///
618///     // Derive(Handler) generates client methods, which call the
619///     // remote handler provided an actor instance,
620///     // the destination actor, and the method arguments.
621///
622///     shopping_list_actor.add(&client, "milk".into()).await?;
623///     shopping_list_actor.add(&client, "eggs".into()).await?;
624///
625///     println!(
626///         "got milk? {}",
627///         shopping_list_actor.exists(&client, "milk".into()).await?
628///     );
629///     println!(
630///         "got yoghurt? {}",
631///         shopping_list_actor
632///             .exists(&client, "yoghurt".into())
633///             .await?
634///     );
635///
636///     shopping_list_actor.remove(&client, "milk".into()).await?;
637///     println!(
638///         "got milk now? {}",
639///         shopping_list_actor.exists(&client, "milk".into()).await?
640///     );
641///
642///     println!(
643///         "shopping list: {:?}",
644///         shopping_list_actor.list(&client).await?
645///     );
646///
647///     let _ = proc
648///         .destroy_and_wait(Duration::from_secs(1), "example cleanup")
649///         .await?;
650///     Ok(())
651/// }
652/// ```
653#[proc_macro_derive(Handler, attributes(reply))]
654pub fn derive_handler(input: TokenStream) -> TokenStream {
655    let input = parse_macro_input!(input as DeriveInput);
656    let name: Ident = input.ident.clone();
657    let (_, ty_generics, _) = input.generics.split_for_impl();
658
659    let messages = match parse_messages(input.clone()) {
660        Ok(messages) => messages,
661        Err(err) => return TokenStream::from(err.to_compile_error()),
662    };
663
664    // Trait definition methods for the handler.
665    let mut handler_trait_methods = Vec::new();
666
667    // The arms of the match used in the message dispatcher.
668    let mut match_arms = Vec::new();
669
670    // Trait implemented by clients.
671    let mut client_trait_methods = Vec::new();
672
673    let global_log_level = parse_log_level(&input.attrs).ok().unwrap_or(None);
674
675    for message in &messages {
676        match message {
677            Message::Call {
678                variant,
679                reply_port,
680                return_type,
681                log_level,
682            } => {
683                let (arg_names, arg_types): (Vec<_>, Vec<_>) = message.args().into_iter().unzip();
684                let variant_name_snake = variant.snake_name();
685                let variant_name_snake_deprecated =
686                    format_ident!("{}_deprecated", variant_name_snake);
687                let enum_name = variant.enum_name();
688                let _variant_qualified_name = variant.qualified_name();
689                let log_level = match (&global_log_level, log_level) {
690                    (_, Some(local)) => local.clone(),
691                    (Some(global), None) => global.clone(),
692                    _ => Ident::new("DEBUG", Span::call_site()),
693                };
694                let _log_level = if reply_port.is_handle {
695                    quote! {
696                        tracing::Level::#log_level
697                    }
698                } else {
699                    quote! {
700                        tracing::Level::TRACE
701                    }
702                };
703                let log_message = quote! {
704                        hyperactor::metrics::ACTOR_MESSAGES_RECEIVED.add(1, hyperactor::kv_pairs!(
705                            "rpc" => "call",
706                            "actor_id" => hyperactor::context::Mailbox::mailbox(cx).actor_addr().to_string(),
707                            "message_type" => stringify!(#enum_name),
708                            "variant" => stringify!(#variant_name_snake),
709                        ));
710                };
711
712                handler_trait_methods.push(quote! {
713                    #[doc = "The generated handler method for this enum variant."]
714                    async fn #variant_name_snake(
715                        &mut self,
716                        cx: &hyperactor::Context<Self>,
717                        #(#arg_names: #arg_types),*)
718                        -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error>;
719                });
720
721                client_trait_methods.push(quote! {
722                    #[doc = "The generated client method for this enum variant."]
723                    async fn #variant_name_snake(
724                        &self,
725                        cx: &impl hyperactor::context::Actor,
726                        #(#arg_names: #arg_types),*)
727                        -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error>;
728
729                    #[doc = "The DEPRECATED DO NOT USE generated client method for this enum variant."]
730                    async fn #variant_name_snake_deprecated(
731                        &self,
732                        cx: &impl hyperactor::context::Actor,
733                        #(#arg_names: #arg_types),*)
734                        -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error>;
735                });
736
737                let (reply_port_arg, _) = message.reply_port_arg().unwrap();
738                let constructor = variant.constructor();
739                let result_ident = Ident::new("result", Span::mixed_site());
740                let construct_result_future = quote! { use hyperactor::Message; let #result_ident = self.#variant_name_snake(cx, #(#arg_names),*).await?; };
741                match_arms.push(quote! {
742                    #constructor => {
743                        #log_message
744                        // TODO: should we propagate this error (to supervision), or send it back as an "RPC error"?
745                        // This would require Result<Result<..., in order to handle RPC errors.
746                        #construct_result_future
747                        use hyperactor::Endpoint as _;
748                        #reply_port_arg.post(cx, #result_ident);
749                        Ok(())
750                    }
751                });
752            }
753            Message::OneWay { variant, log_level } => {
754                let (arg_names, arg_types): (Vec<_>, Vec<_>) = message.args().into_iter().unzip();
755                let variant_name_snake = variant.snake_name();
756                let variant_name_snake_deprecated =
757                    format_ident!("{}_deprecated", variant_name_snake);
758                let enum_name = variant.enum_name();
759                let log_level = match (&global_log_level, log_level) {
760                    (_, Some(local)) => local.clone(),
761                    (Some(global), None) => global.clone(),
762                    _ => Ident::new("TRACE", Span::call_site()),
763                };
764                let _log_level = quote! {
765                    tracing::Level::#log_level
766                };
767                let log_message = quote! {
768                        hyperactor::metrics::ACTOR_MESSAGES_RECEIVED.add(1, hyperactor::kv_pairs!(
769                            "rpc" => "call",
770                            "actor_id" => hyperactor::context::Mailbox::mailbox(cx).actor_addr().to_string(),
771                            "message_type" => stringify!(#enum_name),
772                            "variant" => stringify!(#variant_name_snake),
773                        ));
774                };
775
776                handler_trait_methods.push(quote! {
777                    #[doc = "The generated handler method for this enum variant."]
778                    async fn #variant_name_snake(
779                        &mut self,
780                        cx: &hyperactor::Context<Self>,
781                        #(#arg_names: #arg_types),*)
782                        -> Result<(), hyperactor::internal_macro_support::anyhow::Error>;
783                });
784
785                client_trait_methods.push(quote! {
786                    #[doc = "The generated client method for this enum variant."]
787                    async fn #variant_name_snake(
788                        &self,
789                        cx: &impl hyperactor::context::Actor,
790                        #(#arg_names: #arg_types),*)
791                        -> Result<(), hyperactor::internal_macro_support::anyhow::Error>;
792
793                    #[doc = "The DEPRECATED DO NOT USE generated client method for this enum variant."]
794                    async fn #variant_name_snake_deprecated(
795                        &self,
796                        cx: &impl hyperactor::context::Actor,
797                        #(#arg_names: #arg_types),*)
798                        -> Result<(), hyperactor::internal_macro_support::anyhow::Error>;
799                });
800
801                let constructor = variant.constructor();
802
803                match_arms.push(quote! {
804                    #constructor => {
805                        #log_message
806                        self.#variant_name_snake(cx, #(#arg_names),*).await
807                    },
808                });
809            }
810        }
811    }
812
813    let handler_trait_name = format_ident!("{}Handler", name);
814    let client_trait_name = format_ident!("{}Client", name);
815
816    // We impose additional constraints on the generics in the implementation;
817    // but the trait itself should not impose additional constraints:
818
819    let mut handler_generics = input.generics.clone();
820    for param in handler_generics.type_params_mut() {
821        param.bounds.push(syn::parse_quote!(serde::Serialize));
822        param
823            .bounds
824            .push(syn::parse_quote!(for<'de> serde::Deserialize<'de>));
825        param.bounds.push(syn::parse_quote!(Send));
826        param.bounds.push(syn::parse_quote!(Sync));
827        param.bounds.push(syn::parse_quote!(std::fmt::Debug));
828        param.bounds.push(syn::parse_quote!(typeuri::Named));
829    }
830    let (handler_impl_generics, _, _) = handler_generics.split_for_impl();
831    let (client_impl_generics, _, _) = input.generics.split_for_impl();
832
833    let expanded = quote! {
834        #[doc = "The custom handler trait for this message type."]
835        #[hyperactor::internal_macro_support::async_trait::async_trait]
836        pub trait #handler_trait_name #handler_impl_generics: hyperactor::Actor + Send + Sync  {
837            #(#handler_trait_methods)*
838
839            #[doc = "Handle the next message."]
840            async fn handle(
841                &mut self,
842                cx: &hyperactor::Context<Self>,
843                message: #name #ty_generics,
844            ) -> hyperactor::internal_macro_support::anyhow::Result<()>  {
845                 // Dispatch based on message type.
846                 match message {
847                     #(#match_arms)*
848                }
849            }
850        }
851
852        #[doc = "The custom client trait for this message type."]
853        #[hyperactor::internal_macro_support::async_trait::async_trait]
854        pub trait #client_trait_name #client_impl_generics: Send + Sync  {
855            #(#client_trait_methods)*
856        }
857    };
858
859    TokenStream::from(expanded)
860}
861
862/// Derives a client implementation on `ActorHandle<Actor>`.
863/// See [`Handler`] documentation for details.
864#[proc_macro_derive(HandleClient, attributes(log_level))]
865pub fn derive_handle_client(input: TokenStream) -> TokenStream {
866    derive_client(input, true)
867}
868
869/// Derives a client implementation on `ActorRef<Actor>`.
870/// See [`Handler`] documentation for details.
871#[proc_macro_derive(RefClient, attributes(log_level))]
872pub fn derive_ref_client(input: TokenStream) -> TokenStream {
873    derive_client(input, false)
874}
875
876fn derive_client(input: TokenStream, is_handle: bool) -> TokenStream {
877    let input = parse_macro_input!(input as DeriveInput);
878    let name = input.ident.clone();
879
880    let messages = match parse_messages(input.clone()) {
881        Ok(messages) => messages,
882        Err(err) => return TokenStream::from(err.to_compile_error()),
883    };
884
885    // The client implementation methods.
886    let mut impl_methods = Vec::new();
887
888    let send_message = quote! { hyperactor::Endpoint::post(self, cx, message); };
889    let global_log_level = parse_log_level(&input.attrs).ok().unwrap_or(None);
890
891    for message in &messages {
892        match message {
893            Message::Call {
894                variant,
895                reply_port,
896                return_type,
897                log_level,
898            } => {
899                let (arg_names, arg_types): (Vec<_>, Vec<_>) = message.args().into_iter().unzip();
900                let variant_name_snake = variant.snake_name();
901                let variant_name_snake_deprecated =
902                    format_ident!("{}_deprecated", variant_name_snake);
903                let enum_name = variant.enum_name();
904
905                let (reply_port_arg, _) = message.reply_port_arg().unwrap();
906                let constructor = variant.constructor();
907                let log_level = match (&global_log_level, log_level) {
908                    (_, Some(local)) => local.clone(),
909                    (Some(global), None) => global.clone(),
910                    _ => Ident::new("DEBUG", Span::call_site()),
911                };
912                let log_level = if is_handle {
913                    quote! {
914                        tracing::Level::#log_level
915                    }
916                } else {
917                    quote! {
918                        tracing::Level::TRACE
919                    }
920                };
921                let log_message = quote! {
922                        hyperactor::metrics::ACTOR_MESSAGES_SENT.add(1, hyperactor::kv_pairs!(
923                            "rpc" => "call",
924                            "actor_id" => hyperactor::context::Mailbox::mailbox(cx).actor_addr().to_string(),
925                            "message_type" => stringify!(#enum_name),
926                            "variant" => stringify!(#variant_name_snake),
927                        ));
928
929                };
930                let open_port = reply_port.open_op();
931                let rx_mod = reply_port.rx_modifier();
932                if reply_port.is_handle {
933                    impl_methods.push(quote! {
934                        #[hyperactor::instrument(level=#log_level, rpc = "call", message_type=#name)]
935                        async fn #variant_name_snake(
936                            &self,
937                            cx: &impl hyperactor::context::Actor,
938                            #(#arg_names: #arg_types),*)
939                            -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error> {
940                            let (#reply_port_arg, #rx_mod reply_receiver) =
941                                #open_port::<#return_type>(cx);
942                            let message = #constructor;
943                            #log_message;
944                            #send_message;
945                            reply_receiver.recv().await.map_err(hyperactor::internal_macro_support::anyhow::Error::from)
946                        }
947
948                        #[hyperactor::instrument(level=#log_level, rpc = "call", message_type=#name)]
949                        async fn #variant_name_snake_deprecated(
950                            &self,
951                            cx: &impl hyperactor::context::Actor,
952                            #(#arg_names: #arg_types),*)
953                            -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error> {
954                            let (#reply_port_arg, #rx_mod reply_receiver) =
955                                #open_port::<#return_type>(cx);
956                            let message = #constructor;
957                            #log_message;
958                            #send_message;
959                            reply_receiver.recv().await.map_err(hyperactor::internal_macro_support::anyhow::Error::from)
960                        }
961                    });
962                } else {
963                    impl_methods.push(quote! {
964                        #[hyperactor::instrument(level=#log_level, rpc="call", message_type=#name)]
965                        async fn #variant_name_snake(
966                            &self,
967                            cx: &impl hyperactor::context::Actor,
968                            #(#arg_names: #arg_types),*)
969                            -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error> {
970                            let (#reply_port_arg, #rx_mod reply_receiver) =
971                                #open_port::<#return_type>(cx);
972                            let #reply_port_arg = #reply_port_arg.bind();
973                            let message = #constructor;
974                            #log_message;
975                            #send_message;
976                            reply_receiver.recv().await.map_err(hyperactor::internal_macro_support::anyhow::Error::from)
977                        }
978
979                        #[hyperactor::instrument(level=#log_level, rpc="call", message_type=#name)]
980                        async fn #variant_name_snake_deprecated(
981                            &self,
982                            cx: &impl hyperactor::context::Actor,
983                            #(#arg_names: #arg_types),*)
984                            -> Result<#return_type, hyperactor::internal_macro_support::anyhow::Error> {
985                            let (#reply_port_arg, #rx_mod reply_receiver) =
986                                #open_port::<#return_type>(cx);
987                            let #reply_port_arg = #reply_port_arg.bind();
988                            let message = #constructor;
989                            #log_message;
990                            #send_message;
991                            reply_receiver.recv().await.map_err(hyperactor::internal_macro_support::anyhow::Error::from)
992                        }
993                    });
994                }
995            }
996            Message::OneWay { variant, log_level } => {
997                let (arg_names, arg_types): (Vec<_>, Vec<_>) = message.args().into_iter().unzip();
998                let variant_name_snake = variant.snake_name();
999                let variant_name_snake_deprecated =
1000                    format_ident!("{}_deprecated", variant_name_snake);
1001                let enum_name = variant.enum_name();
1002                let constructor = variant.constructor();
1003                let log_level = match (&global_log_level, log_level) {
1004                    (_, Some(local)) => local.clone(),
1005                    (Some(global), None) => global.clone(),
1006                    _ => Ident::new("DEBUG", Span::call_site()),
1007                };
1008                let _log_level = if is_handle {
1009                    quote! {
1010                        tracing::Level::TRACE
1011                    }
1012                } else {
1013                    quote! {
1014                        tracing::Level::#log_level
1015                    }
1016                };
1017                let log_message = quote! {
1018                    hyperactor::metrics::ACTOR_MESSAGES_SENT.add(1, hyperactor::kv_pairs!(
1019                        "rpc" => "oneway",
1020                        "actor_id" => self.actor_addr().to_string(),
1021                        "message_type" => stringify!(#enum_name),
1022                        "variant" => stringify!(#variant_name_snake),
1023                    ));
1024                };
1025                impl_methods.push(quote! {
1026                    async fn #variant_name_snake(
1027                        &self,
1028                        cx: &impl hyperactor::context::Actor,
1029                        #(#arg_names: #arg_types),*)
1030                        -> Result<(), hyperactor::internal_macro_support::anyhow::Error> {
1031                        let message = #constructor;
1032                        #log_message;
1033                        #send_message;
1034                        Ok(())
1035                    }
1036
1037                    async fn #variant_name_snake_deprecated(
1038                        &self,
1039                        cx: &impl hyperactor::context::Actor,
1040                        #(#arg_names: #arg_types),*)
1041                        -> Result<(), hyperactor::internal_macro_support::anyhow::Error> {
1042                        let message = #constructor;
1043                        #log_message;
1044                        #send_message;
1045                        Ok(())
1046                    }
1047                });
1048            }
1049        }
1050    }
1051
1052    let trait_name = format_ident!("{}Client", name);
1053
1054    let (_, ty_generics, _) = input.generics.split_for_impl();
1055
1056    // Add a new generic parameter 'A'
1057    let actor_ident = Ident::new("A", proc_macro2::Span::from(proc_macro::Span::def_site()));
1058    let mut trait_generics = input.generics.clone();
1059    trait_generics.params.insert(
1060        0,
1061        syn::GenericParam::Type(syn::TypeParam {
1062            ident: actor_ident.clone(),
1063            attrs: vec![],
1064            colon_token: None,
1065            bounds: Punctuated::new(),
1066            eq_token: None,
1067            default: None,
1068        }),
1069    );
1070
1071    for param in trait_generics.type_params_mut() {
1072        if param.ident == actor_ident {
1073            continue;
1074        }
1075        param.bounds.push(syn::parse_quote!(serde::Serialize));
1076        param
1077            .bounds
1078            .push(syn::parse_quote!(for<'de> serde::Deserialize<'de>));
1079        param.bounds.push(syn::parse_quote!(Send));
1080        param.bounds.push(syn::parse_quote!(Sync));
1081        param.bounds.push(syn::parse_quote!(std::fmt::Debug));
1082        param.bounds.push(syn::parse_quote!(typeuri::Named));
1083    }
1084
1085    let (impl_generics, _, _) = trait_generics.split_for_impl();
1086
1087    let expanded = if is_handle {
1088        quote! {
1089            #[hyperactor::internal_macro_support::async_trait::async_trait]
1090            impl #impl_generics #trait_name #ty_generics for hyperactor::ActorHandle<#actor_ident>
1091              where #actor_ident: hyperactor::Handler<#name #ty_generics> {
1092                #(#impl_methods)*
1093            }
1094        }
1095    } else {
1096        quote! {
1097            #[hyperactor::internal_macro_support::async_trait::async_trait]
1098            impl #impl_generics #trait_name #ty_generics for hyperactor::ActorRef<#actor_ident>
1099              where #actor_ident: hyperactor::actor::RemoteHandles<#name #ty_generics> {
1100                #(#impl_methods)*
1101            }
1102        }
1103    };
1104
1105    TokenStream::from(expanded)
1106}
1107
1108const HANDLE_ARGUMENT_ERROR: &str = indoc! {r#"
1109`handle` expects the message type that is being handled
1110
1111= help: use `#[handle(MessageType)]`
1112"#};
1113
1114/// Install a [`Handler`] that routes messages of the provided type to this handler trait implementation.
1115#[proc_macro_attribute]
1116pub fn handle(attr: TokenStream, item: TokenStream) -> TokenStream {
1117    let attr_args = parse_macro_input!(attr with Punctuated::<syn::PathSegment, syn::Token![,]>::parse_terminated);
1118    if attr_args.len() != 1 {
1119        return TokenStream::from(
1120            syn::Error::new_spanned(attr_args, HANDLE_ARGUMENT_ERROR).to_compile_error(),
1121        );
1122    }
1123
1124    let message_type = attr_args.first().unwrap();
1125    let input = parse_macro_input!(item as ItemImpl);
1126
1127    let self_type = match *input.self_ty {
1128        syn::Type::Path(ref type_path) => {
1129            let segment = type_path.path.segments.last().unwrap();
1130            segment.clone() //ident.clone()
1131        }
1132        _ => {
1133            return TokenStream::from(
1134                syn::Error::new_spanned(input.self_ty, "`handle` argument must be a type")
1135                    .to_compile_error(),
1136            );
1137        }
1138    };
1139
1140    let trait_name = match input.trait_ {
1141        Some((_, ref trait_path, _)) => trait_path.segments.last().unwrap().clone(),
1142        None => {
1143            return TokenStream::from(
1144                syn::Error::new_spanned(input.self_ty, "no trait in implementation block")
1145                    .to_compile_error(),
1146            );
1147        }
1148    };
1149
1150    let expanded = quote! {
1151        #input
1152
1153        #[hyperactor::internal_macro_support::async_trait::async_trait]
1154        impl hyperactor::Handler<#message_type> for #self_type {
1155            async fn handle(
1156                &mut self,
1157                cx: &hyperactor::Context<Self>,
1158                message: #message_type,
1159            ) -> hyperactor::internal_macro_support::anyhow::Result<()> {
1160                <Self as #trait_name>::handle(self, cx, message).await
1161            }
1162        }
1163    };
1164
1165    TokenStream::from(expanded)
1166}
1167
1168/// Use this macro in place of tracing::instrument to prevent spamming our tracing table.
1169/// We set a default level of INFO while always setting ERROR if the function returns Result::Err giving us
1170/// consistent and high quality structured logs. Because this wraps around tracing::instrument, all parameters
1171/// mentioned in https://fburl.com/9jlkb5q4 should be valid. For functions that don't return a [`Result`] type, use
1172/// [`instrument_infallible`]
1173///
1174/// ```
1175/// #[telemetry::instrument]
1176/// async fn yolo() -> anyhow::Result<i32> {
1177///     Ok(420)
1178/// }
1179/// ```
1180#[proc_macro_attribute]
1181pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream {
1182    let args =
1183        parse_macro_input!(args with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated);
1184    let input = parse_macro_input!(input as ItemFn);
1185    let output = quote! {
1186        #[hyperactor::internal_macro_support::tracing::instrument(err, skip_all, #args)]
1187        #input
1188    };
1189
1190    TokenStream::from(output)
1191}
1192
1193/// Use this macro in place of tracing::instrument to prevent spamming our tracing table.
1194/// Because this wraps around tracing::instrument, all parameters mentioned in
1195/// https://fburl.com/9jlkb5q4 should be valid.
1196///
1197/// ```
1198/// #[telemetry::instrument]
1199/// async fn yolo() -> i32 {
1200///     420
1201/// }
1202/// ```
1203#[proc_macro_attribute]
1204pub fn instrument_infallible(args: TokenStream, input: TokenStream) -> TokenStream {
1205    let args =
1206        parse_macro_input!(args with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated);
1207    let input = parse_macro_input!(input as ItemFn);
1208
1209    let output = quote! {
1210        #[hyperactor::internal_macro_support::tracing::instrument(skip_all, #args)]
1211        #input
1212    };
1213
1214    TokenStream::from(output)
1215}
1216
1217struct HandlerSpec {
1218    ty: Type,
1219}
1220
1221impl Parse for HandlerSpec {
1222    fn parse(input: ParseStream) -> syn::Result<Self> {
1223        let ty: Type = input.parse()?;
1224
1225        if input.is_empty() || input.peek(Token![,]) {
1226            Ok(HandlerSpec { ty })
1227        } else {
1228            // Something unexpected follows the type
1229            let unexpected: proc_macro2::TokenTree = input.parse()?;
1230            Err(syn::Error::new_spanned(
1231                unexpected,
1232                "unexpected token after type — use the bare message type",
1233            ))
1234        }
1235    }
1236}
1237
1238impl HandlerSpec {
1239    fn message_types(handlers: Vec<HandlerSpec>) -> Vec<Type> {
1240        handlers.into_iter().map(|handler| handler.ty).collect()
1241    }
1242}
1243
1244fn named_impl(data_type_name: &Ident, generics: &syn::Generics) -> proc_macro2::TokenStream {
1245    let generics_with_bounds = generics_with_named_bounds(generics);
1246    let type_params: Vec<_> = generics.type_params().collect();
1247    let has_generics = !type_params.is_empty();
1248
1249    let (impl_generics_with_bounds, _, _) = generics_with_bounds.split_for_impl();
1250    let (_, ty_generics, where_clause) = generics.split_for_impl();
1251
1252    let (typename_impl, typehash_impl) = if has_generics {
1253        let placeholders = vec!["{}"; type_params.len()].join(", ");
1254        let placeholders_format_string = format!("<{}>", placeholders);
1255        let format_string = quote! {
1256            concat!(
1257                std::module_path!(),
1258                "::",
1259                stringify!(#data_type_name),
1260                #placeholders_format_string
1261            )
1262        };
1263        let type_param_idents: Vec<_> = type_params.iter().map(|param| &param.ident).collect();
1264        (
1265            quote! {
1266                typeuri::intern_typename!(Self, #format_string, #(#type_param_idents),*)
1267            },
1268            quote! {
1269                typeuri::cityhasher::hash(Self::typename())
1270            },
1271        )
1272    } else {
1273        (
1274            quote! {
1275                concat!(std::module_path!(), "::", stringify!(#data_type_name))
1276            },
1277            quote! {
1278                static TYPEHASH: std::sync::LazyLock<u64> = std::sync::LazyLock::new(|| {
1279                    typeuri::cityhasher::hash(<#data_type_name as typeuri::Named>::typename())
1280                });
1281                *TYPEHASH
1282            },
1283        )
1284    };
1285
1286    quote! {
1287        impl #impl_generics_with_bounds typeuri::Named for #data_type_name #ty_generics #where_clause {
1288            fn typename() -> &'static str {
1289                #typename_impl
1290            }
1291
1292            fn typehash() -> u64 {
1293                #typehash_impl
1294            }
1295        }
1296    }
1297}
1298
1299fn generics_with_named_bounds(generics: &syn::Generics) -> syn::Generics {
1300    let mut generics = generics.clone();
1301    for param in generics.type_params_mut() {
1302        param.bounds.push(syn::parse_quote!(typeuri::Named));
1303    }
1304    generics
1305}
1306
1307fn generics_with_predicates(
1308    generics: &syn::Generics,
1309    predicates: impl IntoIterator<Item = WherePredicate>,
1310) -> syn::Generics {
1311    let mut generics = generics.clone();
1312    generics.make_where_clause().predicates.extend(predicates);
1313    generics
1314}
1315
1316/// Attribute Struct for [`fn export`] macro.
1317struct ExportAttr {
1318    handlers: Vec<HandlerSpec>,
1319}
1320
1321impl Parse for ExportAttr {
1322    fn parse(input: ParseStream) -> syn::Result<Self> {
1323        if input.is_empty() {
1324            return Ok(Self {
1325                handlers: Vec::new(),
1326            });
1327        }
1328
1329        let compatibility_form = {
1330            let fork = input.fork();
1331            fork.parse::<Ident>().is_ok() && fork.parse::<Token![=]>().is_ok()
1332        };
1333
1334        if !compatibility_form {
1335            let handlers = input
1336                .parse_terminated(HandlerSpec::parse, Token![,])?
1337                .into_iter()
1338                .collect();
1339            return Ok(Self { handlers });
1340        }
1341
1342        let mut handlers: Vec<HandlerSpec> = vec![];
1343
1344        while !input.is_empty() {
1345            let key: Ident = input.parse()?;
1346            input.parse::<Token![=]>()?;
1347
1348            if key == "spawn" {
1349                let expr: Expr = input.parse()?;
1350                return Err(syn::Error::new_spanned(
1351                    expr,
1352                    "`spawn = true` is no longer supported; use `#[spawnable]` on concrete actor declarations or `hyperactor::register_spawnable!(ConcreteType)` for generic instantiations",
1353                ));
1354            } else if key == "handlers" {
1355                let content;
1356                bracketed!(content in input);
1357                let raw_handlers = content.parse_terminated(HandlerSpec::parse, Token![,])?;
1358                handlers = raw_handlers.into_iter().collect();
1359            } else {
1360                return Err(syn::Error::new_spanned(
1361                    key,
1362                    "unexpected key in `#[export(...)]`. Use direct handler lists, or the compatibility key `handlers`",
1363                ));
1364            }
1365
1366            // optional trailing comma
1367            let _ = input.parse::<Token![,]>();
1368        }
1369
1370        Ok(ExportAttr { handlers })
1371    }
1372}
1373
1374/// Exports handlers for this actor. The set of exported handlers
1375/// determine the messages that may be sent to remote references of
1376/// the actor ([`hyperaxtor::ActorRef`]). Only messages that implement
1377/// [`hyperactor::RemoteMessage`] may be exported.
1378///
1379/// # Example
1380///
1381/// In the following example, `MyActor` exports handlers for two message types,
1382/// `MyMessage` and `MyOtherMessage`. Consequently, `ActorRef`s of the actor's
1383/// type may dispatch messages of these types.
1384///
1385/// ```ignore
1386/// #[export(MyMessage, MyOtherMessage)]
1387/// struct MyActor {}
1388/// ```
1389#[proc_macro_attribute]
1390pub fn export(attr: TokenStream, item: TokenStream) -> TokenStream {
1391    let input: DeriveInput = parse_macro_input!(item as DeriveInput);
1392    let data_type_name = &input.ident;
1393    let (_, ty_generics, _) = input.generics.split_for_impl();
1394    let named_generics = generics_with_named_bounds(&input.generics);
1395    let (named_impl_generics, named_ty_generics, named_where_clause) =
1396        named_generics.split_for_impl();
1397
1398    let ExportAttr { handlers } = parse_macro_input!(attr as ExportAttr);
1399
1400    let mut handles = Vec::new();
1401    let mut bindings = Vec::new();
1402    let mut bind_predicates = Vec::new();
1403    let actor_ty: Type = syn::parse_quote!(#data_type_name #ty_generics);
1404
1405    for HandlerSpec { ty } in &handlers {
1406        let message_generics = generics_with_predicates(
1407            &named_generics,
1408            [syn::parse_quote!(#ty: hyperactor::RemoteMessage)],
1409        );
1410        let (message_impl_generics, message_ty_generics, message_where_clause) =
1411            message_generics.split_for_impl();
1412        handles.push(quote! {
1413            impl #message_impl_generics hyperactor::actor::RemoteHandles<#ty>
1414                for #data_type_name #message_ty_generics #message_where_clause {}
1415            impl #message_impl_generics hyperactor::remote::Accepts<#ty>
1416                for #data_type_name #message_ty_generics #message_where_clause {}
1417        });
1418        bindings.push(quote! {
1419            ports.bind::<#ty>();
1420        });
1421        bind_predicates.push(syn::parse_quote!(#ty: hyperactor::RemoteMessage));
1422        bind_predicates.push(syn::parse_quote!(#actor_ty: hyperactor::Handler<#ty>));
1423    }
1424
1425    let bind_generics = generics_with_predicates(&named_generics, bind_predicates);
1426    let (bind_impl_generics, bind_ty_generics, bind_where_clause) = bind_generics.split_for_impl();
1427    let named_impl = named_impl(data_type_name, &input.generics);
1428
1429    let expanded = quote! {
1430        #input
1431
1432        impl #named_impl_generics hyperactor::actor::Referable for #data_type_name #named_ty_generics #named_where_clause {}
1433
1434        #(#handles)*
1435
1436        // Always export the `IntrospectMessage` type.
1437        impl #named_impl_generics hyperactor::actor::RemoteHandles<hyperactor::introspect::IntrospectMessage> for #data_type_name #named_ty_generics #named_where_clause {}
1438        impl #named_impl_generics hyperactor::remote::Accepts<hyperactor::introspect::IntrospectMessage> for #data_type_name #named_ty_generics #named_where_clause {}
1439
1440        impl #bind_impl_generics hyperactor::actor::Binds<#data_type_name #bind_ty_generics> for #data_type_name #bind_ty_generics #bind_where_clause {
1441            fn bind(ports: &hyperactor::proc::HandlerPorts<Self>) {
1442                #(#bindings)*
1443            }
1444        }
1445
1446        #named_impl
1447    };
1448
1449    TokenStream::from(expanded)
1450}
1451
1452/// Marks a concrete actor declaration as remotely spawnable.
1453///
1454/// When combined with `#[export(...)]`, place `#[spawnable]` below `#[export]`.
1455#[proc_macro_attribute]
1456pub fn spawnable(attr: TokenStream, item: TokenStream) -> TokenStream {
1457    if !attr.is_empty() {
1458        return syn::Error::new(Span::call_site(), "`#[spawnable]` does not take arguments")
1459            .to_compile_error()
1460            .into();
1461    }
1462
1463    let input: DeriveInput = parse_macro_input!(item as DeriveInput);
1464    if !matches!(input.data, Data::Struct(_)) {
1465        return syn::Error::new(
1466            input.span(),
1467            "`#[spawnable]` only supports struct actor declarations",
1468        )
1469        .to_compile_error()
1470        .into();
1471    }
1472
1473    if !input.generics.params.is_empty() {
1474        return syn::Error::new(
1475            input.generics.span(),
1476            "generic actor families cannot use `#[spawnable]`; use `hyperactor::register_spawnable!(ConcreteType)` instead",
1477        )
1478        .to_compile_error()
1479        .into();
1480    }
1481
1482    let data_type_name = &input.ident;
1483    quote! {
1484        #input
1485        hyperactor::register_spawnable!(#data_type_name);
1486    }
1487    .into()
1488}
1489
1490/// Represents the full input to [`fn behavior`].
1491struct BehaviorInput {
1492    behavior: Ident,
1493    generics: syn::Generics,
1494    handlers: Vec<HandlerSpec>,
1495}
1496
1497impl syn::parse::Parse for BehaviorInput {
1498    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
1499        let behavior: Ident = input.parse()?;
1500        let generics: syn::Generics = input.parse()?;
1501        let _: Token![,] = input.parse()?;
1502        let raw_handlers = input.parse_terminated(HandlerSpec::parse, Token![,])?;
1503        let handlers = raw_handlers.into_iter().collect();
1504        Ok(BehaviorInput {
1505            behavior,
1506            generics,
1507            handlers,
1508        })
1509    }
1510}
1511
1512/// Create a [`Referable`] definition, handling a specific set of message types.
1513/// Behaviors are used to create an [`ActorRef`] without having to depend on the
1514/// actor's implementation. Casts are supported for any remote message handled
1515/// by the actor.
1516///
1517/// ```
1518/// hyperactor::behavior!(TestActorBehavior, TestMessage, (), MyGeneric<()>, u64,);
1519/// ```
1520///
1521/// This macro also supports generic behaviors:
1522/// ```
1523/// hyperactor::behavior!(TestBehavior<T>, Message<T>, u64,);
1524/// ```
1525#[proc_macro]
1526pub fn behavior(input: TokenStream) -> TokenStream {
1527    let BehaviorInput {
1528        behavior,
1529        generics,
1530        handlers,
1531    } = parse_macro_input!(input as BehaviorInput);
1532    let tys = HandlerSpec::message_types(handlers);
1533
1534    // Add bounds to generics for Named, Serialize, Deserialize
1535    let mut bounded_generics = generics.clone();
1536    for param in bounded_generics.type_params_mut() {
1537        param.bounds.push(syn::parse_quote!(typeuri::Named));
1538        param.bounds.push(syn::parse_quote!(serde::Serialize));
1539        param.bounds.push(syn::parse_quote!(std::marker::Send));
1540        param.bounds.push(syn::parse_quote!(std::marker::Sync));
1541        param.bounds.push(syn::parse_quote!(std::fmt::Debug));
1542        // Note: lifetime parameters are not *actually* hygienic.
1543        // https://github.com/rust-lang/rust/issues/54727
1544        let lifetime =
1545            syn::Lifetime::new("'hyperactor_behavior_de", proc_macro2::Span::mixed_site());
1546        param
1547            .bounds
1548            .push(syn::parse_quote!(for<#lifetime> serde::Deserialize<#lifetime>));
1549    }
1550
1551    // Split the generics for use in different contexts
1552    let (impl_generics, ty_generics, where_clause) = bounded_generics.split_for_impl();
1553
1554    // Create a combined generics for the Binds impl that includes both A and the behavior's generics
1555    let mut binds_generics = bounded_generics.clone();
1556    binds_generics.params.insert(
1557        0,
1558        syn::GenericParam::Type(syn::TypeParam {
1559            attrs: vec![],
1560            ident: Ident::new("A", proc_macro2::Span::call_site()),
1561            colon_token: None,
1562            bounds: Punctuated::new(),
1563            eq_token: None,
1564            default: None,
1565        }),
1566    );
1567    let (binds_impl_generics, _, _) = binds_generics.split_for_impl();
1568
1569    // Determine typename and typehash implementation based on whether we have generics
1570    let type_params: Vec<_> = bounded_generics.type_params().collect();
1571    let has_generics = !type_params.is_empty();
1572
1573    let (typename_impl, typehash_impl) = if has_generics {
1574        // Create format string with placeholders for each generic parameter
1575        let placeholders = vec!["{}"; type_params.len()].join(", ");
1576        let placeholders_format_string = format!("<{}>", placeholders);
1577        let format_string = quote! { concat!(std::module_path!(), "::", stringify!(#behavior), #placeholders_format_string) };
1578
1579        let type_param_idents: Vec<_> = type_params.iter().map(|p| &p.ident).collect();
1580        (
1581            quote! {
1582                typeuri::intern_typename!(Self, #format_string, #(#type_param_idents),*)
1583            },
1584            quote! {
1585                typeuri::cityhasher::hash(Self::typename())
1586            },
1587        )
1588    } else {
1589        (
1590            quote! {
1591                concat!(std::module_path!(), "::", stringify!(#behavior))
1592            },
1593            quote! {
1594                static TYPEHASH: std::sync::LazyLock<u64> = std::sync::LazyLock::new(|| {
1595                    typeuri::cityhasher::hash(<#behavior as typeuri::Named>::typename())
1596                });
1597                *TYPEHASH
1598            },
1599        )
1600    };
1601
1602    let type_param_idents = generics.type_params().map(|p| &p.ident).collect::<Vec<_>>();
1603
1604    let expanded = quote! {
1605        #[doc = "The generated behavior struct."]
1606        #[derive(Debug, serde::Serialize, serde::Deserialize)]
1607        pub struct #behavior #impl_generics #where_clause {
1608            _phantom: std::marker::PhantomData<(#(#type_param_idents),*)>
1609        }
1610
1611        impl #impl_generics typeuri::Named for #behavior #ty_generics #where_clause {
1612            fn typename() -> &'static str {
1613                #typename_impl
1614            }
1615
1616            fn typehash() -> u64 {
1617                #typehash_impl
1618            }
1619        }
1620
1621        impl #impl_generics hyperactor::actor::Referable for #behavior #ty_generics #where_clause {}
1622
1623        impl #binds_impl_generics hyperactor::actor::Binds<A> for #behavior #ty_generics
1624        where
1625            A: hyperactor::Actor #(+ hyperactor::Handler<#tys>)*,
1626            #where_clause
1627        {
1628            fn bind(ports: &hyperactor::proc::HandlerPorts<A>) {
1629                #(
1630                    ports.bind::<#tys>();
1631                )*
1632            }
1633        }
1634
1635        #(
1636            impl #impl_generics hyperactor::actor::RemoteHandles<#tys> for #behavior #ty_generics #where_clause {}
1637            impl #impl_generics hyperactor::remote::Accepts<#tys> for #behavior #ty_generics #where_clause {}
1638        )*
1639    };
1640
1641    TokenStream::from(expanded)
1642}
1643
1644// Helper function for common parsing and validation
1645fn parse_observe_function(
1646    attr: TokenStream,
1647    item: TokenStream,
1648) -> syn::Result<(ItemFn, String, String)> {
1649    let input = syn::parse::<ItemFn>(item)?;
1650
1651    if input.sig.asyncness.is_none() {
1652        return Err(syn::Error::new(
1653            input.sig.span(),
1654            "observe macros can only be applied to async functions",
1655        ));
1656    }
1657
1658    let fn_name_str = input.sig.ident.to_string();
1659    let module_name_str = syn::parse::<syn::LitStr>(attr)?.value();
1660
1661    Ok((input, fn_name_str, module_name_str))
1662}
1663
1664// Helper function for creating telemetry identifiers and setup code
1665fn create_telemetry_setup(
1666    module_name_str: &str,
1667    fn_name_str: &str,
1668    include_error: bool,
1669) -> (Ident, Ident, Option<Ident>, proc_macro2::TokenStream) {
1670    let module_and_fn = format!("{}_{}", module_name_str, fn_name_str);
1671    let latency_ident = Ident::new("latency", Span::from(proc_macro::Span::def_site()));
1672
1673    let success_ident = Ident::new("success", Span::from(proc_macro::Span::def_site()));
1674
1675    let error_ident = if include_error {
1676        Some(Ident::new(
1677            "error",
1678            Span::from(proc_macro::Span::def_site()),
1679        ))
1680    } else {
1681        None
1682    };
1683
1684    let error_declaration = if let Some(ref error_ident) = error_ident {
1685        quote! {
1686            hyperactor_telemetry::declare_static_counter!(#error_ident, concat!(#module_and_fn, ".error"));
1687        }
1688    } else {
1689        quote! {}
1690    };
1691
1692    let setup_code = quote! {
1693        use hyperactor_telemetry;
1694        hyperactor_telemetry::declare_static_timer!(#latency_ident, concat!(#module_and_fn, ".latency"), hyperactor_telemetry::TimeUnit::Micros);
1695        hyperactor_telemetry::declare_static_counter!(#success_ident, concat!(#module_and_fn, ".success"));
1696        #error_declaration
1697    };
1698
1699    (latency_ident, success_ident, error_ident, setup_code)
1700}
1701
1702/// A procedural macro that automatically injects telemetry code into async functions
1703/// that return a Result type.
1704///
1705/// This macro wraps async functions and adds instrumentation to measure:
1706/// 1. Latency - how long the function takes to execute
1707/// 2. Error counter - function error count
1708/// 3. Success counter - function completion count
1709///
1710/// # Example
1711///
1712/// ```rust
1713/// use hyperactor_actor::observe_result;
1714///
1715/// #[observe_result("my_module")]
1716/// async fn process_request(user_id: &str) -> Result<String, Error> {
1717///     // Function implementation
1718///     // Telemetry will be automatically collected
1719/// }
1720/// ```
1721#[proc_macro_attribute]
1722pub fn observe_result(attr: TokenStream, item: TokenStream) -> TokenStream {
1723    let (input, fn_name_str, module_name_str) = match parse_observe_function(attr, item) {
1724        Ok(parsed) => parsed,
1725        Err(err) => return err.to_compile_error().into(),
1726    };
1727
1728    let fn_name = &input.sig.ident;
1729    let vis = &input.vis;
1730    let args = &input.sig.inputs;
1731    let return_type = &input.sig.output;
1732    let body = &input.block;
1733    let attrs = &input.attrs;
1734    let generics = &input.sig.generics;
1735
1736    let (latency_ident, success_ident, error_ident, telemetry_setup) =
1737        create_telemetry_setup(&module_name_str, &fn_name_str, true);
1738    let error_ident = error_ident.unwrap();
1739
1740    let result_ident = Ident::new("result", Span::from(proc_macro::Span::def_site()));
1741
1742    // Generate the instrumented function
1743    let expanded = quote! {
1744        #(#attrs)*
1745        #vis async fn #fn_name #generics(#args) #return_type {
1746            #telemetry_setup
1747
1748            let kv_pairs = hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone());
1749            let _timer = #latency_ident.start(kv_pairs);
1750
1751            let #result_ident = async #body.await;
1752
1753            match &#result_ident {
1754                Ok(_) => {
1755                    #success_ident.add(
1756                        1,
1757                        hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1758                    );
1759                }
1760                Err(_) => {
1761                    #error_ident.add(
1762                        1,
1763                        hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1764                    );
1765                }
1766            }
1767
1768            #result_ident
1769        }
1770    };
1771
1772    expanded.into()
1773}
1774
1775/// A procedural macro that automatically injects telemetry code into async functions
1776/// that do not return a Result type.
1777///
1778/// This macro wraps async functions and adds instrumentation to measure:
1779/// 1. Latency - how long the function takes to execute
1780/// 2. Success counter - function completion count
1781///
1782/// # Example
1783///
1784/// ```rust
1785/// use hyperactor_actor::observe_async;
1786///
1787/// #[observe_async("my_module")]
1788/// async fn process_data(data: &str) -> String {
1789///     // Function implementation
1790///     // Telemetry will be automatically collected
1791/// }
1792/// ```
1793#[proc_macro_attribute]
1794pub fn observe_async(attr: TokenStream, item: TokenStream) -> TokenStream {
1795    let (input, fn_name_str, module_name_str) = match parse_observe_function(attr, item) {
1796        Ok(parsed) => parsed,
1797        Err(err) => return err.to_compile_error().into(),
1798    };
1799
1800    let fn_name = &input.sig.ident;
1801    let vis = &input.vis;
1802    let args = &input.sig.inputs;
1803    let return_type = &input.sig.output;
1804    let body = &input.block;
1805    let attrs = &input.attrs;
1806    let generics = &input.sig.generics;
1807
1808    let (latency_ident, success_ident, _, telemetry_setup) =
1809        create_telemetry_setup(&module_name_str, &fn_name_str, false);
1810
1811    let return_ident = Ident::new("ret", Span::from(proc_macro::Span::def_site()));
1812
1813    // Generate the instrumented function
1814    let expanded = quote! {
1815        #(#attrs)*
1816        #vis async fn #fn_name #generics(#args) #return_type {
1817            #telemetry_setup
1818
1819            let kv_pairs = hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone());
1820            let _timer = #latency_ident.start(kv_pairs);
1821
1822            let #return_ident = async #body.await;
1823
1824            #success_ident.add(
1825                1,
1826                hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1827            );
1828            #return_ident
1829        }
1830    };
1831
1832    expanded.into()
1833}
1834
1835fn validate_label(s: &str) -> Result<(), String> {
1836    if s.is_empty() {
1837        return Err("label must not be empty".to_string());
1838    }
1839    if s.len() > 63 {
1840        return Err("label exceeds 63 characters".to_string());
1841    }
1842    let first = s.as_bytes()[0];
1843    if !first.is_ascii_lowercase() {
1844        return Err("label must start with a lowercase letter".to_string());
1845    }
1846    let last = s.as_bytes()[s.len() - 1];
1847    if !last.is_ascii_lowercase() && !last.is_ascii_digit() {
1848        return Err("label must end with a lowercase letter or digit".to_string());
1849    }
1850    for ch in s.chars() {
1851        if !ch.is_ascii_lowercase() && !ch.is_ascii_digit() && ch != '-' {
1852            return Err(format!("label contains invalid character '{ch}'"));
1853        }
1854    }
1855    Ok(())
1856}
1857
1858fn validate_hex_uid(s: &str) -> Result<u64, String> {
1859    if s.is_empty() || s.len() > 16 {
1860        return Err(format!("hex uid must be 1-16 hex characters, got '{s}'"));
1861    }
1862    for ch in s.chars() {
1863        if !ch.is_ascii_hexdigit() {
1864            return Err(format!("hex uid contains invalid character '{ch}'"));
1865        }
1866    }
1867    u64::from_str_radix(s, 16).map_err(|e| format!("invalid hex uid '{s}': {e}"))
1868}
1869
1870/// Compile-time validated [`hyperactor::id::Uid`] construction.
1871///
1872/// Accepts two forms:
1873/// - `uid!(_my-singleton)` — a singleton Uid
1874/// - `uid!(d5d54d7201103869)` — an instance Uid
1875#[proc_macro]
1876pub fn uid(input: TokenStream) -> TokenStream {
1877    let input2: proc_macro2::TokenStream = input.into();
1878    let combined: String = input2.into_iter().map(|tt| tt.to_string()).collect();
1879
1880    if combined.is_empty() {
1881        return TokenStream::from(quote! { compile_error!("uid! macro requires an argument") });
1882    }
1883
1884    // Singleton: starts with '_'
1885    if let Some(rest) = combined.strip_prefix('_') {
1886        return match validate_label(rest) {
1887            Ok(()) => TokenStream::from(quote! {
1888                hyperactor::id::Uid::Singleton(
1889                    hyperactor::id::Label::new(#rest).unwrap()
1890                )
1891            }),
1892            Err(e) => {
1893                let msg = format!("invalid singleton uid: {e}");
1894                TokenStream::from(quote! { compile_error!(#msg) })
1895            }
1896        };
1897    }
1898
1899    // Instance: bare hex
1900    match validate_hex_uid(&combined) {
1901        Ok(uid_val) => TokenStream::from(quote! {
1902            hyperactor::id::Uid::Instance(#uid_val, None)
1903        }),
1904        Err(e) => {
1905            let msg = format!("invalid uid: {e}");
1906            TokenStream::from(quote! { compile_error!(#msg) })
1907        }
1908    }
1909}
hyperactor_macros/lib.rs

hyperactor_macros/
lib.rs
