hyperactor_config/
attrs.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//! Attribute dictionary for type-safe, heterogeneous key-value storage with serde support.
10//!
11//! This module provides `Attrs`, a type-safe dictionary that can store heterogeneous values
12//! and serialize/deserialize them using serde. All stored values must implement
13//! `AttrValue` to ensure the entire dictionary can be serialized.
14//!
15//! Keys are automatically registered at compile time using the `declare_attrs!` macro and the
16//! inventory crate, eliminating the need for manual registry management.
17//!
18//! # Basic Usage
19//!
20//! ```
21//! use std::time::Duration;
22//!
23//! use hyperactor_config::attrs::Attrs;
24//! use hyperactor_config::attrs::declare_attrs;
25//!
26//! // Declare keys with their associated types
27//! declare_attrs! {
28//!    /// Request timeout
29//!    attr TIMEOUT: Duration;
30//!
31//!   /// Maximum retry count
32//!   attr MAX_RETRIES: u32 = 3;  // with default value
33//! }
34//!
35//! let mut attrs = Attrs::new();
36//! attrs.set(TIMEOUT, Duration::from_secs(30));
37//!
38//! assert_eq!(attrs.get(TIMEOUT), Some(&Duration::from_secs(30)));
39//! assert_eq!(attrs.get(MAX_RETRIES), Some(&3));
40//! ```
41//!
42//! # Serialization
43//!
44//! `Attrs` can be serialized to and deserialized automatically:
45//!
46//! ```
47//! use std::time::Duration;
48//!
49//! use hyperactor_config::attrs::Attrs;
50//! use hyperactor_config::attrs::declare_attrs;
51//!
52//! declare_attrs! {
53//!   /// Request timeout
54//!   pub attr TIMEOUT: Duration;
55//! }
56//!
57//! let mut attrs = Attrs::new();
58//! attrs.set(TIMEOUT, Duration::from_secs(30));
59//!
60//! // Serialize to JSON
61//! let json = serde_json::to_string(&attrs).unwrap();
62//!
63//! // Deserialize from JSON (no manual registry needed!)
64//! let deserialized: Attrs = serde_json::from_str(&json).unwrap();
65//!
66//! assert_eq!(deserialized.get(TIMEOUT), Some(&Duration::from_secs(30)));
67//! ```
68//!
69//! ## Meta attributes
70//!
71//! An attribute can be assigned a set of attribute values,
72//! associated with the attribute key. These are specified by
73//! @-annotations in the `declare_attrs!` macro:
74//!
75//! ```
76//! use std::time::Duration;
77//!
78//! use hyperactor_config::attrs::Attrs;
79//! use hyperactor_config::attrs::declare_attrs;
80//!
81//! declare_attrs! {
82//!   /// Is experimental?
83//!   pub attr EXPERIMENTAL: bool;
84//!
85//!   /// Request timeout
86//!   @meta(EXPERIMENTAL = true)
87//!   pub attr TIMEOUT: Duration;
88//! }
89//!
90//! assert!(TIMEOUT.attrs().get(EXPERIMENTAL).unwrap());
91//! ```
92//!
93//! Meta attributes can be used to provide more generic functionality
94//! on top of the basic attributes. For example, a library can use
95//! meta-attributes to specify the behavior of an attribute.
96
97use std::any::Any;
98use std::collections::HashMap;
99use std::fmt::Display;
100use std::ops::Index;
101use std::ops::IndexMut;
102use std::str::FromStr;
103use std::sync::LazyLock;
104
105use chrono::DateTime;
106use chrono::Utc;
107use erased_serde::Deserializer as ErasedDeserializer;
108use erased_serde::Serialize as ErasedSerialize;
109use serde::Deserialize;
110use serde::Deserializer;
111use serde::Serialize;
112use serde::Serializer;
113use serde::de::DeserializeOwned;
114use serde::de::MapAccess;
115use serde::de::Visitor;
116use serde::ser::SerializeMap;
117use typeuri::Named;
118
119// Information about an attribute key, used for automatic registration.
120// This needs to be public to be accessible from other crates, but it is
121// not part of the public API.
122#[doc(hidden)]
123pub struct AttrKeyInfo {
124    /// Name of the key
125    pub name: &'static str,
126    /// Function to get the type hash of the associated value type
127    pub typehash: fn() -> u64,
128    /// Deserializer function that deserializes directly from any deserializer
129    pub deserialize_erased:
130        fn(&mut dyn ErasedDeserializer) -> Result<Box<dyn SerializableValue>, erased_serde::Error>,
131    /// Meta-attributes.
132    pub meta: &'static LazyLock<Attrs>,
133    /// Display an attribute value using AttrValue::display.
134    pub display: fn(&dyn SerializableValue) -> String,
135    /// Parse an attribute value using AttrValue::parse.
136    pub parse: fn(&str) -> Result<Box<dyn SerializableValue>, anyhow::Error>,
137    /// Default value for the attribute, if any.
138    pub default: Option<&'static dyn SerializableValue>,
139    /// A reference to the relevant key object with the associated
140    /// type parameter erased. Can be downcast to a concrete Key<T>.
141    pub erased: &'static dyn ErasedKey,
142}
143
144inventory::collect!(AttrKeyInfo);
145
146/// A typed key for the attribute dictionary.
147///
148/// Each key is associated with a specific type T and has a unique name.
149/// Keys are typically created using the `declare_attrs!` macro which ensures they have
150/// static lifetime and automatically registers them for serialization.
151pub struct Key<T: 'static> {
152    name: &'static str,
153    default_value: Option<&'static T>,
154    attrs: &'static LazyLock<Attrs>,
155}
156
157impl<T> Key<T> {
158    /// Returns the name of this key.
159    pub fn name(&self) -> &'static str {
160        self.name
161    }
162}
163
164impl<T: Named + 'static> Key<T> {
165    /// Creates a new key with the given name.
166    pub const fn new(
167        name: &'static str,
168        default_value: Option<&'static T>,
169        attrs: &'static LazyLock<Attrs>,
170    ) -> Self {
171        Self {
172            name,
173            default_value,
174            attrs,
175        }
176    }
177
178    /// Returns a reference to the default value for this key, if one exists.
179    pub fn default(&self) -> Option<&'static T> {
180        self.default_value
181    }
182
183    /// Returns whether this key has a default value.
184    pub fn has_default(&self) -> bool {
185        self.default_value.is_some()
186    }
187
188    /// Returns the type hash of the associated value type.
189    pub fn typehash(&self) -> u64 {
190        T::typehash()
191    }
192
193    /// The attributes associated with this key.
194    pub fn attrs(&self) -> &'static LazyLock<Attrs> {
195        self.attrs
196    }
197}
198
199impl<T: 'static> Clone for Key<T> {
200    fn clone(&self) -> Self {
201        // Use Copy.
202        *self
203    }
204}
205
206impl<T: 'static> Copy for Key<T> {}
207
208/// A trait for type-erased keys.
209pub trait ErasedKey: Any + Send + Sync + 'static {
210    /// The name of the key.
211    fn name(&self) -> &'static str;
212
213    /// The typehash of the key's associated type.
214    fn typehash(&self) -> u64;
215
216    /// The typename of the key's associated type.
217    fn typename(&self) -> &'static str;
218}
219
220impl dyn ErasedKey {
221    /// Downcast a type-erased key to a specific key type.
222    pub fn downcast_ref<T: Named + 'static>(&'static self) -> Option<&'static Key<T>> {
223        (self as &dyn Any).downcast_ref::<Key<T>>()
224    }
225}
226
227impl<T: AttrValue> ErasedKey for Key<T> {
228    fn name(&self) -> &'static str {
229        self.name
230    }
231
232    fn typehash(&self) -> u64 {
233        T::typehash()
234    }
235
236    fn typename(&self) -> &'static str {
237        T::typename()
238    }
239}
240
241// Enable attr[key] syntax.
242impl<T: AttrValue> Index<Key<T>> for Attrs {
243    type Output = T;
244
245    fn index(&self, key: Key<T>) -> &Self::Output {
246        self.get(key).unwrap()
247    }
248}
249
250// TODO: separately type keys with defaults, so that we can statically enforce that indexmut is only
251// called on keys with defaults.
252impl<T: AttrValue> IndexMut<Key<T>> for Attrs {
253    fn index_mut(&mut self, key: Key<T>) -> &mut Self::Output {
254        self.get_mut(key).unwrap()
255    }
256}
257
258/// This trait must be implemented by all attribute values. In addition to enforcing
259/// the supertrait `Named + Sized + Serialize + DeserializeOwned + Send + Sync + Clone`,
260/// `AttrValue` requires that the type be representable in "display" format.
261///
262/// `AttrValue` includes its own `display` and `parse` so that behavior can be tailored
263/// for attribute purposes specifically, allowing common types like `Duration` to be used
264/// without modification.
265///
266/// This crate includes a derive macro for AttrValue, which uses the type's
267/// `std::string::ToString` for display, and `std::str::FromStr` for parsing.
268pub trait AttrValue:
269    Named + Sized + Serialize + DeserializeOwned + Send + Sync + Clone + 'static
270{
271    /// Display the value, typically using [`std::fmt::Display`].
272    /// This is called to show the output in human-readable form.
273    fn display(&self) -> String;
274
275    /// Parse a value from a string, typically using [`std::str::FromStr`].
276    fn parse(value: &str) -> Result<Self, anyhow::Error>;
277}
278
279/// Macro to implement AttrValue for types that implement ToString and FromStr.
280///
281/// This macro provides a convenient way to implement AttrValue for types that already
282/// have string conversion capabilities through the standard ToString and FromStr traits.
283///
284/// # Usage
285///
286/// ```ignore
287/// impl_attrvalue!(i32, u64, f64);
288/// ```
289///
290/// This will generate AttrValue implementations for i32, u64, and f64 that use
291/// their ToString and FromStr implementations for display and parsing.
292#[macro_export]
293macro_rules! impl_attrvalue {
294    ($($ty:ty),+ $(,)?) => {
295        $(
296            impl $crate::attrs::AttrValue for $ty {
297                fn display(&self) -> String {
298                    self.to_string()
299                }
300
301                fn parse(value: &str) -> Result<Self, anyhow::Error> {
302                    value.parse().map_err(|e| anyhow::anyhow!("failed to parse {}: {}", stringify!($ty), e))
303                }
304            }
305        )+
306    };
307}
308
309// pub use impl_attrvalue;
310
311// Implement AttrValue for common standard library types
312impl_attrvalue!(
313    i8,
314    i16,
315    i32,
316    i64,
317    i128,
318    isize,
319    u8,
320    u16,
321    u32,
322    u64,
323    u128,
324    usize,
325    f32,
326    f64,
327    String,
328    std::net::IpAddr,
329    std::net::Ipv4Addr,
330    std::net::Ipv6Addr,
331);
332
333impl AttrValue for std::time::Duration {
334    fn display(&self) -> String {
335        humantime::format_duration(*self).to_string()
336    }
337
338    fn parse(value: &str) -> Result<Self, anyhow::Error> {
339        Ok(humantime::parse_duration(value)?)
340    }
341}
342
343impl AttrValue for std::time::SystemTime {
344    fn display(&self) -> String {
345        let datetime: DateTime<Utc> = (*self).into();
346        datetime.to_rfc3339()
347    }
348
349    fn parse(value: &str) -> Result<Self, anyhow::Error> {
350        let datetime = DateTime::parse_from_rfc3339(value)?;
351        Ok(datetime.into())
352    }
353}
354
355impl<T, E> AttrValue for std::ops::Range<T>
356where
357    T: Named
358        + Display
359        + FromStr<Err = E>
360        + Send
361        + Sync
362        + Serialize
363        + DeserializeOwned
364        + Clone
365        + 'static,
366    E: Into<anyhow::Error> + Send + Sync + 'static,
367{
368    fn display(&self) -> String {
369        format!("{}..{}", self.start, self.end)
370    }
371
372    fn parse(value: &str) -> Result<Self, anyhow::Error> {
373        let (start, end) = value.split_once("..").ok_or_else(|| {
374            anyhow::anyhow!("expected range in format `start..end`, got `{}`", value)
375        })?;
376        let start = start.parse().map_err(|e: E| e.into())?;
377        let end = end.parse().map_err(|e: E| e.into())?;
378        Ok(start..end)
379    }
380}
381
382impl AttrValue for bool {
383    fn display(&self) -> String {
384        if *self { 1.to_string() } else { 0.to_string() }
385    }
386
387    fn parse(value: &str) -> Result<Self, anyhow::Error> {
388        let value = value.to_ascii_lowercase();
389        match value.as_str() {
390            "0" | "false" => Ok(false),
391            "1" | "true" => Ok(true),
392            _ => Err(anyhow::anyhow!(
393                "expected `0`, `1`, `true` or `false`, got `{}`",
394                value
395            )),
396        }
397    }
398}
399
400// Internal trait for type-erased serialization
401#[doc(hidden)]
402pub trait SerializableValue: Send + Sync {
403    /// Get a reference to this value as Any for downcasting
404    fn as_any(&self) -> &dyn Any;
405    /// Get a mutable reference to this value as Any for downcasting
406    fn as_any_mut(&mut self) -> &mut dyn Any;
407    /// Get a reference to this value as an erased serializable trait object
408    fn as_erased_serialize(&self) -> &dyn ErasedSerialize;
409    /// Clone the underlying value, retaining dyn compatibility.
410    fn cloned(&self) -> Box<dyn SerializableValue>;
411    /// Display the value
412    fn display(&self) -> String;
413}
414
415impl<T: AttrValue> SerializableValue for T {
416    fn as_any(&self) -> &dyn Any {
417        self
418    }
419
420    fn as_any_mut(&mut self) -> &mut dyn Any {
421        self
422    }
423
424    fn as_erased_serialize(&self) -> &dyn ErasedSerialize {
425        self
426    }
427
428    fn cloned(&self) -> Box<dyn SerializableValue> {
429        Box::new(self.clone())
430    }
431
432    fn display(&self) -> String {
433        self.display()
434    }
435}
436
437/// A heterogeneous, strongly-typed attribute dictionary with serialization support.
438///
439/// This dictionary stores key-value pairs where:
440/// - Keys are type-safe and must be predefined with their associated types
441/// - Values must implement [`AttrValue`]
442/// - The entire dictionary can be serialized to/from JSON automatically
443///
444/// # Type Safety
445///
446/// The dictionary enforces type safety at compile time. You cannot retrieve a value
447/// with the wrong type, and the compiler will catch such errors.
448///
449/// # Serialization
450///
451/// The dictionary can be serialized using serde. During serialization, each value
452/// is serialized with its key name. During deserialization, the automatically registered
453/// key information is used to determine the correct type for each value.
454pub struct Attrs {
455    values: HashMap<&'static str, Box<dyn SerializableValue>>,
456}
457
458impl Attrs {
459    /// Create a new empty attribute dictionary.
460    pub fn new() -> Self {
461        Self {
462            values: HashMap::new(),
463        }
464    }
465
466    /// Set a value for the given key.
467    pub fn set<T: AttrValue>(&mut self, key: Key<T>, value: T) {
468        self.values.insert(key.name, Box::new(value));
469    }
470
471    fn maybe_set_from_default<T: AttrValue>(&mut self, key: Key<T>) {
472        if self.contains_key(key) {
473            return;
474        }
475        let Some(default) = key.default() else { return };
476        self.set(key, default.clone());
477    }
478
479    /// Get a value for the given key, returning None if not present. If the key has a default value,
480    /// that is returned instead.
481    pub fn get<T: AttrValue>(&self, key: Key<T>) -> Option<&T> {
482        self.values
483            .get(key.name)
484            .and_then(|value| value.as_any().downcast_ref::<T>())
485            .or_else(|| key.default())
486    }
487
488    /// Get a mutable reference to a value for the given key. If the key has a default value, it is
489    /// first set, and then returned as a mutable reference.
490    pub fn get_mut<T: AttrValue>(&mut self, key: Key<T>) -> Option<&mut T> {
491        self.maybe_set_from_default(key);
492        self.values
493            .get_mut(key.name)
494            .and_then(|value| value.as_any_mut().downcast_mut::<T>())
495    }
496
497    /// Remove a value for the given key, returning it if present.
498    pub fn remove<T: AttrValue>(&mut self, key: Key<T>) -> bool {
499        // TODO: return value (this is tricky because of the type erasure)
500        self.values.remove(key.name).is_some()
501    }
502
503    /// Checks if the given key exists in the dictionary.
504    pub fn contains_key<T: AttrValue>(&self, key: Key<T>) -> bool {
505        self.values.contains_key(key.name)
506    }
507
508    /// Returns the number of key-value pairs in the dictionary.
509    pub fn len(&self) -> usize {
510        self.values.len()
511    }
512
513    /// Returns true if the dictionary is empty.
514    pub fn is_empty(&self) -> bool {
515        self.values.is_empty()
516    }
517
518    /// Clear all key-value pairs from the dictionary.
519    pub fn clear(&mut self) {
520        self.values.clear();
521    }
522
523    // Internal methods for config guard support
524    /// Take a value by key name, returning the boxed value if present
525    pub fn remove_value<T: 'static>(&mut self, key: Key<T>) -> Option<Box<dyn SerializableValue>> {
526        self.values.remove(key.name)
527    }
528
529    /// Restore a value by key name
530    pub fn insert_value<T: 'static>(&mut self, key: Key<T>, value: Box<dyn SerializableValue>) {
531        self.values.insert(key.name, value);
532    }
533
534    /// Restore a value by key name
535    pub fn insert_value_by_name_unchecked(
536        &mut self,
537        name: &'static str,
538        value: Box<dyn SerializableValue>,
539    ) {
540        self.values.insert(name, value);
541    }
542
543    /// Internal getter by key name for explicitly-set values (no
544    /// defaults).
545    pub fn get_value_by_name(&self, name: &'static str) -> Option<&dyn SerializableValue> {
546        self.values.get(name).map(|b| b.as_ref())
547    }
548
549    /// Merge all attributes from `other` into this set, consuming
550    /// `other`.
551    ///
552    /// For each key in `other`, moves its value into `self`,
553    /// overwriting any existing value for the same key.
554    pub fn merge(&mut self, other: Attrs) {
555        self.values.extend(other.values);
556    }
557}
558
559impl Clone for Attrs {
560    fn clone(&self) -> Self {
561        let mut values = HashMap::new();
562        for (key, value) in &self.values {
563            values.insert(*key, value.cloned());
564        }
565        Self { values }
566    }
567}
568
569impl std::fmt::Display for Attrs {
570    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
571        let mut first = true;
572        for (key, value) in &self.values {
573            if first {
574                first = false;
575            } else {
576                write!(f, ",")?;
577            }
578            write!(f, "{}={}", key, value.display())?
579        }
580        Ok(())
581    }
582}
583
584impl std::fmt::Debug for Attrs {
585    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
586        // Create a map of key names to their JSON representation for debugging
587        let mut debug_map = std::collections::BTreeMap::new();
588        for (key, value) in &self.values {
589            match serde_json::to_string(value.as_erased_serialize()) {
590                Ok(json) => {
591                    debug_map.insert(*key, json);
592                }
593                Err(_) => {
594                    debug_map.insert(*key, "<serialization error>".to_string());
595                }
596            }
597        }
598
599        f.debug_struct("Attrs").field("values", &debug_map).finish()
600    }
601}
602
603impl Default for Attrs {
604    fn default() -> Self {
605        Self::new()
606    }
607}
608
609impl Serialize for Attrs {
610    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
611    where
612        S: Serializer,
613    {
614        let mut map = serializer.serialize_map(Some(self.values.len()))?;
615
616        for (key_name, value) in &self.values {
617            map.serialize_entry(key_name, value.as_erased_serialize())?;
618        }
619
620        map.end()
621    }
622}
623
624struct AttrsVisitor;
625
626impl<'de> Visitor<'de> for AttrsVisitor {
627    type Value = Attrs;
628
629    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
630        formatter.write_str("a map of attribute keys to their serialized values")
631    }
632
633    fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
634    where
635        M: MapAccess<'de>,
636    {
637        static KEYS_BY_NAME: std::sync::LazyLock<HashMap<&'static str, &'static AttrKeyInfo>> =
638            std::sync::LazyLock::new(|| {
639                inventory::iter::<AttrKeyInfo>()
640                    .map(|info| (info.name, info))
641                    .collect()
642            });
643        let keys_by_name = &*KEYS_BY_NAME;
644
645        let exe_name = std::env::current_exe()
646            .ok()
647            .map(|p| p.display().to_string())
648            .unwrap_or_else(|| "<unknown-exe>".to_string());
649
650        let mut attrs = Attrs::new();
651        while let Some(key_name) = access.next_key::<String>()? {
652            let Some(&key) = keys_by_name.get(key_name.as_str()) else {
653                // We hit an attribute key that this binary doesn't
654                // know about.
655                //
656                // In JSON we'd just deserialize that value into
657                // `IgnoredAny` and move on. With bincode we *can't*
658                // do that safely:
659                //
660                // - We don't know this key's value type.
661                // - That means we don't know how many bytes to
662                //   consume for the value.
663                // - If we guess wrong or try `IgnoredAny`, bincode
664                //   would need `Deserializer::deserialize_any()` to
665                //   skip it, but bincode refuses because it can't
666                //   know how many bytes to advance.
667                //
668                // Result: we cannot safely "skip" the unknown value
669                // without risking desync of the remaining stream. So
670                // we abort here and surface which key caused it, and
671                // the caller must strip it before sending.
672                access.next_value::<serde::de::IgnoredAny>().map_err(|_| {
673                    serde::de::Error::custom(format!(
674                        "unknown attr key '{}' on binary '{}'; \
675                         this binary doesn't know this key and cannot skip its value safely under bincode",
676                        key_name, exe_name,
677                    ))
678                })?;
679                continue;
680            };
681
682            // Create a seed to deserialize the value using erased_serde
683            let seed = ValueDeserializeSeed {
684                deserialize_erased: key.deserialize_erased,
685            };
686            match access.next_value_seed(seed) {
687                Ok(value) => {
688                    attrs.values.insert(key.name, value);
689                }
690                Err(err) => {
691                    return Err(serde::de::Error::custom(format!(
692                        "failed to deserialize value for key {}: {}",
693                        key_name, err
694                    )));
695                }
696            }
697        }
698
699        Ok(attrs)
700    }
701}
702
703/// Helper struct to deserialize values using erased_serde
704struct ValueDeserializeSeed {
705    deserialize_erased:
706        fn(&mut dyn ErasedDeserializer) -> Result<Box<dyn SerializableValue>, erased_serde::Error>,
707}
708
709impl<'de> serde::de::DeserializeSeed<'de> for ValueDeserializeSeed {
710    type Value = Box<dyn SerializableValue>;
711
712    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
713    where
714        D: serde::de::Deserializer<'de>,
715    {
716        let mut erased = <dyn erased_serde::Deserializer>::erase(deserializer);
717        (self.deserialize_erased)(&mut erased).map_err(serde::de::Error::custom)
718    }
719}
720
721impl<'de> Deserialize<'de> for Attrs {
722    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
723    where
724        D: Deserializer<'de>,
725    {
726        deserializer.deserialize_map(AttrsVisitor)
727    }
728}
729
730// Converts an ASCII string to lowercase at compile time.
731// Returns a const string with lowercase ASCII characters.
732#[doc(hidden)]
733pub const fn ascii_to_lowercase_const<const N: usize>(input: &str) -> [u8; N] {
734    let bytes = input.as_bytes();
735    let mut result = [0u8; N];
736    let mut i = 0;
737
738    while i < bytes.len() && i < N {
739        let byte = bytes[i];
740        if byte >= b'A' && byte <= b'Z' {
741            result[i] = byte + 32; // Convert to lowercase
742        } else {
743            result[i] = byte;
744        }
745        i += 1;
746    }
747
748    result
749}
750
751// Macro to generate a const lowercase string at compile time
752#[doc(hidden)]
753#[macro_export]
754macro_rules! const_ascii_lowercase {
755    ($s:expr) => {{
756        const INPUT: &str = $s;
757        const LEN: usize = INPUT.len();
758        const BYTES: [u8; LEN] = $crate::attrs::ascii_to_lowercase_const::<LEN>(INPUT);
759        // Safety: We're converting ASCII to ASCII, so it's valid UTF-8
760        unsafe { std::str::from_utf8_unchecked(&BYTES) }
761    }};
762}
763
764/// Macro to check check that a trait is implemented, generating a
765/// nice error message if it isn't.
766#[doc(hidden)]
767#[macro_export]
768macro_rules! assert_impl {
769    ($ty:ty, $trait:path) => {
770        const _: fn() = || {
771            fn check<T: $trait>() {}
772            check::<$ty>();
773        };
774    };
775}
776
777/// Declares attribute keys using a lazy_static! style syntax.
778///
779/// # Syntax
780///
781/// ```ignore
782/// declare_attrs! {
783///     /// Documentation for the key (default visibility).
784///     attr KEY_NAME: Type = default_value;
785///
786///     /// Another key (default value is optional)
787///     pub attr ANOTHER_KEY: AnotherType;
788/// }
789/// ```
790///
791/// # Arguments
792///
793/// * Optional visibility modifier (`pub`, `pub(crate)`, etc.)
794/// * `attr` keyword (required)
795/// * Key name (identifier)
796/// * Type of values this key can store
797/// * Optional default value
798///
799/// # Example
800///
801/// ```
802/// use std::time::Duration;
803///
804/// use hyperactor_config::attrs::Attrs;
805/// use hyperactor_config::attrs::declare_attrs;
806///
807/// declare_attrs! {
808///     /// Timeout for RPC operations
809///     pub attr TIMEOUT: Duration = Duration::from_secs(30);
810///
811///     /// Maximum number of retry attempts (no default specified)
812///     attr MAX_RETRIES: u32;
813/// }
814///
815/// let mut attrs = Attrs::new();
816/// assert_eq!(attrs.get(TIMEOUT), Some(&Duration::from_secs(30)));
817/// attrs.set(MAX_RETRIES, 5);
818/// ```
819#[macro_export]
820macro_rules! declare_attrs {
821    // Handle multiple attribute keys with optional default values and optional meta attributes
822    ($(
823        $(#[$attr:meta])*
824        $(@meta($($meta_key:ident = $meta_value:expr),* $(,)?))*
825        $vis:vis attr $name:ident: $type:ty $(= $default:expr)?;
826    )*) => {
827        $(
828            $crate::declare_attrs! {
829                @single
830                $(@meta($($meta_key = $meta_value),*))*
831                $(#[$attr])* ;
832                $vis attr $name: $type $(= $default)?;
833            }
834        )*
835    };
836
837    // Handle single attribute key with default value and meta attributes
838    (@single $(@meta($($meta_key:ident = $meta_value:expr),* $(,)?))* $(#[$attr:meta])* ; $vis:vis attr $name:ident: $type:ty = $default:expr;) => {
839        $crate::assert_impl!($type, $crate::attrs::AttrValue);
840
841        // Create a static default value
842        $crate::paste! {
843            static [<$name _DEFAULT>]: $type = $default;
844            static [<$name _META_ATTRS>]: std::sync::LazyLock<$crate::attrs::Attrs> =
845                std::sync::LazyLock::new(|| {
846                    #[allow(unused_mut)]
847                    let mut attrs = $crate::attrs::Attrs::new();
848                    $($(
849                        attrs.set($meta_key, $meta_value);
850                    )*)*
851                    attrs
852                });
853        }
854
855        $(#[$attr])*
856        $vis static $name: $crate::attrs::Key<$type> = {
857            $crate::assert_impl!($type, $crate::attrs::AttrValue);
858
859            const FULL_NAME: &str = concat!(std::module_path!(), "::", stringify!($name));
860            const LOWER_NAME: &str = $crate::const_ascii_lowercase!(FULL_NAME);
861            $crate::paste! {
862                $crate::attrs::Key::new(
863                    LOWER_NAME,
864                    Some(&[<$name _DEFAULT>]),
865                    $crate::paste! { &[<$name _META_ATTRS>] },
866                )
867            }
868        };
869
870        // Register the key for serialization
871        $crate::submit! {
872            $crate::attrs::AttrKeyInfo {
873                name: {
874                    const FULL_NAME: &str = concat!(std::module_path!(), "::", stringify!($name));
875                    $crate::const_ascii_lowercase!(FULL_NAME)
876                },
877                typehash: <$type as $crate::typeuri::Named>::typehash,
878                deserialize_erased: |deserializer| {
879                    let value: $type = erased_serde::deserialize(deserializer)?;
880                    Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
881                },
882                meta: $crate::paste! { &[<$name _META_ATTRS>] },
883                display: |value: &dyn $crate::attrs::SerializableValue| {
884                    let value = value.as_any().downcast_ref::<$type>().unwrap();
885                    $crate::attrs::AttrValue::display(value)
886                },
887                parse: |value: &str| {
888                    let value: $type = $crate::attrs::AttrValue::parse(value)?;
889                    Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
890                },
891                default: Some($crate::paste! { &[<$name _DEFAULT>] }),
892                erased: &$name,
893            }
894        }
895    };
896
897    // Handle single attribute key without default value but with meta attributes
898    (@single $(@meta($($meta_key:ident = $meta_value:expr),* $(,)?))* $(#[$attr:meta])* ; $vis:vis attr $name:ident: $type:ty;) => {
899        $crate::assert_impl!($type, $crate::attrs::AttrValue);
900
901        $crate::paste! {
902            static [<$name _META_ATTRS>]: std::sync::LazyLock<$crate::attrs::Attrs> =
903            std::sync::LazyLock::new(|| {
904                #[allow(unused_mut)]
905                let mut attrs = $crate::attrs::Attrs::new();
906                $($(
907                    // Note: This assumes meta keys are already declared somewhere
908                    // The user needs to ensure the meta keys exist and are in scope
909                    attrs.set($meta_key, $meta_value);
910                )*)*
911                attrs
912            });
913        }
914
915        $(#[$attr])*
916        $vis static $name: $crate::attrs::Key<$type> = {
917            const FULL_NAME: &str = concat!(std::module_path!(), "::", stringify!($name));
918            const LOWER_NAME: &str = $crate::const_ascii_lowercase!(FULL_NAME);
919            $crate::attrs::Key::new(LOWER_NAME, None, $crate::paste! { &[<$name _META_ATTRS>] })
920        };
921
922
923        // Register the key for serialization
924        $crate::submit! {
925            $crate::attrs::AttrKeyInfo {
926                name: {
927                    const FULL_NAME: &str = concat!(std::module_path!(), "::", stringify!($name));
928                    $crate::const_ascii_lowercase!(FULL_NAME)
929                },
930                typehash: <$type as $crate::typeuri::Named>::typehash,
931                deserialize_erased: |deserializer| {
932                    let value: $type = erased_serde::deserialize(deserializer)?;
933                    Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
934                },
935                meta: $crate::paste! { &[<$name _META_ATTRS>] },
936                display: |value: &dyn $crate::attrs::SerializableValue| {
937                    let value = value.as_any().downcast_ref::<$type>().unwrap();
938                    $crate::attrs::AttrValue::display(value)
939                },
940                parse: |value: &str| {
941                    let value: $type = $crate::attrs::AttrValue::parse(value)?;
942                    Ok(Box::new(value) as Box<dyn $crate::attrs::SerializableValue>)
943                },
944                default: None,
945                erased: &$name,
946            }
947        }
948    };
949}
950
951pub use declare_attrs;
952
953#[cfg(test)]
954mod tests {
955    use std::time::Duration;
956
957    use super::*;
958
959    declare_attrs! {
960        attr TEST_TIMEOUT: Duration;
961        attr TEST_COUNT: u32;
962        @meta(TEST_COUNT = 42)
963        pub attr TEST_NAME: String;
964    }
965
966    #[test]
967    fn test_basic_operations() {
968        let mut attrs = Attrs::new();
969
970        // Test setting and getting values
971        attrs.set(TEST_TIMEOUT, Duration::from_secs(5));
972        attrs.set(TEST_COUNT, 42u32);
973        attrs.set(TEST_NAME, "test".to_string());
974
975        assert_eq!(attrs.get(TEST_TIMEOUT), Some(&Duration::from_secs(5)));
976        assert_eq!(attrs.get(TEST_COUNT), Some(&42u32));
977        assert_eq!(attrs.get(TEST_NAME), Some(&"test".to_string()));
978
979        // Test contains_key
980        assert!(attrs.contains_key(TEST_TIMEOUT));
981        assert!(attrs.contains_key(TEST_COUNT));
982        assert!(attrs.contains_key(TEST_NAME));
983
984        // Test len
985        assert_eq!(attrs.len(), 3);
986        assert!(!attrs.is_empty());
987
988        // Meta attribute:
989        assert_eq!(TEST_NAME.attrs().get(TEST_COUNT).unwrap(), &42u32);
990    }
991
992    #[test]
993    fn test_get_mut() {
994        let mut attrs = Attrs::new();
995        attrs.set(TEST_COUNT, 10u32);
996
997        if let Some(count) = attrs.get_mut(TEST_COUNT) {
998            *count += 5;
999        }
1000
1001        assert_eq!(attrs.get(TEST_COUNT), Some(&15u32));
1002    }
1003
1004    #[test]
1005    fn test_remove() {
1006        let mut attrs = Attrs::new();
1007        attrs.set(TEST_COUNT, 42u32);
1008
1009        let removed = attrs.remove(TEST_COUNT);
1010        assert!(removed);
1011        assert_eq!(attrs.get(TEST_COUNT), None);
1012        assert!(!attrs.contains_key(TEST_COUNT));
1013    }
1014
1015    #[test]
1016    fn test_clear() {
1017        let mut attrs = Attrs::new();
1018        attrs.set(TEST_TIMEOUT, Duration::from_secs(1));
1019        attrs.set(TEST_COUNT, 42u32);
1020
1021        attrs.clear();
1022        assert!(attrs.is_empty());
1023        assert_eq!(attrs.len(), 0);
1024    }
1025
1026    #[test]
1027    fn test_key_properties() {
1028        assert_eq!(
1029            TEST_TIMEOUT.name(),
1030            "hyperactor_config::attrs::tests::test_timeout"
1031        );
1032    }
1033
1034    #[test]
1035    fn test_serialization() {
1036        let mut attrs = Attrs::new();
1037        attrs.set(TEST_TIMEOUT, Duration::from_secs(5));
1038        attrs.set(TEST_COUNT, 42u32);
1039        attrs.set(TEST_NAME, "test".to_string());
1040
1041        // Test serialization
1042        let serialized = serde_json::to_string(&attrs).expect("Failed to serialize");
1043
1044        // The serialized string should contain the key names and their values
1045        assert!(serialized.contains("hyperactor_config::attrs::tests::test_timeout"));
1046        assert!(serialized.contains("hyperactor_config::attrs::tests::test_count"));
1047        assert!(serialized.contains("hyperactor_config::attrs::tests::test_name"));
1048    }
1049
1050    #[test]
1051    fn test_deserialization() {
1052        // Create original attrs
1053        let mut original_attrs = Attrs::new();
1054        original_attrs.set(TEST_TIMEOUT, Duration::from_secs(5));
1055        original_attrs.set(TEST_COUNT, 42u32);
1056        original_attrs.set(TEST_NAME, "test".to_string());
1057
1058        // Serialize
1059        let serialized = serde_json::to_string(&original_attrs).expect("Failed to serialize");
1060
1061        // Deserialize (no manual registry needed!)
1062        let deserialized_attrs: Attrs =
1063            serde_json::from_str(&serialized).expect("Failed to deserialize");
1064
1065        // Verify the deserialized values
1066        assert_eq!(
1067            deserialized_attrs.get(TEST_TIMEOUT),
1068            Some(&Duration::from_secs(5))
1069        );
1070        assert_eq!(deserialized_attrs.get(TEST_COUNT), Some(&42u32));
1071        assert_eq!(deserialized_attrs.get(TEST_NAME), Some(&"test".to_string()));
1072    }
1073
1074    #[test]
1075    fn test_roundtrip_serialization() {
1076        // Create original attrs
1077        let mut original = Attrs::new();
1078        original.set(TEST_TIMEOUT, Duration::from_secs(10));
1079        original.set(TEST_COUNT, 5u32);
1080        original.set(TEST_NAME, "test-service".to_string());
1081
1082        // Serialize
1083        let serialized = serde_json::to_string(&original).unwrap();
1084
1085        // Deserialize
1086        let deserialized: Attrs = serde_json::from_str(&serialized).unwrap();
1087
1088        // Verify round-trip worked
1089        assert_eq!(
1090            deserialized.get(TEST_TIMEOUT),
1091            Some(&Duration::from_secs(10))
1092        );
1093        assert_eq!(deserialized.get(TEST_COUNT), Some(&5u32));
1094        assert_eq!(
1095            deserialized.get(TEST_NAME),
1096            Some(&"test-service".to_string())
1097        );
1098    }
1099
1100    #[test]
1101    fn test_empty_attrs_serialization() {
1102        let attrs = Attrs::new();
1103        let serialized = serde_json::to_string(&attrs).unwrap();
1104
1105        // Empty attrs should serialize to empty JSON object
1106        assert_eq!(serialized, "{}");
1107
1108        let deserialized: Attrs = serde_json::from_str(&serialized).unwrap();
1109
1110        assert!(deserialized.is_empty());
1111    }
1112
1113    #[test]
1114    fn test_format_independence() {
1115        // Test that proves we're using the serializer directly, not JSON internally
1116        let mut attrs = Attrs::new();
1117        attrs.set(TEST_COUNT, 42u32);
1118        attrs.set(TEST_NAME, "test".to_string());
1119
1120        // Serialize to different formats
1121        let json_output = serde_json::to_string(&attrs).unwrap();
1122        let yaml_output = serde_yaml::to_string(&attrs).unwrap();
1123
1124        // JSON should have colons and quotes
1125        assert!(json_output.contains(":"));
1126        assert!(json_output.contains("\""));
1127
1128        // JSON should serialize numbers as numbers, not strings
1129        assert!(json_output.contains("42"));
1130        assert!(!json_output.contains("\"42\""));
1131
1132        // YAML should have colons but different formatting
1133        assert!(yaml_output.contains(":"));
1134        assert!(yaml_output.contains("42"));
1135
1136        // YAML shouldn't quote simple strings or numbers
1137        assert!(!yaml_output.contains("\"42\""));
1138
1139        // The outputs should be different (proving different serializers were used)
1140        assert_ne!(json_output, yaml_output);
1141
1142        // Verify that both can be deserialized correctly
1143        let from_json: Attrs = serde_json::from_str(&json_output).unwrap();
1144        let from_yaml: Attrs = serde_yaml::from_str(&yaml_output).unwrap();
1145
1146        assert_eq!(from_json.get(TEST_COUNT), Some(&42u32));
1147        assert_eq!(from_yaml.get(TEST_COUNT), Some(&42u32));
1148        assert_eq!(from_json.get(TEST_NAME), Some(&"test".to_string()));
1149        assert_eq!(from_yaml.get(TEST_NAME), Some(&"test".to_string()));
1150    }
1151
1152    #[test]
1153    fn test_clone() {
1154        // Create original attrs with multiple types
1155        let mut original = Attrs::new();
1156        original.set(TEST_COUNT, 42u32);
1157        original.set(TEST_NAME, "test".to_string());
1158        original.set(TEST_TIMEOUT, std::time::Duration::from_secs(10));
1159
1160        // Clone the attrs
1161        let cloned = original.clone();
1162
1163        // Verify that the clone has the same values
1164        assert_eq!(cloned.get(TEST_COUNT), Some(&42u32));
1165        assert_eq!(cloned.get(TEST_NAME), Some(&"test".to_string()));
1166        assert_eq!(
1167            cloned.get(TEST_TIMEOUT),
1168            Some(&std::time::Duration::from_secs(10))
1169        );
1170
1171        // Verify that modifications to the original don't affect the clone
1172        original.set(TEST_COUNT, 100u32);
1173        assert_eq!(original.get(TEST_COUNT), Some(&100u32));
1174        assert_eq!(cloned.get(TEST_COUNT), Some(&42u32)); // Clone should be unchanged
1175
1176        // Verify that modifications to the clone don't affect the original
1177        let mut cloned_mut = cloned.clone();
1178        cloned_mut.set(TEST_NAME, "modified".to_string());
1179        assert_eq!(cloned_mut.get(TEST_NAME), Some(&"modified".to_string()));
1180        assert_eq!(original.get(TEST_NAME), Some(&"test".to_string())); // Original should be unchanged
1181    }
1182
1183    #[test]
1184    fn test_debug_with_json() {
1185        let mut attrs = Attrs::new();
1186        attrs.set(TEST_COUNT, 42u32);
1187        attrs.set(TEST_NAME, "test".to_string());
1188
1189        // Test that Debug implementation works and contains JSON representations
1190        let debug_output = format!("{:?}", attrs);
1191
1192        // Should contain the struct name
1193        assert!(debug_output.contains("Attrs"));
1194
1195        // Should contain JSON representations of the values
1196        assert!(debug_output.contains("42"));
1197
1198        // Should contain the key names
1199        assert!(debug_output.contains("hyperactor_config::attrs::tests::test_count"));
1200        assert!(debug_output.contains("hyperactor_config::attrs::tests::test_name"));
1201
1202        // For strings, the JSON representation should be the escaped version
1203        // Let's check that the test string is actually present in some form
1204        assert!(debug_output.contains("test"));
1205    }
1206
1207    declare_attrs! {
1208        /// With default...
1209        attr TIMEOUT_WITH_DEFAULT: Duration = Duration::from_secs(10);
1210
1211        /// Just to ensure visibilty is parsed.
1212        pub(crate) attr CRATE_LOCAL_ATTR: String;
1213    }
1214
1215    #[test]
1216    fn test_defaults() {
1217        assert!(TIMEOUT_WITH_DEFAULT.has_default());
1218        assert!(!CRATE_LOCAL_ATTR.has_default());
1219
1220        assert_eq!(
1221            Attrs::new().get(TIMEOUT_WITH_DEFAULT),
1222            Some(&Duration::from_secs(10))
1223        );
1224    }
1225
1226    #[test]
1227    fn test_indexing() {
1228        let mut attrs = Attrs::new();
1229
1230        assert_eq!(attrs[TIMEOUT_WITH_DEFAULT], Duration::from_secs(10));
1231        attrs[TIMEOUT_WITH_DEFAULT] = Duration::from_secs(100);
1232        assert_eq!(attrs[TIMEOUT_WITH_DEFAULT], Duration::from_secs(100));
1233
1234        attrs.set(CRATE_LOCAL_ATTR, "test".to_string());
1235        assert_eq!(attrs[CRATE_LOCAL_ATTR], "test".to_string());
1236    }
1237
1238    #[test]
1239    fn attrs_deserialize_unknown_key_is_error() {
1240        // Build a real Attrs, but inject a key that this binary does
1241        // NOT know about. We do that with
1242        // insert_value_by_name_unchecked(), which bypasses the
1243        // declare_attrs!/inventory registration path.
1244        //
1245        // Then:
1246        //   1. Serialize that Attrs with bincode (the real wire
1247        //      format).
1248        //   2. Attempt to bincode-deserialize those bytes back into
1249        //      Attrs.
1250        //
1251        // During deserialization, AttrsVisitor::visit_map() will:
1252        //   - read the key string
1253        //   - fail to find it in KEYS_BY_NAME (the compiled-in
1254        //     registry)
1255        //   - immediately error instead of trying to skip the value,
1256        //     because with bincode we can't safely consume an unknown
1257        //     typed value without risking stream desync.
1258        //
1259        // This reproduces exactly what happens when a parent proc
1260        // sends an Attrs containing a key the child binary wasn't
1261        // built with.
1262
1263        // Definitely not declared in this crate's inventory:
1264        let bad_key: &'static str = "monarch_hyperactor::pytokio::unawaited_pytokio_traceback";
1265
1266        // Make an Attrs that pretends to have that key. u32 already
1267        // implements AttrValue -> SerializableValue, so we can just
1268        // box a 0u32.
1269        let mut attrs = Attrs::new();
1270        attrs.insert_value_by_name_unchecked(bad_key, Box::new(0u32));
1271
1272        // Serialize this Attrs using bincode (same codec we use on
1273        // the wire).
1274        let wire_bytes = bincode::serialize(&attrs).unwrap();
1275
1276        // Now try to decode those bytes back into Attrs. This should
1277        // hit the unknown-key branch and return Err.
1278        let err = bincode::deserialize::<Attrs>(&wire_bytes)
1279            .expect_err("should error on unknown attr key");
1280
1281        let exe_str = std::env::current_exe()
1282            .ok()
1283            .map(|p| p.display().to_string())
1284            .unwrap_or_else(|| "<unknown-exe>".to_string());
1285        let msg = format!("{err}");
1286
1287        assert!(msg.contains("unknown attr key"), "got: {msg}");
1288        assert!(msg.contains(bad_key), "got: {msg}");
1289        assert!(msg.contains(&exe_str), "got: {msg}");
1290    }
1291}
hyperactor_config/attrs.rs

hyperactor_config/
attrs.rs
