monarch_hyperactor/
operation_context.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//! Operation-context carrier for the Python endpoint machinery.
10//!
11//! Stamps `OPERATION_ENDPOINT` / `OPERATION_ADVERB` on outgoing
12//! requests so failure surfaces (mailbox log, undeliverable
13//! formatter) can name the operation.
14//!
15//! ## Operation-context invariants (OC-*)
16//!
17//! - **OC-1 (request-originated).** Stamped at the caller's
18//!   request-send site. Only the caller knows the qualified
19//!   endpoint and adverb.
20//!
21//! - **OC-2 (filtered carry).** Only keys marked
22//!   `@meta(OPERATION_CONTEXT_HEADER = true)` ride this path.
23//!   Adding another `OPERATION_*` key joins the vocabulary by
24//!   declaration alone.
25//!
26//! - **OC-3 (carry continuity).** Captured headers are preserved
27//!   onto the reply envelope at the `Port` capture site via
28//!   `send_with_headers`.
29
30use hyperactor::mailbox::headers::OPERATION_ADVERB;
31use hyperactor::mailbox::headers::OPERATION_ENDPOINT;
32use hyperactor_config::Attrs;
33use hyperactor_config::Flattrs;
34use hyperactor_config::attrs::OPERATION_CONTEXT_HEADER;
35use hyperactor_config::attrs::stamp_marked_attrs_into_flattrs;
36
37/// Build an `Attrs` carrier populated with the operation-context keys
38/// an endpoint-send site wants to stamp onto its outgoing request.
39///
40/// Callers supply the values they know in scope at the send site
41/// (qualified endpoint name, adverb token). Any `None` is simply
42/// omitted — the marker-driven stamp filters on presence.
43pub fn build_operation_context_attrs(
44    endpoint: Option<String>,
45    adverb: Option<&'static str>,
46) -> Attrs {
47    let mut attrs = Attrs::new();
48    if let Some(ep) = endpoint {
49        attrs.set(OPERATION_ENDPOINT, ep);
50    }
51    if let Some(a) = adverb {
52        attrs.set(OPERATION_ADVERB, a.to_string());
53    }
54    attrs
55}
56
57/// Stamp the operation-context subset of `attrs` onto `headers` using
58/// the shared marker-driven mechanism. Only entries whose declared
59/// key carries `@meta(OPERATION_CONTEXT_HEADER = true)` are written;
60/// anything else in `attrs` is silently skipped (OC-2).
61pub fn stamp_operation_context(headers: &mut Flattrs, attrs: &Attrs) {
62    stamp_marked_attrs_into_flattrs(headers, attrs, OPERATION_CONTEXT_HEADER);
63}
64
65#[cfg(test)]
66mod tests {
67    use hyperactor_config::attrs::copy_marked_flattrs;
68
69    use super::*;
70
71    /// OC-2: the shared marker-driven copy propagates exactly the
72    /// `OPERATION_CONTEXT_HEADER`-marked keys from the source headers
73    /// onto the destination. Unrelated declared `Flattrs` entries
74    /// (e.g. `SEND_TIMESTAMP`) are excluded even though they're
75    /// present on the source.
76    #[test]
77    fn test_oc2_filter_rejects_unmarked_keys() {
78        let mut src = Flattrs::new();
79        src.set(OPERATION_ENDPOINT, "training.buffer.sample()".to_string());
80        src.set(
81            hyperactor::mailbox::headers::SEND_TIMESTAMP,
82            std::time::SystemTime::UNIX_EPOCH,
83        );
84
85        let mut dst = Flattrs::new();
86        copy_marked_flattrs(&mut dst, &src, OPERATION_CONTEXT_HEADER);
87
88        assert_eq!(
89            dst.get(OPERATION_ENDPOINT),
90            Some("training.buffer.sample()".to_string()),
91            "OC-2: marked key must cross the carry"
92        );
93        assert!(
94            dst.get(hyperactor::mailbox::headers::SEND_TIMESTAMP)
95                .is_none(),
96            "OC-2: unmarked key must not cross the carry"
97        );
98    }
99
100    /// OC-1: the caller-side helper assembles the operation-context
101    /// attrs from values known at the request-send site, rather
102    /// than synthesized at the reply site. `build_operation_context_attrs`
103    /// is the vehicle for that origination step.
104    #[test]
105    fn test_oc1_build_attrs_populates_supplied_fields() {
106        let attrs = build_operation_context_attrs(
107            Some("training.buffer.sample()".to_string()),
108            Some("call_one"),
109        );
110        assert_eq!(
111            attrs.get(OPERATION_ENDPOINT),
112            Some(&"training.buffer.sample()".to_string())
113        );
114        assert_eq!(attrs.get(OPERATION_ADVERB), Some(&"call_one".to_string()));
115    }
116
117    /// OC-1: fields the request site cannot supply must not be
118    /// fabricated downstream. `build_operation_context_attrs` omits
119    /// them when the caller passes `None`.
120    #[test]
121    fn test_oc1_build_attrs_omits_none_fields() {
122        let attrs =
123            build_operation_context_attrs(Some("training.buffer.sample()".to_string()), None);
124        assert_eq!(
125            attrs.get(OPERATION_ENDPOINT),
126            Some(&"training.buffer.sample()".to_string())
127        );
128        assert!(attrs.get(OPERATION_ADVERB).is_none());
129    }
130
131    /// OC-2: the stamp helper writes only the
132    /// `OPERATION_CONTEXT_HEADER`-marked keys onto the outgoing
133    /// headers. Same mechanism as the OC-2 copy test above, on the
134    /// `Attrs -> Flattrs` direction.
135    #[test]
136    fn test_oc2_stamp_writes_marked_keys_to_headers() {
137        let attrs = build_operation_context_attrs(
138            Some("training.buffer.sample()".to_string()),
139            Some("call_one"),
140        );
141        let mut headers = Flattrs::new();
142        stamp_operation_context(&mut headers, &attrs);
143
144        assert_eq!(
145            headers.get(OPERATION_ENDPOINT),
146            Some("training.buffer.sample()".to_string())
147        );
148        assert_eq!(headers.get(OPERATION_ADVERB), Some("call_one".to_string()));
149    }
150}
monarch_hyperactor/operation_context.rs

monarch_hyperactor/
operation_context.rs
