Skip to main content

hyperactor/
ordering.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//! This module contains utilities used to help messages are delivered in order
10//! for any given sender and receiver actor pair.
11
12use std::any::TypeId;
13use std::fmt;
14use std::ops::Range;
15use std::sync::Arc;
16use std::sync::Mutex;
17
18use hyperactor_config::AttrValue;
19use hyperactor_config::attrs::declare_attrs;
20use rustc_hash::FxHashMap;
21use serde::Deserialize;
22use serde::Serialize;
23use typeuri::Named;
24use uuid::Uuid;
25
26use crate::ActorAddr;
27use crate::PortAddr;
28
29// Control ports are delivered through runtime-owned channels rather than the
30// actor's work queue, so they must not share the per-actor handler-port
31// sequence counter.
32
33/// Returns true if `id` is the `TypeId` of a bypass-channel message type
34/// (i.e. one that must not be wired through `Ports::get`).
35pub(crate) fn is_bypass_workq_type_id(id: TypeId) -> bool {
36    id == TypeId::of::<crate::introspect::IntrospectMessage>()
37        || id == TypeId::of::<crate::proc::StatusMessage>()
38}
39
40/// Per-session receiver-local ordering snapshot. This does not lock the
41/// work queue or hold sequencing state across awaits.
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Named, AttrValue)]
43pub struct OrderingSessionSnapshot {
44    /// Identifier of the session whose ordering state this snapshot
45    /// describes; matches `SeqInfo::Session { session_id, .. }` on the
46    /// wire.
47    pub session_id: Uuid,
48
49    /// Sender `ActorAddr` captured by the first-non-None-wins
50    /// rule. `None` when every message in this session bypassed the
51    /// normal stamping sites (rare: test fixtures, non-handler routes).
52    pub sender: Option<ActorAddr>,
53
54    /// Highest seq released from the reorder buffer into the actor work
55    /// queue. NOT "handler processed" -- that's downstream of the queue.
56    pub last_released_seq: u64,
57
58    /// `last_released_seq + 1` -- the seq the next contiguous send must
59    /// carry for delivery without further buffering.
60    pub expected_next_seq: u64,
61
62    /// Number of messages currently sitting in the reorder buffer waiting
63    /// for a seq gap to be filled. Zero on healthy in-order sessions.
64    pub buffered_count: usize,
65
66    /// Lowest seq currently buffered. `None` when `buffered_count == 0`.
67    pub oldest_buffered_seq: Option<u64>,
68
69    /// Highest seq currently buffered. `None` when `buffered_count == 0`.
70    pub newest_buffered_seq: Option<u64>,
71}
72
73/// Diagnostic projection of receiver-local ordering state.
74///
75/// Each `OrderingSessionSnapshot` corresponds to one sender session keyed
76/// by `SeqInfo::Session.session_id`. Sequence progress here means
77/// "released into the actor work queue", not "processed by the actor handler".
78///
79/// Carries completeness metadata so callers can distinguish "no stalled
80/// sessions" from "snapshot was partial" from "buffering is disabled".
81#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Named, AttrValue)]
82pub struct OrderingSnapshot {
83    /// Whether reorder buffering is enabled. When `false`, messages
84    /// bypass sequencing, and `sessions` is empty even under load.
85    pub enabled: bool,
86
87    /// Per-session entries, sorted by `session_id` for stable output.
88    /// Includes idle (drained) sessions with `buffered_count == 0`.
89    pub sessions: Vec<OrderingSessionSnapshot>,
90
91    /// Sessions that could not be observed because sequencing state was
92    /// busy when we tried to snapshot. NOT in `sessions`. Zero means
93    /// the snapshot is complete.
94    pub skipped_session_count: usize,
95}
96
97impl OrderingSnapshot {
98    /// True when no session was skipped, i.e., every live session's
99    /// state was successfully captured in `sessions`.
100    pub fn is_complete(&self) -> bool {
101        self.skipped_session_count == 0
102    }
103
104    #[cfg(test)]
105    fn delivery_progress_from(&self, sender: &ActorAddr) -> Option<DeliveryProgress> {
106        if !self.enabled || !self.is_complete() {
107            return None;
108        }
109
110        let largest_dequeueable_sequence = self
111            .sessions
112            .iter()
113            .filter(|session| session.sender.as_ref() == Some(sender))
114            .map(|session| session.last_released_seq)
115            .max()
116            .unwrap_or_default();
117
118        Some(DeliveryProgress {
119            largest_dequeueable_sequence,
120        })
121    }
122
123    /// Delivery progress for one sender sequencing session.
124    pub fn delivery_progress_for_session(&self, session_id: Uuid) -> Option<DeliveryProgress> {
125        if !self.enabled || !self.is_complete() {
126            return None;
127        }
128
129        let largest_dequeueable_sequence = self
130            .sessions
131            .iter()
132            .find(|session| session.session_id == session_id)
133            .map(|session| session.last_released_seq)
134            .unwrap_or_default();
135
136        Some(DeliveryProgress {
137            largest_dequeueable_sequence,
138        })
139    }
140}
141
142/// Receiver-side delivery progress for a sender.
143#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Named)]
144pub struct DeliveryProgress {
145    /// Largest sequence released from receiver-side ordering into the
146    /// actor's work queue.
147    pub largest_dequeueable_sequence: u64,
148}
149
150impl fmt::Display for OrderingSessionSnapshot {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        // Infallible for this struct (all fields serialize cleanly).
153        write!(f, "{}", serde_json::to_string(self).unwrap())
154    }
155}
156
157impl std::str::FromStr for OrderingSessionSnapshot {
158    type Err = serde_json::Error;
159    fn from_str(s: &str) -> Result<Self, Self::Err> {
160        serde_json::from_str(s)
161    }
162}
163
164impl fmt::Display for OrderingSnapshot {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        write!(f, "{}", serde_json::to_string(self).unwrap())
167    }
168}
169
170impl std::str::FromStr for OrderingSnapshot {
171    type Err = serde_json::Error;
172    fn from_str(s: &str) -> Result<Self, Self::Err> {
173        serde_json::from_str(s)
174    }
175}
176
177/// Key for sequence assignment.
178/// Handler ports share a sequence per actor; ephemeral ports get individual
179/// sequences. Control ports are direct.
180#[derive(Clone, Debug, Hash, PartialEq, Eq)]
181pub struct SeqKey(SeqKeyKind);
182
183#[derive(Clone, Debug, Hash, PartialEq, Eq)]
184enum SeqKeyKind {
185    Actor(ActorAddr),
186    Port(PortAddr),
187}
188
189impl SeqKey {
190    /// Return the sequence key shared by an actor's handler ports.
191    pub fn for_handler(actor: &ActorAddr) -> Self {
192        Self(SeqKeyKind::Actor(actor.clone()))
193    }
194
195    /// Return the sequence key for an ephemeral port.
196    pub fn for_ephemeral_port(port: &PortAddr) -> Option<Self> {
197        port.port()
198            .is_ephemeral()
199            .then(|| Self(SeqKeyKind::Port(port.clone())))
200    }
201}
202
203/// A message's sequencer number infomation.
204#[derive(Debug, Serialize, Deserialize, Clone, Named, AttrValue, PartialEq)]
205pub enum SeqInfo {
206    /// Messages with the same session ID should be delivered in order.
207    Session {
208        /// Message's session ID
209        session_id: Uuid,
210        /// Message's sequence number in the given session.
211        seq: u64,
212    },
213    /// This message does not have a seq number and should be delivered
214    /// immediately.
215    Direct,
216}
217
218impl SeqInfo {
219    /// Whether this sequencing metadata is valid for receive-side ordering.
220    pub fn is_valid(&self) -> bool {
221        !matches!(self, Self::Session { seq: 0, .. })
222    }
223}
224
225impl fmt::Display for SeqInfo {
226    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227        match self {
228            Self::Direct => write!(f, "direct"),
229            Self::Session { session_id, seq } => write!(f, "{}:{}", session_id, seq),
230        }
231    }
232}
233
234impl std::str::FromStr for SeqInfo {
235    type Err = anyhow::Error;
236
237    fn from_str(s: &str) -> Result<Self, Self::Err> {
238        if s == "direct" {
239            return Ok(SeqInfo::Direct);
240        }
241
242        let parts: Vec<_> = s.split(':').collect();
243        if parts.len() != 2 {
244            return Err(anyhow::anyhow!("invalid SeqInfo: {}", s));
245        }
246        let session_id: Uuid = parts[0].parse()?;
247        let seq: u64 = parts[1].parse()?;
248        Ok(SeqInfo::Session { session_id, seq })
249    }
250}
251
252declare_attrs! {
253    /// The sender of this message, the session ID, and the message's sequence
254    /// number assigned by this session.
255    pub attr SEQ_INFO: SeqInfo;
256}
257
258/// Used by sender to track the message sequence numbers it sends to each destination.
259/// Each [Sequencer] object has a session id, sequence numbers are scoped by
260/// the (session_id, SeqKey) pair.
261#[derive(Clone, Debug)]
262pub struct Sequencer {
263    session_id: Uuid,
264    // Map's key is the sequence key (actor or port), value is the last seq number.
265    last_seqs: Arc<Mutex<FxHashMap<SeqKey, u64>>>,
266}
267
268impl Sequencer {
269    pub(crate) fn new(session_id: Uuid) -> Self {
270        Self {
271            session_id,
272            last_seqs: Arc::new(Mutex::new(FxHashMap::default())),
273        }
274    }
275
276    /// Assign the next seq for a port, mutate the sequencer with the new seq,
277    /// and return the new seq.
278    ///
279    /// - Handler ports: share the same sequence scheme per actor (keyed by ActorAddr).
280    /// - Ephemeral ports: get individual sequence schemes (keyed by PortAddr).
281    /// - Control ports: bypass receive-side reordering and use [`SeqInfo::Direct`].
282    pub fn assign_seq(&self, port_id: &PortAddr) -> SeqInfo {
283        let Some(key) = Self::seq_key(port_id) else {
284            return SeqInfo::Direct;
285        };
286
287        let mut guard = self.last_seqs.lock().unwrap();
288        let entry = guard.entry(key).or_default();
289        *entry += 1;
290        SeqInfo::Session {
291            session_id: self.session_id,
292            seq: *entry,
293        }
294    }
295
296    /// Last sequence sent to `port_id`.
297    pub fn last_sent(&self, port_id: &PortAddr) -> Option<u64> {
298        let key = Self::seq_key(port_id)?;
299        self.last_seqs.lock().unwrap().get(&key).copied()
300    }
301
302    /// Last sequence sent to any handler port on `actor`.
303    pub fn last_sent_to_actor(&self, actor: &ActorAddr) -> u64 {
304        self.last_seqs
305            .lock()
306            .unwrap()
307            .get(&SeqKey::for_handler(actor))
308            .copied()
309            .unwrap_or_default()
310    }
311
312    fn seq_key(port_id: &PortAddr) -> Option<SeqKey> {
313        if port_id.port().is_control() {
314            return None;
315        }
316
317        if port_id.is_handler_port() {
318            Some(SeqKey::for_handler(&port_id.actor_addr()))
319        } else {
320            SeqKey::for_ephemeral_port(port_id)
321        }
322    }
323
324    /// Assign the next sequence number to each key.
325    ///
326    /// This assigns the same sequence numbers, in the same order, as calling
327    /// [`Self::assign_seq`] for each equivalent destination. Each occurrence
328    /// of a repeated key advances its sequence number.
329    ///
330    /// The result is run-length encoded. Each `(range, seq)` assigns `seq` to
331    /// the input indexes in `range`. The ordered ranges cover `0..keys.len()`
332    /// without gaps. An empty input returns an empty result.
333    pub fn assign_seqs(&self, keys: &[SeqKey]) -> Vec<(Range<usize>, u64)> {
334        let mut runs: Vec<(Range<usize>, u64)> = Vec::new();
335        let mut run_start = 0usize;
336        let mut prev: Option<u64> = None;
337        let mut ordinal = 0;
338        for chunk in keys
339            .chunks(hyperactor_config::global::get(crate::config::SEQUENCER_MAX_LOCKED_KEYS).get())
340        {
341            let mut guard = self.last_seqs.lock().unwrap();
342            for key in chunk {
343                let seq = if let Some(last) = guard.get_mut(key) {
344                    *last += 1;
345                    *last
346                } else {
347                    guard.insert(key.clone(), 1);
348                    1
349                };
350                if let Some(prev_seq) = prev
351                    && prev_seq != seq
352                {
353                    runs.push((run_start..ordinal, prev_seq));
354                    run_start = ordinal;
355                }
356                prev = Some(seq);
357                ordinal += 1;
358            }
359        }
360        if let Some(prev_seq) = prev {
361            runs.push((run_start..keys.len(), prev_seq));
362        }
363        runs
364    }
365
366    /// Id of the session this sequencer belongs to.
367    pub fn session_id(&self) -> Uuid {
368        self.session_id
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use std::collections::HashMap;
375    use std::collections::HashSet;
376    use std::sync::Arc;
377    use std::sync::atomic::AtomicUsize;
378    use std::sync::atomic::Ordering as AtomicOrdering;
379
380    use async_trait::async_trait;
381    use hyperactor_config::Flattrs;
382    use rand::SeedableRng;
383    use rand::rngs::StdRng;
384    use rand::seq::SliceRandom;
385    use timed_test::async_timed_test;
386    use tokio::sync::Barrier;
387    use tokio::sync::oneshot;
388
389    use super::*;
390    use crate as hyperactor;
391    use crate::Actor;
392    use crate::ActorHandle;
393    use crate::ActorRef;
394    use crate::Context;
395    use crate::Endpoint as _;
396    use crate::Handler;
397    use crate::Proc;
398    use crate::config;
399    use crate::introspect::IntrospectMessage;
400    use crate::mailbox::headers::SENDER_ACTOR_ID;
401    use crate::mailbox::headers::stamp_sender_actor_id;
402    use crate::port::ControlPort;
403    use crate::port::Port;
404    use crate::proc::StatusMessage;
405    use crate::testing::ids::test_actor_id;
406
407    /// Test message type 1 for handler port sequencing tests.
408    #[derive(Named)]
409    struct TestMsg1;
410
411    /// Test message type 2 for handler port sequencing tests.
412    #[derive(Named)]
413    struct TestMsg2;
414
415    /// Helper to extract seq from SeqInfo::Session variant (for tests only)
416    fn get_seq(seq_info: SeqInfo) -> u64 {
417        match seq_info {
418            SeqInfo::Session { seq, .. } => seq,
419            SeqInfo::Direct => panic!("expected Session variant, got Direct"),
420        }
421    }
422
423    #[test]
424    fn seq_info_validity_rejects_zero_session_seq() {
425        let session_id = Uuid::now_v7();
426
427        assert!(SeqInfo::Direct.is_valid());
428        assert!(SeqInfo::Session { session_id, seq: 1 }.is_valid());
429        assert!(!SeqInfo::Session { session_id, seq: 0 }.is_valid());
430    }
431
432    #[test]
433    fn assign_seqs_returns_no_runs_for_empty_input() {
434        // GIVEN a sequencer with no assigned keys.
435        let sequencer = Sequencer::new(Uuid::now_v7());
436
437        // WHEN sequence numbers are requested for no keys.
438        let runs = sequencer.assign_seqs(&[]);
439
440        // THEN there are no encoded runs.
441        assert!(runs.is_empty());
442    }
443
444    #[test]
445    fn seq_key_for_ephemeral_port_rejects_other_port_kinds() {
446        // GIVEN one ephemeral, handler, and control port on the same actor.
447        let actor = test_actor_id("worker_0", "worker");
448        let ephemeral = actor.port_addr(Port::from(1));
449        let handler = actor.port_addr(Port::handler::<TestMsg1>());
450        let control = actor.port_addr(Port::control(ControlPort::Status));
451
452        // WHEN each port is passed to the ephemeral-key constructor.
453        let ephemeral_key = SeqKey::for_ephemeral_port(&ephemeral);
454        let handler_key = SeqKey::for_ephemeral_port(&handler);
455        let control_key = SeqKey::for_ephemeral_port(&control);
456
457        // THEN only the ephemeral port produces a key.
458        assert!(ephemeral_key.is_some());
459        assert!(handler_key.is_none());
460        assert!(control_key.is_none());
461    }
462
463    #[test]
464    fn assign_seqs_coalesces_synchronized_keys() {
465        // GIVEN synchronized keys and a chunk size smaller than the input.
466        let config = hyperactor_config::global::lock();
467        let _guard = config.override_key(
468            crate::config::SEQUENCER_MAX_LOCKED_KEYS,
469            hyperactor_config::NonZeroUsize::MIN,
470        );
471        let sequencer = Sequencer::new(Uuid::now_v7());
472        let keys: Vec<_> = (0..3)
473            .map(|i| SeqKey::for_handler(&test_actor_id(&format!("worker_{i}"), "worker")))
474            .collect();
475
476        // WHEN the keys advance together.
477        let single = sequencer.assign_seqs(&keys[..1]);
478        let first = sequencer.assign_seqs(&keys[1..]);
479        let second = sequencer.assign_seqs(&keys[1..]);
480
481        // THEN each synchronized set is represented by one run.
482        assert_eq!(single, vec![(0..1, 1)]);
483        assert_eq!(first, vec![(0..2, 1)]);
484        assert_eq!(second, vec![(0..2, 2)]);
485    }
486
487    #[test]
488    fn assign_seqs_keeps_nonadjacent_equal_values_in_separate_runs() {
489        // GIVEN the first and third keys are one sequence ahead of the second.
490        let sequencer = Sequencer::new(Uuid::now_v7());
491        let keys: Vec<_> = (0..3)
492            .map(|i| SeqKey::for_handler(&test_actor_id(&format!("worker_{i}"), "worker")))
493            .collect();
494        sequencer.assign_seqs(&[keys[0].clone(), keys[2].clone()]);
495
496        // WHEN all three keys are assigned in input order.
497        let runs = sequencer.assign_seqs(&keys);
498
499        // THEN the values [2, 1, 2] remain three ordered runs.
500        assert_eq!(runs, vec![(0..1, 2), (1..2, 1), (2..3, 2)]);
501    }
502
503    #[test]
504    fn assign_seqs_advances_repeated_interleaved_keys() {
505        // GIVEN two new keys repeated in an interleaved input.
506        let sequencer = Sequencer::new(Uuid::now_v7());
507        let a = SeqKey::for_handler(&test_actor_id("worker_0", "worker"));
508        let b = SeqKey::for_handler(&test_actor_id("worker_1", "worker"));
509        let keys = [a.clone(), b.clone(), a.clone(), b, a];
510
511        // WHEN the batch is assigned.
512        let runs = sequencer.assign_seqs(&keys);
513
514        // THEN every occurrence advances its key and adjacent equal values coalesce.
515        assert_eq!(runs, vec![(0..2, 1), (2..4, 2), (4..5, 3)]);
516    }
517
518    #[test]
519    fn assign_seqs_matches_assign_seq_after_mixed_history() {
520        // GIVEN scalar and batch sequencers with the same mixed history and a
521        // batch chunk size smaller than the requested assignment.
522        let config = hyperactor_config::global::lock();
523        let _guard = config.override_key(
524            crate::config::SEQUENCER_MAX_LOCKED_KEYS,
525            hyperactor_config::NonZeroUsize::new(2).expect("2 is non-zero"),
526        );
527        let scalar = Sequencer::new(Uuid::now_v7());
528        let batch = Sequencer::new(Uuid::now_v7());
529
530        let actor_0 = test_actor_id("worker_0", "worker");
531        let actor_1 = test_actor_id("worker_1", "worker");
532        let ports = [
533            actor_0.port_addr(Port::handler::<TestMsg1>()),
534            actor_1.port_addr(Port::handler::<TestMsg1>()),
535            actor_0.port_addr(Port::from(1)),
536        ];
537        for index in [0, 2, 0] {
538            assert_eq!(
539                get_seq(scalar.assign_seq(&ports[index])),
540                get_seq(batch.assign_seq(&ports[index]))
541            );
542        }
543
544        // WHEN one sequencer assigns an interleaved batch and the other uses scalars.
545        let indexes = [0, 1, 0, 2, 1];
546        let expected: Vec<_> = indexes
547            .iter()
548            .map(|&index| get_seq(scalar.assign_seq(&ports[index])))
549            .collect();
550        // `[(0..2, 1), (2..5, 2)]` -> `[1, 1, 2, 2, 2]`.
551        let actual: Vec<_> = {
552            let keys: Vec<_> = indexes
553                .iter()
554                .map(|&index| {
555                    let port = &ports[index];
556                    if port.is_handler_port() {
557                        SeqKey::for_handler(&port.actor_addr())
558                    } else {
559                        SeqKey::for_ephemeral_port(port).expect("test port must be ephemeral")
560                    }
561                })
562                .collect();
563
564            batch
565                .assign_seqs(&keys)
566                .into_iter()
567                .flat_map(|(range, seq)| range.map(move |_| seq))
568                .collect()
569        };
570
571        // THEN both APIs assign the same values and leave the same subsequent state.
572        assert_eq!(actual, expected);
573        for port in &ports {
574            assert_eq!(
575                get_seq(scalar.assign_seq(port)),
576                get_seq(batch.assign_seq(port))
577            );
578        }
579    }
580
581    #[test]
582    fn test_sequencer_clone() {
583        let sequencer = Sequencer {
584            session_id: Uuid::now_v7(),
585            last_seqs: Arc::new(Mutex::new(FxHashMap::default())),
586        };
587
588        let actor_ref: ActorAddr = test_actor_id("test_0", "test");
589        let port_ref = actor_ref.port_addr(Port::from(1));
590
591        // Modify original sequencer
592        sequencer.assign_seq(&port_ref);
593        sequencer.assign_seq(&port_ref);
594
595        // Clone should share the same state
596        let cloned_sequencer = sequencer.clone();
597        assert_eq!(sequencer.session_id(), cloned_sequencer.session_id(),);
598        assert_eq!(get_seq(cloned_sequencer.assign_seq(&port_ref)), 3);
599    }
600
601    #[test]
602    fn test_sequencer_handler_ports_share_sequence() {
603        let sequencer = Sequencer {
604            session_id: Uuid::now_v7(),
605            last_seqs: Arc::new(Mutex::new(FxHashMap::default())),
606        };
607
608        let actor_ref: ActorAddr = test_actor_id("worker_0", "worker");
609        // Two different handler ports for the same actor.
610        let handler_port_1 = actor_ref.port_addr(Port::handler::<TestMsg1>());
611        let handler_port_2 = actor_ref.port_addr(Port::handler::<TestMsg2>());
612
613        // Handler ports should share a sequence (keyed by ActorAddr)
614        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_1)), 1);
615        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_2)), 2); // continues from 1
616        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_1)), 3);
617
618        // Handler ports from a different actor get their own shared sequence
619        let actor_ref_2: ActorAddr = test_actor_id("worker_1", "worker");
620        let handler_port_3 = actor_ref_2.port_addr(Port::handler::<TestMsg1>());
621        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_3)), 1); // independent from actor_ref
622    }
623
624    #[test]
625    fn test_sequencer_non_handler_ports_have_independent_sequences() {
626        let sequencer = Sequencer {
627            session_id: Uuid::now_v7(),
628            last_seqs: Arc::new(Mutex::new(FxHashMap::default())),
629        };
630
631        let actor_ref_0: ActorAddr = test_actor_id("worker_0", "worker");
632        let actor_ref_1: ActorAddr = test_actor_id("worker_1", "worker");
633
634        // Ephemeral ports from the same actor.
635        let port_1 = actor_ref_0.port_addr(Port::from(1));
636        let port_2 = actor_ref_0.port_addr(Port::from(2));
637
638        // Non-handler ports should have independent sequences (keyed by PortAddr)
639        assert_eq!(get_seq(sequencer.assign_seq(&port_1)), 1);
640        assert_eq!(get_seq(sequencer.assign_seq(&port_2)), 1); // independent, starts at 1
641        assert_eq!(get_seq(sequencer.assign_seq(&port_1)), 2);
642        assert_eq!(get_seq(sequencer.assign_seq(&port_2)), 2);
643
644        // Non-handler ports from different actors are also independent
645        let port_3 = actor_ref_1.port_addr(Port::from(1));
646        assert_eq!(get_seq(sequencer.assign_seq(&port_3)), 1); // independent from port_1
647        assert_eq!(get_seq(sequencer.assign_seq(&port_1)), 3);
648        assert_eq!(get_seq(sequencer.assign_seq(&port_3)), 2);
649    }
650
651    #[test]
652    fn test_sequencer_mixed_handler_and_non_handler_ports() {
653        let sequencer = Sequencer {
654            session_id: Uuid::now_v7(),
655            last_seqs: Arc::new(Mutex::new(FxHashMap::default())),
656        };
657
658        let actor_ref: ActorAddr = test_actor_id("worker_0", "worker");
659
660        // Handler ports (share sequence per actor)
661        let handler_port_1 = actor_ref.port_addr(Port::handler::<TestMsg1>());
662        let handler_port_2 = actor_ref.port_addr(Port::handler::<TestMsg2>());
663
664        // Non-handler ports (independent sequences per port)
665        let non_handler_port_1 = actor_ref.port_addr(Port::from(1));
666        let non_handler_port_2 = actor_ref.port_addr(Port::from(2));
667
668        // Interleave sends to all port types
669        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_1)), 1);
670        assert_eq!(get_seq(sequencer.assign_seq(&non_handler_port_1)), 1); // independent
671        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_2)), 2); // continues handler sequence
672        assert_eq!(get_seq(sequencer.assign_seq(&non_handler_port_2)), 1); // independent
673        assert_eq!(get_seq(sequencer.assign_seq(&non_handler_port_1)), 2); // continues its own
674        assert_eq!(get_seq(sequencer.assign_seq(&handler_port_1)), 3); // continues handler sequence
675        assert_eq!(get_seq(sequencer.assign_seq(&non_handler_port_2)), 2); // continues its own
676    }
677
678    #[test]
679    fn bypass_registry_introspect_message() {
680        assert!(is_bypass_workq_type_id(TypeId::of::<IntrospectMessage>()));
681    }
682
683    #[test]
684    fn bypass_registry_status_message() {
685        assert!(is_bypass_workq_type_id(TypeId::of::<StatusMessage>()));
686    }
687
688    #[test]
689    fn control_port_uses_direct_seq_info() {
690        // Control messages bypass the actor work queue and therefore must not
691        // allocate session sequence numbers that a receive-side reorder buffer
692        // would interpret.
693        let sequencer = Sequencer::new(Uuid::now_v7());
694        let actor_ref: ActorAddr = test_actor_id("agent_0", "proc_agent");
695
696        let introspect_port = actor_ref.port_addr(Port::control(ControlPort::Introspect));
697        let regular_actor_port = actor_ref.port_addr(Port::handler::<TestMsg1>());
698
699        assert_eq!(sequencer.assign_seq(&introspect_port), SeqInfo::Direct);
700        assert_eq!(get_seq(sequencer.assign_seq(&regular_actor_port)), 1);
701        assert_eq!(sequencer.assign_seq(&introspect_port), SeqInfo::Direct);
702        assert_eq!(get_seq(sequencer.assign_seq(&regular_actor_port)), 2);
703    }
704
705    #[test]
706    fn test_sequencer_reports_last_sent() {
707        let sequencer = Sequencer::new(Uuid::now_v7());
708        let actor_ref: ActorAddr = test_actor_id("test_0", "test");
709        let handler_port = actor_ref.port_addr(Port::handler::<TestMsg1>());
710        let other_handler_port = actor_ref.port_addr(Port::handler::<TestMsg2>());
711        let control_port = actor_ref.port_addr(Port::control(ControlPort::Status));
712
713        assert_eq!(sequencer.last_sent(&handler_port), None);
714        assert_eq!(get_seq(sequencer.assign_seq(&handler_port)), 1);
715        assert_eq!(get_seq(sequencer.assign_seq(&other_handler_port)), 2);
716        assert_eq!(sequencer.last_sent(&handler_port), Some(2));
717        assert_eq!(sequencer.last_sent(&control_port), None);
718    }
719
720    #[test]
721    fn test_ordering_snapshot_reports_delivery_progress_from_sender() {
722        let sender = test_actor_id("sender", "test");
723        let other = test_actor_id("other", "test");
724        let snapshot = OrderingSnapshot {
725            enabled: true,
726            sessions: vec![
727                OrderingSessionSnapshot {
728                    session_id: Uuid::from_u128(1),
729                    sender: Some(sender.clone()),
730                    last_released_seq: 3,
731                    expected_next_seq: 4,
732                    buffered_count: 0,
733                    oldest_buffered_seq: None,
734                    newest_buffered_seq: None,
735                },
736                OrderingSessionSnapshot {
737                    session_id: Uuid::from_u128(2),
738                    sender: Some(sender.clone()),
739                    last_released_seq: 5,
740                    expected_next_seq: 6,
741                    buffered_count: 1,
742                    oldest_buffered_seq: Some(7),
743                    newest_buffered_seq: Some(7),
744                },
745                OrderingSessionSnapshot {
746                    session_id: Uuid::from_u128(3),
747                    sender: Some(other.clone()),
748                    last_released_seq: 8,
749                    expected_next_seq: 9,
750                    buffered_count: 0,
751                    oldest_buffered_seq: None,
752                    newest_buffered_seq: None,
753                },
754            ],
755            skipped_session_count: 1,
756        };
757
758        assert_eq!(snapshot.delivery_progress_from(&sender), None,);
759
760        let disabled_snapshot = OrderingSnapshot {
761            enabled: false,
762            sessions: Vec::new(),
763            skipped_session_count: 0,
764        };
765        assert_eq!(disabled_snapshot.delivery_progress_from(&sender), None);
766
767        let snapshot = OrderingSnapshot {
768            skipped_session_count: 0,
769            ..snapshot
770        };
771        assert_eq!(
772            snapshot.delivery_progress_from(&sender),
773            Some(DeliveryProgress {
774                largest_dequeueable_sequence: 5,
775            })
776        );
777        assert_eq!(
778            snapshot.delivery_progress_from(&test_actor_id("missing", "test")),
779            Some(DeliveryProgress {
780                largest_dequeueable_sequence: 0,
781            })
782        );
783        assert_eq!(
784            snapshot.delivery_progress_for_session(Uuid::from_u128(1)),
785            Some(DeliveryProgress {
786                largest_dequeueable_sequence: 3,
787            })
788        );
789        assert_eq!(
790            snapshot.delivery_progress_for_session(Uuid::from_u128(2)),
791            Some(DeliveryProgress {
792                largest_dequeueable_sequence: 5,
793            })
794        );
795        assert_eq!(
796            snapshot.delivery_progress_for_session(Uuid::from_u128(4)),
797            Some(DeliveryProgress {
798                largest_dequeueable_sequence: 0,
799            })
800        );
801    }
802
803    /// Sequencer-level test for the debug-skip helper's underlying behavior:
804    /// `assign_seq` called `count` times advances the per-(actor,dest)
805    /// counter, so a subsequent `assign_seq` returns `count + 1` not 1.
806    /// The `Instance<A>::debug_skip_next_ordering_seq` method is a tiny
807    /// for-loop of `assign_seq` calls; its correctness reduces to this.
808    #[test]
809    fn test_sequencer_skip_advances_counter() {
810        let sequencer = Sequencer::new(Uuid::now_v7());
811        let actor_ref: ActorAddr = test_actor_id("test_0", "test");
812        let dest = actor_ref.port_addr(Port::handler::<TestMsg1>());
813
814        // Assign once: seq 1.
815        assert_eq!(get_seq(sequencer.assign_seq(&dest)), 1);
816
817        // Skip 2 (the helper would loop assign_seq).
818        for _ in 0..2 {
819            let _ = sequencer.assign_seq(&dest);
820        }
821
822        // Next assignment is seq 4, not 2.
823        assert_eq!(get_seq(sequencer.assign_seq(&dest)), 4);
824    }
825
826    /// Pins the exact `AttrValue` contract used by `declare_attrs!`
827    /// storage: `AttrValue::display(&self) -> String` then
828    /// `AttrValue::parse(&str) -> Result<Self, _>`.
829    #[test]
830    fn test_snapshot_attr_roundtrip() {
831        let addr: ActorAddr = test_actor_id("a", "client");
832        let snap = OrderingSnapshot {
833            enabled: true,
834            sessions: vec![OrderingSessionSnapshot {
835                session_id: Uuid::from_u128(42),
836                sender: Some(addr),
837                last_released_seq: 7,
838                expected_next_seq: 8,
839                buffered_count: 2,
840                oldest_buffered_seq: Some(9),
841                newest_buffered_seq: Some(11),
842            }],
843            skipped_session_count: 3,
844        };
845        let s = AttrValue::display(&snap);
846        let parsed = <OrderingSnapshot as AttrValue>::parse(&s).unwrap();
847        assert_eq!(snap, parsed);
848    }
849
850    // End-to-end ordering tests with real actor mailboxes. SenderActors
851    // send normally to a ChaosActor. ChaosActor forwards the same messages
852    // to ReceiverActor, but in shuffled batches and with optional duplicate
853    // replay. ReceiverActor only records what its handler observes; its
854    // sequenced handler receiver restores per-session order.
855
856    #[derive(Clone, Debug, Serialize, Deserialize, Named)]
857    struct Frame {
858        sender_idx: u32,
859        payload_idx: u64,
860    }
861
862    #[derive(Clone, Debug, Serialize, Deserialize, Named)]
863    struct SenderDone {
864        sender_idx: u32,
865    }
866
867    // What ReceiverActor observed after its sequenced handler receiver released a
868    // frame: the preserved ordering header, the sender-owner header, and the
869    // test payload.
870    type ReceivedFrame = (SeqInfo, Option<ActorAddr>, Frame);
871
872    // Start is only sent to a local ActorHandle, so it can carry local-only
873    // test handles like ActorHandle and Arc<Barrier>. Do not export it.
874    #[derive(Clone, Debug)]
875    struct Start {
876        count: u64,
877        target: ActorHandle<ChaosActor>,
878        start_barrier: Option<Arc<Barrier>>,
879        yield_between_frames: bool,
880    }
881
882    #[derive(Debug)]
883    struct SenderActor {
884        sender_idx: u32,
885    }
886
887    #[async_trait]
888    impl Actor for SenderActor {}
889
890    #[async_trait]
891    impl Handler<Start> for SenderActor {
892        async fn handle(&mut self, cx: &Context<Self>, msg: Start) -> Result<(), anyhow::Error> {
893            let Start {
894                count,
895                target,
896                start_barrier,
897                yield_between_frames,
898            } = msg;
899            if let Some(b) = start_barrier {
900                b.wait().await;
901            }
902            for payload_idx in 0..count {
903                target.post(
904                    cx,
905                    Frame {
906                        sender_idx: self.sender_idx,
907                        payload_idx,
908                    },
909                );
910                if yield_between_frames {
911                    tokio::task::yield_now().await;
912                }
913            }
914            target.post(
915                cx,
916                SenderDone {
917                    sender_idx: self.sender_idx,
918                },
919            );
920            Ok(())
921        }
922    }
923
924    // Which original frames should be replayed after their batch has been
925    // forwarded once.
926    #[derive(Debug, Clone, Copy)]
927    enum DuplicatePolicy {
928        None,
929        EveryNth { stride: usize },
930    }
931
932    #[derive(Debug, Default)]
933    struct ChaosStats {
934        // Number of original frames that reveal a per-session inversion.
935        // The counter is bumped on the lower sequence number, after a higher
936        // sequence number from the same SEQ_INFO session was already
937        // forwarded.
938        out_of_order_original_forwards: AtomicUsize,
939
940        // Number of original frames selected for replay.
941        duplicates_selected: AtomicUsize,
942
943        // Number of selected replays actually forwarded.
944        late_duplicate_forwards: AtomicUsize,
945
946        // Sessions that saw at least one inversion.
947        out_of_order_sessions: Mutex<HashSet<Uuid>>,
948
949        // Every frame ChaosActor forwarded, in forward order. Printed on
950        // failures so the panic includes the chaos-side trace.
951        forwarded_trace: Mutex<Vec<ForwardedEntry>>,
952
953        // Frames observed by ReceiverActor after it already sent the result
954        // snapshot. Non-zero is a failure. Zero is diagnostic only because
955        // the test does not wait for a final drain.
956        receiver_overflow: AtomicUsize,
957    }
958
959    // Used only in panic output.
960    #[allow(dead_code)]
961    #[derive(Debug, Clone)]
962    struct ForwardedEntry {
963        session_id: Uuid,
964        seq: u64,
965        sender_idx: u32,
966        is_duplicate: bool,
967    }
968
969    #[derive(Debug)]
970    #[hyperactor::export(handlers = [Frame, SenderDone])]
971    struct ChaosActor {
972        window: Vec<(SeqInfo, Flattrs, Frame)>,
973        window_size: usize,
974        target_port: PortAddr,
975        rng: StdRng,
976        expected_dones: u32,
977        done_count: u32,
978        expected_count_per_sender: u64,
979        duplicate_policy: DuplicatePolicy,
980        duplicate_cursor: usize,
981        stats: Arc<ChaosStats>,
982        session_owners: HashMap<Uuid, ActorAddr>,
983        max_forwarded_seq_by_session: HashMap<Uuid, u64>,
984    }
985
986    #[async_trait]
987    impl Actor for ChaosActor {}
988
989    impl ChaosActor {
990        // Forward the current batch to the receiver. The batch is shuffled,
991        // then adjusted if needed so at least one session is out of order
992        // whenever the batch has enough frames to make that possible.
993        fn flush(&mut self, cx: &Context<Self>) -> Result<(), anyhow::Error> {
994            self.window.shuffle(&mut self.rng);
995
996            // A shuffled batch can accidentally still be ordered. If so,
997            // pick the first session in window order that has at least two
998            // frames and swap its first two frames.
999            let mut first_seen_order: Vec<Uuid> = Vec::new();
1000            let mut session_indices: HashMap<Uuid, Vec<usize>> = HashMap::new();
1001            for (i, (seq_info, _, _)) in self.window.iter().enumerate() {
1002                if let SeqInfo::Session { session_id, .. } = seq_info {
1003                    if !session_indices.contains_key(session_id) {
1004                        first_seen_order.push(*session_id);
1005                    }
1006                    session_indices.entry(*session_id).or_default().push(i);
1007                }
1008            }
1009            for sid in &first_seen_order {
1010                let indices = &session_indices[sid];
1011                if indices.len() < 2 {
1012                    continue;
1013                }
1014                let mut already_inverted = false;
1015                for w in indices.windows(2) {
1016                    if seq_of(&self.window[w[0]].0) > seq_of(&self.window[w[1]].0) {
1017                        already_inverted = true;
1018                        break;
1019                    }
1020                }
1021                if !already_inverted {
1022                    self.window.swap(indices[0], indices[1]);
1023                    break;
1024                }
1025            }
1026
1027            // Send each shuffled frame once. Some frames are cloned for a
1028            // later replay.
1029            let mut duplicates: Vec<(SeqInfo, Flattrs, Frame)> = Vec::new();
1030            let entries: Vec<(SeqInfo, Flattrs, Frame)> = std::mem::take(&mut self.window);
1031            for (seq_info, headers, frame) in entries {
1032                let (session_id, seq) = match &seq_info {
1033                    SeqInfo::Session { session_id, seq } => (*session_id, *seq),
1034                    SeqInfo::Direct => panic!("Direct SeqInfo at flush"),
1035                };
1036                // This frame proves an inversion if a higher sequence number
1037                // from the same session was already forwarded.
1038                let prev_max = self
1039                    .max_forwarded_seq_by_session
1040                    .get(&session_id)
1041                    .copied()
1042                    .unwrap_or(0);
1043                if seq < prev_max {
1044                    self.stats
1045                        .out_of_order_original_forwards
1046                        .fetch_add(1, AtomicOrdering::Relaxed);
1047                    self.stats
1048                        .out_of_order_sessions
1049                        .lock()
1050                        .unwrap()
1051                        .insert(session_id);
1052                }
1053                self.max_forwarded_seq_by_session
1054                    .insert(session_id, prev_max.max(seq));
1055                // Select every Nth original frame for later replay.
1056                let should_duplicate = match self.duplicate_policy {
1057                    DuplicatePolicy::None => false,
1058                    DuplicatePolicy::EveryNth { stride } => {
1059                        assert!(stride > 0, "DuplicatePolicy::EveryNth requires stride > 0");
1060                        self.duplicate_cursor = self.duplicate_cursor.wrapping_add(1);
1061                        self.duplicate_cursor.is_multiple_of(stride)
1062                    }
1063                };
1064                if should_duplicate {
1065                    duplicates.push((seq_info.clone(), headers.clone(), frame.clone()));
1066                    self.stats
1067                        .duplicates_selected
1068                        .fetch_add(1, AtomicOrdering::Relaxed);
1069                }
1070                // Forward with the original sequence header and original
1071                // sender owner, like a routing hop would.
1072                let owner = self
1073                    .session_owners
1074                    .get(&session_id)
1075                    .expect("session owner missing at original forward")
1076                    .clone();
1077                let dest = self.target_port.clone();
1078                let mut outbound = headers;
1079                stamp_sender_actor_id(&mut outbound, &seq_info, &dest, &owner);
1080                cx.post_with_external_seq_info(dest, outbound, wirevalue::Any::serialize(&frame)?);
1081                self.stats
1082                    .forwarded_trace
1083                    .lock()
1084                    .unwrap()
1085                    .push(ForwardedEntry {
1086                        session_id,
1087                        seq,
1088                        sender_idx: frame.sender_idx,
1089                        is_duplicate: false,
1090                    });
1091            }
1092
1093            // Replay selected frames after the original batch. A duplicate
1094            // that arrives while its original is still buffered is dropped
1095            // as a duplicate; these replays are intended to hit the
1096            // late-duplicate drop path instead.
1097            for (seq_info, headers, frame) in duplicates {
1098                let (session_id, seq) = match &seq_info {
1099                    SeqInfo::Session { session_id, seq } => (*session_id, *seq),
1100                    SeqInfo::Direct => unreachable!(),
1101                };
1102                let owner = self
1103                    .session_owners
1104                    .get(&session_id)
1105                    .expect("session owner missing at duplicate forward")
1106                    .clone();
1107                let dest = self.target_port.clone();
1108                let mut outbound = headers;
1109                stamp_sender_actor_id(&mut outbound, &seq_info, &dest, &owner);
1110                cx.post_with_external_seq_info(dest, outbound, wirevalue::Any::serialize(&frame)?);
1111                self.stats
1112                    .late_duplicate_forwards
1113                    .fetch_add(1, AtomicOrdering::Relaxed);
1114                self.stats
1115                    .forwarded_trace
1116                    .lock()
1117                    .unwrap()
1118                    .push(ForwardedEntry {
1119                        session_id,
1120                        seq,
1121                        sender_idx: frame.sender_idx,
1122                        is_duplicate: true,
1123                    });
1124            }
1125
1126            Ok(())
1127        }
1128    }
1129
1130    fn seq_of(s: &SeqInfo) -> u64 {
1131        match s {
1132            SeqInfo::Session { seq, .. } => *seq,
1133            SeqInfo::Direct => panic!("seq_of called on SeqInfo::Direct"),
1134        }
1135    }
1136
1137    #[async_trait]
1138    impl Handler<Frame> for ChaosActor {
1139        async fn handle(&mut self, cx: &Context<Self>, frame: Frame) -> Result<(), anyhow::Error> {
1140            let seq_info = cx.headers().get(SEQ_INFO);
1141            let (session_id, seq) = match &seq_info {
1142                Some(SeqInfo::Session { session_id, seq }) => (*session_id, *seq),
1143                Some(SeqInfo::Direct) => {
1144                    panic!("chaos inbound has SeqInfo::Direct; bind misconfigured")
1145                }
1146                None => panic!("chaos inbound missing SEQ_INFO"),
1147            };
1148            let sender_addr = cx.headers().get(SENDER_ACTOR_ID);
1149            if seq <= 4 {
1150                assert!(
1151                    sender_addr.is_some(),
1152                    "missing SENDER_ACTOR_ID on early-session chaos inbound (seq={seq})",
1153                );
1154            }
1155            if let Some(addr) = &sender_addr {
1156                match self.session_owners.entry(session_id) {
1157                    std::collections::hash_map::Entry::Vacant(v) => {
1158                        v.insert(addr.clone());
1159                    }
1160                    std::collections::hash_map::Entry::Occupied(o) => {
1161                        assert_eq!(
1162                            o.get(),
1163                            addr,
1164                            "session owner changed mid-stream for session_id={session_id}",
1165                        );
1166                    }
1167                }
1168            }
1169            let inbound_headers = cx.headers().clone();
1170            self.window
1171                .push((seq_info.unwrap(), inbound_headers, frame));
1172            if self.window.len() >= self.window_size {
1173                self.flush(cx)?;
1174            }
1175            Ok(())
1176        }
1177    }
1178
1179    #[async_trait]
1180    impl Handler<SenderDone> for ChaosActor {
1181        async fn handle(
1182            &mut self,
1183            cx: &Context<Self>,
1184            done: SenderDone,
1185        ) -> Result<(), anyhow::Error> {
1186            let seq_info = cx.headers().get(SEQ_INFO);
1187            let (session_id, seq) = match &seq_info {
1188                Some(SeqInfo::Session { session_id, seq }) => (*session_id, *seq),
1189                Some(SeqInfo::Direct) => {
1190                    panic!("SenderDone arrived as SeqInfo::Direct; bind misconfigured")
1191                }
1192                None => panic!("SenderDone missing SEQ_INFO"),
1193            };
1194            // The sender posts all Frames and then SenderDone to the same
1195            // actor. They should share one sequence stream.
1196            assert_eq!(
1197                seq,
1198                self.expected_count_per_sender + 1,
1199                "SenderDone(sender_idx={}) seq={seq}; expected {} (count+1)",
1200                done.sender_idx,
1201                self.expected_count_per_sender + 1,
1202            );
1203            assert!(
1204                self.session_owners.contains_key(&session_id),
1205                "SenderDone(sender_idx={}) for session_id={session_id} \
1206                 before any Frame from that session",
1207                done.sender_idx,
1208            );
1209            // SenderDone is usually past the early messages that carry the
1210            // sender-owner header. If the header is present, it must still
1211            // match the session owner learned from the Frames.
1212            if let Some(addr) = cx.headers().get(SENDER_ACTOR_ID) {
1213                let owner = self.session_owners.get(&session_id).unwrap();
1214                assert_eq!(
1215                    &addr, owner,
1216                    "SenderDone SENDER_ACTOR_ID inconsistent with session_owner \
1217                     for session_id={session_id}",
1218                );
1219            }
1220            assert!(
1221                self.done_count < self.expected_dones,
1222                "extra SenderDone(sender_idx={}); done_count already at expected_dones={}",
1223                done.sender_idx,
1224                self.expected_dones,
1225            );
1226            self.done_count += 1;
1227            if self.done_count == self.expected_dones && !self.window.is_empty() {
1228                self.flush(cx)?;
1229            }
1230            Ok(())
1231        }
1232    }
1233
1234    #[derive(Debug)]
1235    #[hyperactor::export(handlers = [Frame])]
1236    struct ReceiverActor {
1237        received: Vec<ReceivedFrame>,
1238        expected_total: usize,
1239        done: Option<oneshot::Sender<Vec<ReceivedFrame>>>,
1240        stats: Arc<ChaosStats>,
1241    }
1242
1243    #[async_trait]
1244    impl Actor for ReceiverActor {}
1245
1246    #[async_trait]
1247    impl Handler<Frame> for ReceiverActor {
1248        async fn handle(&mut self, cx: &Context<Self>, frame: Frame) -> Result<(), anyhow::Error> {
1249            let seq_info = match cx.headers().get(SEQ_INFO) {
1250                Some(SeqInfo::Session { session_id, seq }) => SeqInfo::Session { session_id, seq },
1251                Some(SeqInfo::Direct) => {
1252                    panic!("receiver inbound has SeqInfo::Direct; chaos forward bypassed SEQ_INFO")
1253                }
1254                None => panic!("receiver inbound missing SEQ_INFO"),
1255            };
1256            let sender_addr = cx.headers().get(SENDER_ACTOR_ID);
1257            // Anything after the snapshot is an extra delivery. Count it, but
1258            // keep the original snapshot unchanged for assertions.
1259            if self.done.is_none() {
1260                self.stats
1261                    .receiver_overflow
1262                    .fetch_add(1, AtomicOrdering::Relaxed);
1263                return Ok(());
1264            }
1265            self.received.push((seq_info, sender_addr, frame));
1266            if self.received.len() == self.expected_total {
1267                let tx = self.done.take().expect("done sender already consumed");
1268                let snapshot = std::mem::take(&mut self.received);
1269                let _ = tx.send(snapshot);
1270            }
1271            Ok(())
1272        }
1273    }
1274
1275    // Checks the receiver-side invariants shared by all chaos tests. Panics
1276    // include the ChaosActor forward trace.
1277    fn assert_received_protocol_correct(
1278        received: &[ReceivedFrame],
1279        expected_senders: usize,
1280        expected_count_per_sender: u64,
1281        sender_addrs: &[ActorAddr],
1282        stats: &ChaosStats,
1283    ) {
1284        let expected_total = expected_senders * (expected_count_per_sender as usize);
1285        let dump_trace = || -> String {
1286            let trace = stats.forwarded_trace.lock().unwrap();
1287            format!("forwarded_trace ({} entries) = {:#?}", trace.len(), *trace,)
1288        };
1289
1290        assert_eq!(
1291            received.len(),
1292            expected_total,
1293            "expected {} frames at receiver, got {}; {}",
1294            expected_total,
1295            received.len(),
1296            dump_trace(),
1297        );
1298
1299        // This catches extra deliveries after the receiver snapshot.
1300        let overflow = stats.receiver_overflow.load(AtomicOrdering::Acquire);
1301        assert_eq!(
1302            overflow,
1303            0,
1304            "stats.receiver_overflow = {overflow}; duplicate leaked past \
1305             sequenced receiver drop branch after snapshot; {}",
1306            dump_trace(),
1307        );
1308
1309        // Each logical sender should map to exactly one wire session.
1310        let mut sender_session: HashMap<u32, Uuid> = HashMap::new();
1311        for (seq_info, _, frame) in received {
1312            let session_id = match seq_info {
1313                SeqInfo::Session { session_id, .. } => *session_id,
1314                SeqInfo::Direct => {
1315                    panic!("receiver captured Direct SEQ_INFO; {}", dump_trace())
1316                }
1317            };
1318            match sender_session.entry(frame.sender_idx) {
1319                std::collections::hash_map::Entry::Vacant(v) => {
1320                    v.insert(session_id);
1321                }
1322                std::collections::hash_map::Entry::Occupied(o) => {
1323                    assert_eq!(
1324                        o.get(),
1325                        &session_id,
1326                        "sender_idx={} mapped to two distinct sessions ({} and {}); {}",
1327                        frame.sender_idx,
1328                        o.get(),
1329                        session_id,
1330                        dump_trace(),
1331                    );
1332                }
1333            }
1334        }
1335
1336        // And each wire session should belong to exactly one logical sender.
1337        assert_eq!(
1338            sender_session.len(),
1339            expected_senders,
1340            "expected {} distinct sender_idx, got {}; {}",
1341            expected_senders,
1342            sender_session.len(),
1343            dump_trace(),
1344        );
1345        let distinct_sessions: HashSet<Uuid> = sender_session.values().copied().collect();
1346        assert_eq!(
1347            distinct_sessions.len(),
1348            expected_senders,
1349            "sender_idx -> session_id is not a bijection (got {} distinct sessions for {} senders); {}",
1350            distinct_sessions.len(),
1351            expected_senders,
1352            dump_trace(),
1353        );
1354
1355        // Each sender's delivered subsequence should be exactly 0..K payloads
1356        // and 1..=K sequence numbers.
1357        for sender_idx in 0..(expected_senders as u32) {
1358            let subseq: Vec<&ReceivedFrame> = received
1359                .iter()
1360                .filter(|(_, _, f)| f.sender_idx == sender_idx)
1361                .collect();
1362            assert_eq!(
1363                subseq.len() as u64,
1364                expected_count_per_sender,
1365                "sender_idx={sender_idx}: expected {} frames, got {}; {}",
1366                expected_count_per_sender,
1367                subseq.len(),
1368                dump_trace(),
1369            );
1370            for (i, (seq_info, _, frame)) in subseq.iter().enumerate() {
1371                let expected_payload = i as u64;
1372                assert_eq!(
1373                    frame.payload_idx,
1374                    expected_payload,
1375                    "sender_idx={sender_idx}, position {i}: payload_idx={} (expected {}); {}",
1376                    frame.payload_idx,
1377                    expected_payload,
1378                    dump_trace(),
1379                );
1380                let expected_seq = (i as u64) + 1;
1381                let actual_seq = seq_of(seq_info);
1382                assert_eq!(
1383                    actual_seq,
1384                    expected_seq,
1385                    "sender_idx={sender_idx}, position {i}: SEQ_INFO.seq={} (expected {}); {}",
1386                    actual_seq,
1387                    expected_seq,
1388                    dump_trace(),
1389                );
1390            }
1391        }
1392
1393        // Early messages should identify the actor that owns the sequence.
1394        for (seq_info, sender_addr_opt, frame) in received {
1395            let seq = seq_of(seq_info);
1396            if seq > 4 {
1397                continue;
1398            }
1399            let expected_addr = &sender_addrs[frame.sender_idx as usize];
1400            let captured = sender_addr_opt.as_ref().unwrap_or_else(|| {
1401                panic!(
1402                    "sender_idx={}, seq={seq}: SENDER_ACTOR_ID is None at early-session; {}",
1403                    frame.sender_idx,
1404                    dump_trace(),
1405                )
1406            });
1407            assert_eq!(
1408                captured,
1409                expected_addr,
1410                "sender_idx={}, seq={seq}: SENDER_ACTOR_ID {captured:?} != expected {expected_addr:?}; {}",
1411                frame.sender_idx,
1412                dump_trace(),
1413            );
1414        }
1415    }
1416
1417    // One sender, shuffled batches, no duplicates.
1418    #[async_timed_test(timeout_secs = 30)]
1419    async fn test_chaos_single_sender_preserves_order() {
1420        let config = hyperactor_config::global::lock();
1421        let _g = config.override_key(config::ENABLE_DEST_ACTOR_REORDERING_BUFFER, true);
1422
1423        let proc = Proc::isolated();
1424        let client = proc.client("client");
1425
1426        let stats = Arc::new(ChaosStats::default());
1427        let n: usize = 1;
1428        let k: u64 = 50;
1429
1430        let (done_tx, done_rx) = oneshot::channel();
1431        let receiver_handle = proc.spawn_with_label(
1432            "receiver",
1433            ReceiverActor {
1434                received: Vec::new(),
1435                expected_total: n * (k as usize),
1436                done: Some(done_tx),
1437                stats: stats.clone(),
1438            },
1439        );
1440        let receiver_ref: ActorRef<ReceiverActor> = receiver_handle.bind();
1441        let target_port = receiver_ref.port::<Frame>().port_addr().clone();
1442
1443        let chaos_handle = proc.spawn_with_label(
1444            "chaos",
1445            ChaosActor {
1446                window: Vec::new(),
1447                window_size: 10,
1448                target_port,
1449                rng: StdRng::seed_from_u64(0xC4A0_5EED),
1450                expected_dones: n as u32,
1451                done_count: 0,
1452                expected_count_per_sender: k,
1453                duplicate_policy: DuplicatePolicy::None,
1454                duplicate_cursor: 0,
1455                stats: stats.clone(),
1456                session_owners: HashMap::new(),
1457                max_forwarded_seq_by_session: HashMap::new(),
1458            },
1459        );
1460        // Bind exported handler ports so local handle posts carry
1461        // SEQ_INFO::Session instead of SeqInfo::Direct.
1462        let _bound_chaos_ref: ActorRef<ChaosActor> = chaos_handle.bind();
1463
1464        let sender_handle = proc.spawn_with_label("sender0", SenderActor { sender_idx: 0 });
1465        let sender_addr = sender_handle.actor_addr().clone();
1466
1467        sender_handle.post(
1468            &client,
1469            Start {
1470                count: k,
1471                target: chaos_handle.clone(),
1472                start_barrier: None,
1473                yield_between_frames: false,
1474            },
1475        );
1476
1477        let received = match tokio::time::timeout(std::time::Duration::from_secs(5), done_rx).await
1478        {
1479            Ok(Ok(v)) => v,
1480            Ok(Err(_)) => panic!("done_rx sender dropped before sending; receiver crashed?"),
1481            Err(_) => {
1482                let trace = stats.forwarded_trace.lock().unwrap();
1483                panic!(
1484                    "timed out waiting for receiver; forwarded_trace ({} entries) = {:#?}",
1485                    trace.len(),
1486                    *trace,
1487                );
1488            }
1489        };
1490
1491        assert_received_protocol_correct(&received, n, k, &[sender_addr], &stats);
1492
1493        let oo = stats
1494            .out_of_order_original_forwards
1495            .load(AtomicOrdering::Acquire);
1496        assert!(
1497            oo > 0,
1498            "expected out_of_order_original_forwards > 0, got {oo}: \
1499             chaos shuffle never inverted a real frame (shuffle was identity?)",
1500        );
1501    }
1502
1503    // Multiple senders share one chaos actor. Each sender must be delivered
1504    // in order; cross-sender interleaving is unconstrained.
1505    #[async_timed_test(timeout_secs = 30)]
1506    async fn test_chaos_multi_sender_preserves_per_session_order() {
1507        let config = hyperactor_config::global::lock();
1508        let _g = config.override_key(config::ENABLE_DEST_ACTOR_REORDERING_BUFFER, true);
1509
1510        let proc = Proc::isolated();
1511        let client = proc.client("client");
1512
1513        let stats = Arc::new(ChaosStats::default());
1514        let n: usize = 4;
1515        let k: u64 = 25;
1516
1517        let (done_tx, done_rx) = oneshot::channel();
1518        let receiver_handle = proc.spawn_with_label(
1519            "receiver",
1520            ReceiverActor {
1521                received: Vec::new(),
1522                expected_total: n * (k as usize),
1523                done: Some(done_tx),
1524                stats: stats.clone(),
1525            },
1526        );
1527        let receiver_ref: ActorRef<ReceiverActor> = receiver_handle.bind();
1528        let target_port = receiver_ref.port::<Frame>().port_addr().clone();
1529
1530        let chaos_handle = proc.spawn_with_label(
1531            "chaos",
1532            ChaosActor {
1533                window: Vec::new(),
1534                window_size: 10,
1535                target_port,
1536                rng: StdRng::seed_from_u64(0xC4A0_5EED),
1537                expected_dones: n as u32,
1538                done_count: 0,
1539                expected_count_per_sender: k,
1540                duplicate_policy: DuplicatePolicy::None,
1541                duplicate_cursor: 0,
1542                stats: stats.clone(),
1543                session_owners: HashMap::new(),
1544                max_forwarded_seq_by_session: HashMap::new(),
1545            },
1546        );
1547        // Bind exported handler ports so local handle posts carry
1548        // SEQ_INFO::Session instead of SeqInfo::Direct.
1549        let _bound_chaos_ref: ActorRef<ChaosActor> = chaos_handle.bind();
1550
1551        let barrier = Arc::new(Barrier::new(n));
1552        let mut sender_addrs: Vec<ActorAddr> = Vec::with_capacity(n);
1553        let mut sender_handles: Vec<ActorHandle<SenderActor>> = Vec::with_capacity(n);
1554        for sender_idx in 0..n {
1555            let h = proc.spawn_with_label(
1556                &format!("sender{sender_idx}"),
1557                SenderActor {
1558                    sender_idx: sender_idx as u32,
1559                },
1560            );
1561            sender_addrs.push(h.actor_addr().clone());
1562            sender_handles.push(h);
1563        }
1564        for h in &sender_handles {
1565            h.post(
1566                &client,
1567                Start {
1568                    count: k,
1569                    target: chaos_handle.clone(),
1570                    start_barrier: Some(barrier.clone()),
1571                    yield_between_frames: true,
1572                },
1573            );
1574        }
1575
1576        let received = match tokio::time::timeout(std::time::Duration::from_secs(5), done_rx).await
1577        {
1578            Ok(Ok(v)) => v,
1579            Ok(Err(_)) => panic!("done_rx sender dropped before sending; receiver crashed?"),
1580            Err(_) => {
1581                let trace = stats.forwarded_trace.lock().unwrap();
1582                panic!(
1583                    "timed out waiting for receiver; forwarded_trace ({} entries) = {:#?}",
1584                    trace.len(),
1585                    *trace,
1586                );
1587            }
1588        };
1589
1590        assert_received_protocol_correct(&received, n, k, &sender_addrs, &stats);
1591
1592        let oo = stats
1593            .out_of_order_original_forwards
1594            .load(AtomicOrdering::Acquire);
1595        assert!(
1596            oo > 0,
1597            "expected out_of_order_original_forwards > 0, got {oo}",
1598        );
1599        // Keep the diagnostic set consistent with the global counter.
1600        let sessions_with_inversion = stats.out_of_order_sessions.lock().unwrap().len();
1601        assert!(
1602            sessions_with_inversion >= 1,
1603            "expected out_of_order_sessions.len() >= 1, got {sessions_with_inversion}",
1604        );
1605    }
1606
1607    // Same as the multi-sender case, but replay every 10th original frame.
1608    // The receiver should still see each original exactly once.
1609    #[async_timed_test(timeout_secs = 30)]
1610    async fn test_chaos_drops_duplicates() {
1611        let config = hyperactor_config::global::lock();
1612        let _g = config.override_key(config::ENABLE_DEST_ACTOR_REORDERING_BUFFER, true);
1613
1614        let proc = Proc::isolated();
1615        let client = proc.client("client");
1616
1617        let stats = Arc::new(ChaosStats::default());
1618        let n: usize = 4;
1619        let k: u64 = 25;
1620
1621        let (done_tx, done_rx) = oneshot::channel();
1622        let receiver_handle = proc.spawn_with_label(
1623            "receiver",
1624            ReceiverActor {
1625                received: Vec::new(),
1626                expected_total: n * (k as usize),
1627                done: Some(done_tx),
1628                stats: stats.clone(),
1629            },
1630        );
1631        let receiver_ref: ActorRef<ReceiverActor> = receiver_handle.bind();
1632        let target_port = receiver_ref.port::<Frame>().port_addr().clone();
1633
1634        let chaos_handle = proc.spawn_with_label(
1635            "chaos",
1636            ChaosActor {
1637                window: Vec::new(),
1638                window_size: 10,
1639                target_port,
1640                rng: StdRng::seed_from_u64(0xC4A0_5EED),
1641                expected_dones: n as u32,
1642                done_count: 0,
1643                expected_count_per_sender: k,
1644                duplicate_policy: DuplicatePolicy::EveryNth { stride: 10 },
1645                duplicate_cursor: 0,
1646                stats: stats.clone(),
1647                session_owners: HashMap::new(),
1648                max_forwarded_seq_by_session: HashMap::new(),
1649            },
1650        );
1651        // Bind exported handler ports so local handle posts carry
1652        // SEQ_INFO::Session instead of SeqInfo::Direct.
1653        let _bound_chaos_ref: ActorRef<ChaosActor> = chaos_handle.bind();
1654
1655        let barrier = Arc::new(Barrier::new(n));
1656        let mut sender_addrs: Vec<ActorAddr> = Vec::with_capacity(n);
1657        let mut sender_handles: Vec<ActorHandle<SenderActor>> = Vec::with_capacity(n);
1658        for sender_idx in 0..n {
1659            let h = proc.spawn_with_label(
1660                &format!("sender{sender_idx}"),
1661                SenderActor {
1662                    sender_idx: sender_idx as u32,
1663                },
1664            );
1665            sender_addrs.push(h.actor_addr().clone());
1666            sender_handles.push(h);
1667        }
1668        for h in &sender_handles {
1669            h.post(
1670                &client,
1671                Start {
1672                    count: k,
1673                    target: chaos_handle.clone(),
1674                    start_barrier: Some(barrier.clone()),
1675                    yield_between_frames: true,
1676                },
1677            );
1678        }
1679
1680        let received = match tokio::time::timeout(std::time::Duration::from_secs(5), done_rx).await
1681        {
1682            Ok(Ok(v)) => v,
1683            Ok(Err(_)) => panic!("done_rx sender dropped before sending; receiver crashed?"),
1684            Err(_) => {
1685                let trace = stats.forwarded_trace.lock().unwrap();
1686                panic!(
1687                    "timed out waiting for receiver; forwarded_trace ({} entries) = {:#?}",
1688                    trace.len(),
1689                    *trace,
1690                );
1691            }
1692        };
1693
1694        assert_received_protocol_correct(&received, n, k, &sender_addrs, &stats);
1695
1696        let selected = stats.duplicates_selected.load(AtomicOrdering::Acquire);
1697        let emitted = stats.late_duplicate_forwards.load(AtomicOrdering::Acquire);
1698        assert!(
1699            selected > 0,
1700            "expected duplicates_selected > 0, got {selected}: chaos policy never fired",
1701        );
1702        assert_eq!(
1703            emitted, selected,
1704            "duplicate selections did not match duplicate forwards: duplicates_selected={selected}, \
1705             late_duplicate_forwards={emitted}",
1706        );
1707        let oo = stats
1708            .out_of_order_original_forwards
1709            .load(AtomicOrdering::Acquire);
1710        assert!(
1711            oo > 0,
1712            "expected out_of_order_original_forwards > 0, got {oo}",
1713        );
1714    }
1715}