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            default: None,
1067        }),
1068    );
1069
1070    for param in trait_generics.type_params_mut() {
1071        if param.ident == actor_ident {
1072            continue;
1073        }
1074        param.bounds.push(syn::parse_quote!(serde::Serialize));
1075        param
1076            .bounds
1077            .push(syn::parse_quote!(for<'de> serde::Deserialize<'de>));
1078        param.bounds.push(syn::parse_quote!(Send));
1079        param.bounds.push(syn::parse_quote!(Sync));
1080        param.bounds.push(syn::parse_quote!(std::fmt::Debug));
1081        param.bounds.push(syn::parse_quote!(typeuri::Named));
1082    }
1083
1084    let (impl_generics, _, _) = trait_generics.split_for_impl();
1085
1086    let expanded = if is_handle {
1087        quote! {
1088            #[hyperactor::internal_macro_support::async_trait::async_trait]
1089            impl #impl_generics #trait_name #ty_generics for hyperactor::ActorHandle<#actor_ident>
1090              where #actor_ident: hyperactor::Handler<#name #ty_generics> {
1091                #(#impl_methods)*
1092            }
1093        }
1094    } else {
1095        quote! {
1096            #[hyperactor::internal_macro_support::async_trait::async_trait]
1097            impl #impl_generics #trait_name #ty_generics for hyperactor::ActorRef<#actor_ident>
1098              where #actor_ident: hyperactor::actor::RemoteHandles<#name #ty_generics> {
1099                #(#impl_methods)*
1100            }
1101        }
1102    };
1103
1104    TokenStream::from(expanded)
1105}
1106
1107const HANDLE_ARGUMENT_ERROR: &str = indoc! {r#"
1108`handle` expects the message type that is being handled
1109
1110= help: use `#[handle(MessageType)]`
1111"#};
1112
1113/// Install a [`Handler`] that routes messages of the provided type to this handler trait implementation.
1114#[proc_macro_attribute]
1115pub fn handle(attr: TokenStream, item: TokenStream) -> TokenStream {
1116    let attr_args = parse_macro_input!(attr with Punctuated::<syn::PathSegment, syn::Token![,]>::parse_terminated);
1117    if attr_args.len() != 1 {
1118        return TokenStream::from(
1119            syn::Error::new_spanned(attr_args, HANDLE_ARGUMENT_ERROR).to_compile_error(),
1120        );
1121    }
1122
1123    let message_type = attr_args.first().unwrap();
1124    let input = parse_macro_input!(item as ItemImpl);
1125
1126    let self_type = match *input.self_ty {
1127        syn::Type::Path(ref type_path) => {
1128            let segment = type_path.path.segments.last().unwrap();
1129            segment.clone() //ident.clone()
1130        }
1131        _ => {
1132            return TokenStream::from(
1133                syn::Error::new_spanned(input.self_ty, "`handle` argument must be a type")
1134                    .to_compile_error(),
1135            );
1136        }
1137    };
1138
1139    let trait_name = match &input.trait_ {
1140        Some((trait_path, _)) => trait_path.segments.last().unwrap().clone(),
1141        None => {
1142            return TokenStream::from(
1143                syn::Error::new_spanned(input.self_ty, "no trait in implementation block")
1144                    .to_compile_error(),
1145            );
1146        }
1147    };
1148
1149    let expanded = quote! {
1150        #input
1151
1152        #[hyperactor::internal_macro_support::async_trait::async_trait]
1153        impl hyperactor::Handler<#message_type> for #self_type {
1154            async fn handle(
1155                &mut self,
1156                cx: &hyperactor::Context<Self>,
1157                message: #message_type,
1158            ) -> hyperactor::internal_macro_support::anyhow::Result<()> {
1159                <Self as #trait_name>::handle(self, cx, message).await
1160            }
1161        }
1162    };
1163
1164    TokenStream::from(expanded)
1165}
1166
1167/// Use this macro in place of tracing::instrument to prevent spamming our tracing table.
1168/// We set a default level of INFO while always setting ERROR if the function returns Result::Err giving us
1169/// consistent and high quality structured logs. Because this wraps around tracing::instrument, all parameters
1170/// mentioned in https://fburl.com/9jlkb5q4 should be valid. For functions that don't return a [`Result`] type, use
1171/// [`instrument_infallible`]
1172///
1173/// ```
1174/// #[telemetry::instrument]
1175/// async fn yolo() -> anyhow::Result<i32> {
1176///     Ok(420)
1177/// }
1178/// ```
1179#[proc_macro_attribute]
1180pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream {
1181    let args =
1182        parse_macro_input!(args with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated);
1183    let input = parse_macro_input!(input as ItemFn);
1184    let output = quote! {
1185        #[hyperactor::internal_macro_support::tracing::instrument(err, skip_all, #args)]
1186        #input
1187    };
1188
1189    TokenStream::from(output)
1190}
1191
1192/// Use this macro in place of tracing::instrument to prevent spamming our tracing table.
1193/// Because this wraps around tracing::instrument, all parameters mentioned in
1194/// https://fburl.com/9jlkb5q4 should be valid.
1195///
1196/// ```
1197/// #[telemetry::instrument]
1198/// async fn yolo() -> i32 {
1199///     420
1200/// }
1201/// ```
1202#[proc_macro_attribute]
1203pub fn instrument_infallible(args: TokenStream, input: TokenStream) -> TokenStream {
1204    let args =
1205        parse_macro_input!(args with Punctuated::<syn::Expr, syn::Token![,]>::parse_terminated);
1206    let input = parse_macro_input!(input as ItemFn);
1207
1208    let output = quote! {
1209        #[hyperactor::internal_macro_support::tracing::instrument(skip_all, #args)]
1210        #input
1211    };
1212
1213    TokenStream::from(output)
1214}
1215
1216struct HandlerSpec {
1217    ty: Type,
1218}
1219
1220impl Parse for HandlerSpec {
1221    fn parse(input: ParseStream) -> syn::Result<Self> {
1222        let ty: Type = input.parse()?;
1223
1224        if input.is_empty() || input.peek(Token![,]) {
1225            Ok(HandlerSpec { ty })
1226        } else {
1227            // Something unexpected follows the type
1228            let unexpected: proc_macro2::TokenTree = input.parse()?;
1229            Err(syn::Error::new_spanned(
1230                unexpected,
1231                "unexpected token after type — use the bare message type",
1232            ))
1233        }
1234    }
1235}
1236
1237impl HandlerSpec {
1238    fn message_types(handlers: Vec<HandlerSpec>) -> Vec<Type> {
1239        handlers.into_iter().map(|handler| handler.ty).collect()
1240    }
1241}
1242
1243fn named_impl(data_type_name: &Ident, generics: &syn::Generics) -> proc_macro2::TokenStream {
1244    let generics_with_bounds = generics_with_named_bounds(generics);
1245    let type_params: Vec<_> = generics.type_params().collect();
1246    let has_generics = !type_params.is_empty();
1247
1248    let (impl_generics_with_bounds, _, _) = generics_with_bounds.split_for_impl();
1249    let (_, ty_generics, where_clause) = generics.split_for_impl();
1250
1251    let (typename_impl, typehash_impl) = if has_generics {
1252        let placeholders = vec!["{}"; type_params.len()].join(", ");
1253        let placeholders_format_string = format!("<{}>", placeholders);
1254        let format_string = quote! {
1255            concat!(
1256                std::module_path!(),
1257                "::",
1258                stringify!(#data_type_name),
1259                #placeholders_format_string
1260            )
1261        };
1262        let type_param_idents: Vec<_> = type_params.iter().map(|param| &param.ident).collect();
1263        (
1264            quote! {
1265                typeuri::intern_typename!(Self, #format_string, #(#type_param_idents),*)
1266            },
1267            quote! {
1268                typeuri::cityhasher::hash(Self::typename())
1269            },
1270        )
1271    } else {
1272        (
1273            quote! {
1274                concat!(std::module_path!(), "::", stringify!(#data_type_name))
1275            },
1276            quote! {
1277                static TYPEHASH: std::sync::LazyLock<u64> = std::sync::LazyLock::new(|| {
1278                    typeuri::cityhasher::hash(<#data_type_name as typeuri::Named>::typename())
1279                });
1280                *TYPEHASH
1281            },
1282        )
1283    };
1284
1285    quote! {
1286        impl #impl_generics_with_bounds typeuri::Named for #data_type_name #ty_generics #where_clause {
1287            fn typename() -> &'static str {
1288                #typename_impl
1289            }
1290
1291            fn typehash() -> u64 {
1292                #typehash_impl
1293            }
1294        }
1295    }
1296}
1297
1298fn generics_with_named_bounds(generics: &syn::Generics) -> syn::Generics {
1299    let mut generics = generics.clone();
1300    for param in generics.type_params_mut() {
1301        param.bounds.push(syn::parse_quote!(typeuri::Named));
1302    }
1303    generics
1304}
1305
1306fn generics_with_predicates(
1307    generics: &syn::Generics,
1308    predicates: impl IntoIterator<Item = WherePredicate>,
1309) -> syn::Generics {
1310    let mut generics = generics.clone();
1311    generics.make_where_clause().predicates.extend(predicates);
1312    generics
1313}
1314
1315/// Attribute Struct for [`fn export`] macro.
1316struct ExportAttr {
1317    handlers: Vec<HandlerSpec>,
1318}
1319
1320impl Parse for ExportAttr {
1321    fn parse(input: ParseStream) -> syn::Result<Self> {
1322        if input.is_empty() {
1323            return Ok(Self {
1324                handlers: Vec::new(),
1325            });
1326        }
1327
1328        let compatibility_form = {
1329            let fork = input.fork();
1330            fork.parse::<Ident>().is_ok() && fork.parse::<Token![=]>().is_ok()
1331        };
1332
1333        if !compatibility_form {
1334            let handlers = input
1335                .parse_terminated(HandlerSpec::parse, Token![,])?
1336                .into_iter()
1337                .collect();
1338            return Ok(Self { handlers });
1339        }
1340
1341        let mut handlers: Vec<HandlerSpec> = vec![];
1342
1343        while !input.is_empty() {
1344            let key: Ident = input.parse()?;
1345            input.parse::<Token![=]>()?;
1346
1347            if key == "spawn" {
1348                let expr: Expr = input.parse()?;
1349                return Err(syn::Error::new_spanned(
1350                    expr,
1351                    "`spawn = true` is no longer supported; use `#[spawnable]` on concrete actor declarations or `hyperactor::register_spawnable!(ConcreteType)` for generic instantiations",
1352                ));
1353            } else if key == "handlers" {
1354                let content;
1355                bracketed!(content in input);
1356                let raw_handlers = content.parse_terminated(HandlerSpec::parse, Token![,])?;
1357                handlers = raw_handlers.into_iter().collect();
1358            } else {
1359                return Err(syn::Error::new_spanned(
1360                    key,
1361                    "unexpected key in `#[export(...)]`. Use direct handler lists, or the compatibility key `handlers`",
1362                ));
1363            }
1364
1365            // optional trailing comma
1366            let _ = input.parse::<Token![,]>();
1367        }
1368
1369        Ok(ExportAttr { handlers })
1370    }
1371}
1372
1373/// Exports handlers for this actor. The set of exported handlers
1374/// determine the messages that may be sent to remote references of
1375/// the actor ([`hyperaxtor::ActorRef`]). Only messages that implement
1376/// [`hyperactor::RemoteMessage`] may be exported.
1377///
1378/// # Example
1379///
1380/// In the following example, `MyActor` exports handlers for two message types,
1381/// `MyMessage` and `MyOtherMessage`. Consequently, `ActorRef`s of the actor's
1382/// type may dispatch messages of these types.
1383///
1384/// ```ignore
1385/// #[export(MyMessage, MyOtherMessage)]
1386/// struct MyActor {}
1387/// ```
1388#[proc_macro_attribute]
1389pub fn export(attr: TokenStream, item: TokenStream) -> TokenStream {
1390    let input: DeriveInput = parse_macro_input!(item as DeriveInput);
1391    let data_type_name = &input.ident;
1392    let (_, ty_generics, _) = input.generics.split_for_impl();
1393    let named_generics = generics_with_named_bounds(&input.generics);
1394    let (named_impl_generics, named_ty_generics, named_where_clause) =
1395        named_generics.split_for_impl();
1396
1397    let ExportAttr { handlers } = parse_macro_input!(attr as ExportAttr);
1398
1399    let mut handles = Vec::new();
1400    let mut bindings = Vec::new();
1401    let mut bind_predicates = Vec::new();
1402    let actor_ty: Type = syn::parse_quote!(#data_type_name #ty_generics);
1403
1404    for HandlerSpec { ty } in &handlers {
1405        let message_generics = generics_with_predicates(
1406            &named_generics,
1407            [syn::parse_quote!(#ty: hyperactor::RemoteMessage)],
1408        );
1409        let (message_impl_generics, message_ty_generics, message_where_clause) =
1410            message_generics.split_for_impl();
1411        handles.push(quote! {
1412            impl #message_impl_generics hyperactor::actor::RemoteHandles<#ty>
1413                for #data_type_name #message_ty_generics #message_where_clause {}
1414            impl #message_impl_generics hyperactor::remote::Accepts<#ty>
1415                for #data_type_name #message_ty_generics #message_where_clause {}
1416        });
1417        bindings.push(quote! {
1418            ports.bind::<#ty>();
1419        });
1420        bind_predicates.push(syn::parse_quote!(#ty: hyperactor::RemoteMessage));
1421        bind_predicates.push(syn::parse_quote!(#actor_ty: hyperactor::Handler<#ty>));
1422    }
1423
1424    let bind_generics = generics_with_predicates(&named_generics, bind_predicates);
1425    let (bind_impl_generics, bind_ty_generics, bind_where_clause) = bind_generics.split_for_impl();
1426    let named_impl = named_impl(data_type_name, &input.generics);
1427
1428    let expanded = quote! {
1429        #input
1430
1431        impl #named_impl_generics hyperactor::actor::Referable for #data_type_name #named_ty_generics #named_where_clause {}
1432
1433        #(#handles)*
1434
1435        // Always export the `IntrospectMessage` type.
1436        impl #named_impl_generics hyperactor::actor::RemoteHandles<hyperactor::introspect::IntrospectMessage> for #data_type_name #named_ty_generics #named_where_clause {}
1437        impl #named_impl_generics hyperactor::remote::Accepts<hyperactor::introspect::IntrospectMessage> for #data_type_name #named_ty_generics #named_where_clause {}
1438
1439        impl #bind_impl_generics hyperactor::actor::Binds<#data_type_name #bind_ty_generics> for #data_type_name #bind_ty_generics #bind_where_clause {
1440            fn bind(ports: &hyperactor::proc::HandlerPorts<Self>) {
1441                #(#bindings)*
1442            }
1443        }
1444
1445        #named_impl
1446    };
1447
1448    TokenStream::from(expanded)
1449}
1450
1451/// Marks a concrete actor declaration as remotely spawnable.
1452///
1453/// When combined with `#[export(...)]`, place `#[spawnable]` below `#[export]`.
1454#[proc_macro_attribute]
1455pub fn spawnable(attr: TokenStream, item: TokenStream) -> TokenStream {
1456    if !attr.is_empty() {
1457        return syn::Error::new(Span::call_site(), "`#[spawnable]` does not take arguments")
1458            .to_compile_error()
1459            .into();
1460    }
1461
1462    let input: DeriveInput = parse_macro_input!(item as DeriveInput);
1463    if !matches!(input.data, Data::Struct(_)) {
1464        return syn::Error::new(
1465            input.span(),
1466            "`#[spawnable]` only supports struct actor declarations",
1467        )
1468        .to_compile_error()
1469        .into();
1470    }
1471
1472    if !input.generics.params.is_empty() {
1473        return syn::Error::new(
1474            input.generics.span(),
1475            "generic actor families cannot use `#[spawnable]`; use `hyperactor::register_spawnable!(ConcreteType)` instead",
1476        )
1477        .to_compile_error()
1478        .into();
1479    }
1480
1481    let data_type_name = &input.ident;
1482    quote! {
1483        #input
1484        hyperactor::register_spawnable!(#data_type_name);
1485    }
1486    .into()
1487}
1488
1489/// Represents the full input to [`fn behavior`].
1490struct BehaviorInput {
1491    behavior: Ident,
1492    generics: syn::Generics,
1493    handlers: Vec<HandlerSpec>,
1494}
1495
1496impl syn::parse::Parse for BehaviorInput {
1497    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
1498        let behavior: Ident = input.parse()?;
1499        let generics: syn::Generics = input.parse()?;
1500        let _: Token![,] = input.parse()?;
1501        let raw_handlers = input.parse_terminated(HandlerSpec::parse, Token![,])?;
1502        let handlers = raw_handlers.into_iter().collect();
1503        Ok(BehaviorInput {
1504            behavior,
1505            generics,
1506            handlers,
1507        })
1508    }
1509}
1510
1511/// Create a [`Referable`] definition, handling a specific set of message types.
1512/// Behaviors are used to create an [`ActorRef`] without having to depend on the
1513/// actor's implementation. Casts are supported for any remote message handled
1514/// by the actor.
1515///
1516/// ```
1517/// hyperactor::behavior!(TestActorBehavior, TestMessage, (), MyGeneric<()>, u64,);
1518/// ```
1519///
1520/// This macro also supports generic behaviors:
1521/// ```
1522/// hyperactor::behavior!(TestBehavior<T>, Message<T>, u64,);
1523/// ```
1524#[proc_macro]
1525pub fn behavior(input: TokenStream) -> TokenStream {
1526    let BehaviorInput {
1527        behavior,
1528        generics,
1529        handlers,
1530    } = parse_macro_input!(input as BehaviorInput);
1531    let tys = HandlerSpec::message_types(handlers);
1532
1533    // Add bounds to generics for Named, Serialize, Deserialize
1534    let mut bounded_generics = generics.clone();
1535    for param in bounded_generics.type_params_mut() {
1536        param.bounds.push(syn::parse_quote!(typeuri::Named));
1537        param.bounds.push(syn::parse_quote!(serde::Serialize));
1538        param.bounds.push(syn::parse_quote!(std::marker::Send));
1539        param.bounds.push(syn::parse_quote!(std::marker::Sync));
1540        param.bounds.push(syn::parse_quote!(std::fmt::Debug));
1541        // Note: lifetime parameters are not *actually* hygienic.
1542        // https://github.com/rust-lang/rust/issues/54727
1543        let lifetime =
1544            syn::Lifetime::new("'hyperactor_behavior_de", proc_macro2::Span::mixed_site());
1545        param
1546            .bounds
1547            .push(syn::parse_quote!(for<#lifetime> serde::Deserialize<#lifetime>));
1548    }
1549
1550    // Split the generics for use in different contexts
1551    let (impl_generics, ty_generics, where_clause) = bounded_generics.split_for_impl();
1552
1553    // Create a combined generics for the Binds impl that includes both A and the behavior's generics
1554    let mut binds_generics = bounded_generics.clone();
1555    binds_generics.params.insert(
1556        0,
1557        syn::GenericParam::Type(syn::TypeParam {
1558            attrs: vec![],
1559            ident: Ident::new("A", proc_macro2::Span::call_site()),
1560            colon_token: None,
1561            bounds: Punctuated::new(),
1562            default: None,
1563        }),
1564    );
1565    let (binds_impl_generics, _, _) = binds_generics.split_for_impl();
1566
1567    // Determine typename and typehash implementation based on whether we have generics
1568    let type_params: Vec<_> = bounded_generics.type_params().collect();
1569    let has_generics = !type_params.is_empty();
1570
1571    let (typename_impl, typehash_impl) = if has_generics {
1572        // Create format string with placeholders for each generic parameter
1573        let placeholders = vec!["{}"; type_params.len()].join(", ");
1574        let placeholders_format_string = format!("<{}>", placeholders);
1575        let format_string = quote! { concat!(std::module_path!(), "::", stringify!(#behavior), #placeholders_format_string) };
1576
1577        let type_param_idents: Vec<_> = type_params.iter().map(|p| &p.ident).collect();
1578        (
1579            quote! {
1580                typeuri::intern_typename!(Self, #format_string, #(#type_param_idents),*)
1581            },
1582            quote! {
1583                typeuri::cityhasher::hash(Self::typename())
1584            },
1585        )
1586    } else {
1587        (
1588            quote! {
1589                concat!(std::module_path!(), "::", stringify!(#behavior))
1590            },
1591            quote! {
1592                static TYPEHASH: std::sync::LazyLock<u64> = std::sync::LazyLock::new(|| {
1593                    typeuri::cityhasher::hash(<#behavior as typeuri::Named>::typename())
1594                });
1595                *TYPEHASH
1596            },
1597        )
1598    };
1599
1600    let type_param_idents = generics.type_params().map(|p| &p.ident).collect::<Vec<_>>();
1601
1602    let expanded = quote! {
1603        #[doc = "The generated behavior struct."]
1604        #[derive(Debug, serde::Serialize, serde::Deserialize)]
1605        pub struct #behavior #impl_generics #where_clause {
1606            _phantom: std::marker::PhantomData<(#(#type_param_idents),*)>
1607        }
1608
1609        impl #impl_generics typeuri::Named for #behavior #ty_generics #where_clause {
1610            fn typename() -> &'static str {
1611                #typename_impl
1612            }
1613
1614            fn typehash() -> u64 {
1615                #typehash_impl
1616            }
1617        }
1618
1619        impl #impl_generics hyperactor::actor::Referable for #behavior #ty_generics #where_clause {}
1620
1621        impl #binds_impl_generics hyperactor::actor::Binds<A> for #behavior #ty_generics
1622        where
1623            A: hyperactor::Actor #(+ hyperactor::Handler<#tys>)*,
1624            #where_clause
1625        {
1626            fn bind(ports: &hyperactor::proc::HandlerPorts<A>) {
1627                #(
1628                    ports.bind::<#tys>();
1629                )*
1630            }
1631        }
1632
1633        #(
1634            impl #impl_generics hyperactor::actor::RemoteHandles<#tys> for #behavior #ty_generics #where_clause {}
1635            impl #impl_generics hyperactor::remote::Accepts<#tys> for #behavior #ty_generics #where_clause {}
1636        )*
1637    };
1638
1639    TokenStream::from(expanded)
1640}
1641
1642// Helper function for common parsing and validation
1643fn parse_observe_function(
1644    attr: TokenStream,
1645    item: TokenStream,
1646) -> syn::Result<(ItemFn, String, String)> {
1647    let input = syn::parse::<ItemFn>(item)?;
1648
1649    if input.sig.asyncness.is_none() {
1650        return Err(syn::Error::new(
1651            input.sig.span(),
1652            "observe macros can only be applied to async functions",
1653        ));
1654    }
1655
1656    let fn_name_str = input.sig.ident.to_string();
1657    let module_name_str = syn::parse::<syn::LitStr>(attr)?.value();
1658
1659    Ok((input, fn_name_str, module_name_str))
1660}
1661
1662// Helper function for creating telemetry identifiers and setup code
1663fn create_telemetry_setup(
1664    module_name_str: &str,
1665    fn_name_str: &str,
1666    include_error: bool,
1667) -> (Ident, Ident, Option<Ident>, proc_macro2::TokenStream) {
1668    let module_and_fn = format!("{}_{}", module_name_str, fn_name_str);
1669    let latency_ident = Ident::new("latency", Span::from(proc_macro::Span::def_site()));
1670
1671    let success_ident = Ident::new("success", Span::from(proc_macro::Span::def_site()));
1672
1673    let error_ident = if include_error {
1674        Some(Ident::new(
1675            "error",
1676            Span::from(proc_macro::Span::def_site()),
1677        ))
1678    } else {
1679        None
1680    };
1681
1682    let error_declaration = if let Some(ref error_ident) = error_ident {
1683        quote! {
1684            hyperactor_telemetry::declare_static_counter!(#error_ident, concat!(#module_and_fn, ".error"));
1685        }
1686    } else {
1687        quote! {}
1688    };
1689
1690    let setup_code = quote! {
1691        use hyperactor_telemetry;
1692        hyperactor_telemetry::declare_static_timer!(#latency_ident, concat!(#module_and_fn, ".latency"), hyperactor_telemetry::TimeUnit::Micros);
1693        hyperactor_telemetry::declare_static_counter!(#success_ident, concat!(#module_and_fn, ".success"));
1694        #error_declaration
1695    };
1696
1697    (latency_ident, success_ident, error_ident, setup_code)
1698}
1699
1700/// A procedural macro that automatically injects telemetry code into async functions
1701/// that return a Result type.
1702///
1703/// This macro wraps async functions and adds instrumentation to measure:
1704/// 1. Latency - how long the function takes to execute
1705/// 2. Error counter - function error count
1706/// 3. Success counter - function completion count
1707///
1708/// # Example
1709///
1710/// ```rust
1711/// use hyperactor_actor::observe_result;
1712///
1713/// #[observe_result("my_module")]
1714/// async fn process_request(user_id: &str) -> Result<String, Error> {
1715///     // Function implementation
1716///     // Telemetry will be automatically collected
1717/// }
1718/// ```
1719#[proc_macro_attribute]
1720pub fn observe_result(attr: TokenStream, item: TokenStream) -> TokenStream {
1721    let (input, fn_name_str, module_name_str) = match parse_observe_function(attr, item) {
1722        Ok(parsed) => parsed,
1723        Err(err) => return err.to_compile_error().into(),
1724    };
1725
1726    let fn_name = &input.sig.ident;
1727    let vis = &input.vis;
1728    let args = &input.sig.inputs;
1729    let return_type = &input.sig.output;
1730    let body = &input.block;
1731    let attrs = &input.attrs;
1732    let generics = &input.sig.generics;
1733
1734    let (latency_ident, success_ident, error_ident, telemetry_setup) =
1735        create_telemetry_setup(&module_name_str, &fn_name_str, true);
1736    let error_ident = error_ident.unwrap();
1737
1738    let result_ident = Ident::new("result", Span::from(proc_macro::Span::def_site()));
1739
1740    // Generate the instrumented function
1741    let expanded = quote! {
1742        #(#attrs)*
1743        #vis async fn #fn_name #generics(#args) #return_type {
1744            #telemetry_setup
1745
1746            let kv_pairs = hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone());
1747            let _timer = #latency_ident.start(kv_pairs);
1748
1749            let #result_ident = async #body.await;
1750
1751            match &#result_ident {
1752                Ok(_) => {
1753                    #success_ident.add(
1754                        1,
1755                        hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1756                    );
1757                }
1758                Err(_) => {
1759                    #error_ident.add(
1760                        1,
1761                        hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1762                    );
1763                }
1764            }
1765
1766            #result_ident
1767        }
1768    };
1769
1770    expanded.into()
1771}
1772
1773/// A procedural macro that automatically injects telemetry code into async functions
1774/// that do not return a Result type.
1775///
1776/// This macro wraps async functions and adds instrumentation to measure:
1777/// 1. Latency - how long the function takes to execute
1778/// 2. Success counter - function completion count
1779///
1780/// # Example
1781///
1782/// ```rust
1783/// use hyperactor_actor::observe_async;
1784///
1785/// #[observe_async("my_module")]
1786/// async fn process_data(data: &str) -> String {
1787///     // Function implementation
1788///     // Telemetry will be automatically collected
1789/// }
1790/// ```
1791#[proc_macro_attribute]
1792pub fn observe_async(attr: TokenStream, item: TokenStream) -> TokenStream {
1793    let (input, fn_name_str, module_name_str) = match parse_observe_function(attr, item) {
1794        Ok(parsed) => parsed,
1795        Err(err) => return err.to_compile_error().into(),
1796    };
1797
1798    let fn_name = &input.sig.ident;
1799    let vis = &input.vis;
1800    let args = &input.sig.inputs;
1801    let return_type = &input.sig.output;
1802    let body = &input.block;
1803    let attrs = &input.attrs;
1804    let generics = &input.sig.generics;
1805
1806    let (latency_ident, success_ident, _, telemetry_setup) =
1807        create_telemetry_setup(&module_name_str, &fn_name_str, false);
1808
1809    let return_ident = Ident::new("ret", Span::from(proc_macro::Span::def_site()));
1810
1811    // Generate the instrumented function
1812    let expanded = quote! {
1813        #(#attrs)*
1814        #vis async fn #fn_name #generics(#args) #return_type {
1815            #telemetry_setup
1816
1817            let kv_pairs = hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone());
1818            let _timer = #latency_ident.start(kv_pairs);
1819
1820            let #return_ident = async #body.await;
1821
1822            #success_ident.add(
1823                1,
1824                hyperactor_telemetry::kv_pairs!("function" => #fn_name_str.clone())
1825            );
1826            #return_ident
1827        }
1828    };
1829
1830    expanded.into()
1831}
1832
1833fn validate_label(s: &str) -> Result<(), String> {
1834    if s.is_empty() {
1835        return Err("label must not be empty".to_string());
1836    }
1837    if s.len() > 63 {
1838        return Err("label exceeds 63 characters".to_string());
1839    }
1840    let first = s.as_bytes()[0];
1841    if !first.is_ascii_lowercase() {
1842        return Err("label must start with a lowercase letter".to_string());
1843    }
1844    let last = s.as_bytes()[s.len() - 1];
1845    if !last.is_ascii_lowercase() && !last.is_ascii_digit() {
1846        return Err("label must end with a lowercase letter or digit".to_string());
1847    }
1848    for ch in s.chars() {
1849        if !ch.is_ascii_lowercase() && !ch.is_ascii_digit() && ch != '-' {
1850            return Err(format!("label contains invalid character '{ch}'"));
1851        }
1852    }
1853    Ok(())
1854}
1855
1856fn validate_hex_uid(s: &str) -> Result<u64, String> {
1857    if s.is_empty() || s.len() > 16 {
1858        return Err(format!("hex uid must be 1-16 hex characters, got '{s}'"));
1859    }
1860    for ch in s.chars() {
1861        if !ch.is_ascii_hexdigit() {
1862            return Err(format!("hex uid contains invalid character '{ch}'"));
1863        }
1864    }
1865    u64::from_str_radix(s, 16).map_err(|e| format!("invalid hex uid '{s}': {e}"))
1866}
1867
1868/// Compile-time validated [`hyperactor::id::Uid`] construction.
1869///
1870/// Accepts two forms:
1871/// - `uid!(_my-singleton)` — a singleton Uid
1872/// - `uid!(d5d54d7201103869)` — an instance Uid
1873#[proc_macro]
1874pub fn uid(input: TokenStream) -> TokenStream {
1875    let input2: proc_macro2::TokenStream = input.into();
1876    let combined: String = input2.into_iter().map(|tt| tt.to_string()).collect();
1877
1878    if combined.is_empty() {
1879        return TokenStream::from(quote! { compile_error!("uid! macro requires an argument") });
1880    }
1881
1882    // Singleton: starts with '_'
1883    if let Some(rest) = combined.strip_prefix('_') {
1884        return match validate_label(rest) {
1885            Ok(()) => TokenStream::from(quote! {
1886                hyperactor::id::Uid::Singleton(
1887                    hyperactor::id::Label::new(#rest).unwrap()
1888                )
1889            }),
1890            Err(e) => {
1891                let msg = format!("invalid singleton uid: {e}");
1892                TokenStream::from(quote! { compile_error!(#msg) })
1893            }
1894        };
1895    }
1896
1897    // Instance: bare hex
1898    match validate_hex_uid(&combined) {
1899        Ok(uid_val) => TokenStream::from(quote! {
1900            hyperactor::id::Uid::Instance(#uid_val, None)
1901        }),
1902        Err(e) => {
1903            let msg = format!("invalid uid: {e}");
1904            TokenStream::from(quote! { compile_error!(#msg) })
1905        }
1906    }
1907}
hyperactor_macros/lib.rs

hyperactor_macros/
lib.rs
