hyperactor_mesh/

resource.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! This modules defines a set of common message types used for managing resources
//! in hyperactor meshes.
//!
//! # Rank-status positioning (`RSP-*`)
//!
//! Status queries ([`GetRankStatus`], [`WaitRankStatus`]) reply with a
//! [`StatusOverlay`] — a sparse, rank-indexed patch the caller merges into an
//! accumulator over its own view. Positioning that overlay is the contract this
//! registry names; each invariant is cited at the sites that must uphold it.
//!
//! - **RSP-1** — an overlay is positioned at the recipient's dense rank in the
//!   *consuming* view (the one the caller accumulates over), never the rank the
//!   resource was first created under. Reusing the creation rank mis-attributes
//!   or drops the reply over a renumbered or sliced view.
//! - **RSP-2** — the positioning rank travels as the [`Rank`] message field. Cast
//!   sites leave it `Rank::default()` and the cast layer late-binds it in transit
//!   with the delivered dense rank; direct (non-cast) delivery supplies it
//!   explicitly with `Rank::new(view_rank)`.
//! - **RSP-3** — a deferred waiter (a [`WaitRankStatus`] below its threshold, or
//!   one registered before its resource exists) retains its registration rank
//!   until it is flushed; the flush never rewrites it to a creation rank.
//! - **RSP-4** — when a handler reports absence *immediately*, it returns an
//!   empty overlay; the message's [`Rank`] is never overloaded as an absence
//!   signal. `WaitRankStatus` may instead defer when later creation is expected
//!   (RSP-3), replying once the resource exists.
//!
//! `ActorState` streaming supervision must satisfy the analogous positioning
//! contract; extending this registry to cover `ActorState` is future work.

pub mod mesh;

use core::slice::GetDisjointMutIndex as _;
use std::collections::HashMap;
use std::fmt;
use std::fmt::Debug;
use std::hash::Hash;
use std::mem::replace;
use std::mem::take;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ops::Range;
use std::time::Duration;

use enum_as_inner::EnumAsInner;
use hyperactor::HandleClient;
use hyperactor::Handler;
use hyperactor::IdleFlushPortRef;
use hyperactor::PortRef;
use hyperactor::RefClient;
use hyperactor::RemoteMessage;
use hyperactor::mailbox::PortReceiver;
use hyperactor_config::attrs::Attrs;
use ndslice::Region;
use ndslice::ViewExt;
use serde::Deserialize;
use serde::Serialize;
use typeuri::Named;

use crate::StatusOverlay;
use crate::ValueMesh;
use crate::bootstrap;
use crate::bootstrap::BootstrapCommand;
use crate::bootstrap::ProcBind;
use crate::host_mesh::host_agent::ProcState;
use crate::mesh_id::ResourceId;
use crate::proc_agent::ActorSpec;
use crate::proc_agent::ActorState;

/// The current lifecycle status of a resource.
#[derive(
    Clone,
    Debug,
    Serialize,
    Deserialize,
    Named,
    PartialOrd,
    Ord,
    PartialEq,
    Eq,
    Hash,
    EnumAsInner,
    strum::Display
)]
pub enum Status {
    /// The resource does not exist.
    NotExist,
    /// The resource is being created.
    Initializing,
    /// The resource is running.
    Running,
    /// The resource is being stopped.
    Stopping,
    /// The resource is stopped.
    Stopped,
    /// The resource has failed, with an error message.
    #[strum(to_string = "Failed({0})")]
    Failed(String),
    /// The resource has been declared failed after a timeout.
    #[strum(to_string = "Timeout({0:?})")]
    Timeout(Duration),
    /// The resource exists but its status is not known.
    Unknown,
}

impl Status {
    /// Returns whether the status is a terminating status (includes `Stopping`).
    pub fn is_terminating(&self) -> bool {
        matches!(
            self,
            Status::Stopping | Status::Stopped | Status::Failed(_) | Status::Timeout(_)
        )
    }

    /// Tells whether the status represents a failure. A failure is both terminating
    /// (the resource is not running), but also means abnormal exit (the resource
    /// did not stop cleanly).
    pub fn is_failure(&self) -> bool {
        matches!(self, Self::Failed(_) | Self::Timeout(_))
    }

    /// Returns whether the status is fully terminal (the resource has
    /// stopped, failed, or timed out — but NOT merely `Stopping`).
    pub fn is_terminated(&self) -> bool {
        matches!(
            self,
            Status::Stopped | Status::Failed(_) | Status::Timeout(_)
        )
    }

    pub fn is_healthy(&self) -> bool {
        matches!(self, Status::Initializing | Status::Running)
    }

    /// Ensure this status is at least as terminal as `floor`.
    ///
    /// If `floor` is a terminating status (Stopping, Stopped, Failed,
    /// Timeout) and `self` is not, returns `floor`. Otherwise returns
    /// `self` unchanged. This is used to prevent a child resource
    /// from appearing healthier than its parent.
    pub fn clamp_min(self, floor: Status) -> Status {
        if floor.is_terminating() && !self.is_terminating() {
            floor
        } else {
            self
        }
    }
}

impl From<bootstrap::ProcStatus> for Status {
    fn from(status: bootstrap::ProcStatus) -> Self {
        use bootstrap::ProcStatus;
        match status {
            ProcStatus::Starting => Status::Initializing,
            ProcStatus::Running { .. } | ProcStatus::Ready { .. } => Status::Running,
            ProcStatus::Stopping { .. } => Status::Stopping,
            ProcStatus::Stopped { .. } => Status::Stopped,
            ProcStatus::Failed { reason } => Status::Failed(reason),
            ProcStatus::Killed { .. } => Status::Failed(format!("{}", status)),
        }
    }
}

impl From<crate::host::LocalProcStatus> for Status {
    fn from(status: crate::host::LocalProcStatus) -> Self {
        match status {
            crate::host::LocalProcStatus::Stopping => Status::Stopping,
            crate::host::LocalProcStatus::Stopped => Status::Stopped,
        }
    }
}

/// Data type used to communicate ranks.
/// Serialized as a typed multipart part so comm actors can replace instances
/// with the delivered rank while the message is in transit.
#[derive(Clone, Debug, Named, PartialEq, Eq, Default)]
pub struct Rank(pub Option<usize>);
wirevalue::register_type!(Rank);

/// Serialized representation for [`Rank`] multipart parts.
#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
pub struct RankRepr(pub Option<usize>);

impl TryFrom<&Rank> for RankRepr {
    type Error = serde_multipart::Error;

    fn try_from(rank: &Rank) -> serde_multipart::Result<Self> {
        Ok(Self(rank.0))
    }
}

impl TryFrom<RankRepr> for Rank {
    type Error = serde_multipart::Error;

    fn try_from(repr: RankRepr) -> serde_multipart::Result<Self> {
        Ok(Self(repr.0))
    }
}

serde_multipart::part_codec! {
    impl Rank
    {
        type Repr = RankRepr;
    }
}

impl Rank {
    /// Create a new rank with the provided value.
    pub fn new(rank: usize) -> Self {
        Self(Some(rank))
    }

    /// Unwrap the rank; panics if not set.
    pub fn unwrap(&self) -> usize {
        self.0.unwrap()
    }
}

/// Get the status of a resource.
///
/// Delivered either by cast to every rank or by direct point-to-point send to a
/// chosen recipient. Each recipient replies with a sparse status **overlay**.
/// The comm reducer merges overlays (right-wins) and the accumulator applies
/// them to produce **full StatusMesh snapshots** on the receiver side.
#[derive(
    Clone,
    Debug,
    Serialize,
    Deserialize,
    Named,
    Handler,
    HandleClient,
    RefClient
)]
pub struct GetRankStatus {
    /// The resource identifier.
    pub id: ResourceId,
    /// The rank at which the recipient positions its reply overlay, in the
    /// caller's view (RSP-1, RSP-2). Cast sites leave this `Rank::default()` and
    /// the cast layer rewrites it in transit with the delivered dense rank;
    /// direct (non-cast) callers set it explicitly with `Rank::new(...)`.
    pub rank: Rank,
    /// Sparse status updates (overlays) from a rank.
    pub reply: IdleFlushPortRef<StatusOverlay>,
}

/// Like [`GetRankStatus`], but the handler defers its reply until the
/// resource's status is >= `min_status`. This avoids the race where
/// the caller sees `Stopping` before the process has actually exited.
#[derive(
    Clone,
    Debug,
    Serialize,
    Deserialize,
    Named,
    Handler,
    HandleClient,
    RefClient
)]
pub struct WaitRankStatus {
    /// The resource identifier.
    pub id: ResourceId,
    /// The rank at which the recipient positions its reply overlay, in the
    /// caller's view (RSP-1, RSP-2). Cast sites leave this `Rank::default()` and
    /// the cast layer rewrites it in transit with the delivered dense rank;
    /// direct (non-cast) callers set it explicitly with `Rank::new(...)`. A
    /// deferred waiter retains this rank until it is flushed (RSP-3).
    pub rank: Rank,
    /// The minimum status the caller wants to observe.
    /// The handler will not reply until the resource's status
    /// is >= this threshold.
    pub min_status: Status,
    /// Sparse status updates (overlays) from a rank.
    pub reply: IdleFlushPortRef<StatusOverlay>,
}

/// Collect an accumulated [`ValueMesh<T>`] from `rx` until every rank has been
/// reported or `max_idle_time` elapses with no update.
///
/// `reported` decides whether a rank's current value counts as "reported" (e.g.
/// moved off a `NotExist`/`Timeout` placeholder). The accumulator emits the full
/// mesh over the target region, so completion is "every cell reported". Returns
/// `Ok(mesh)` once all ranks are reported, or `Err(mesh)` carrying the latest
/// snapshot (or `fallback` if nothing arrived) on idle timeout / channel close.
pub async fn wait_mesh<T>(
    mut rx: PortReceiver<ValueMesh<T>>,
    max_idle_time: Duration,
    fallback: ValueMesh<T>,
    reported: impl Fn(&T) -> bool,
) -> Result<ValueMesh<T>, ValueMesh<T>>
where
    T: Send + Sync + Clone + 'static,
{
    let mut alarm = hyperactor::time::Alarm::new();
    alarm.arm(max_idle_time);

    // Latest-wins snapshot; `fallback` stands in until the first update arrives.
    let mut snapshot = fallback;

    loop {
        let mut sleeper = alarm.sleeper();
        tokio::select! {
            _ = sleeper.sleep() => return Err(snapshot),
            next = rx.recv() => {
                match next {
                    Ok(mesh) => { snapshot = mesh; }   // latest-wins snapshot
                    Err(_)   => return Err(snapshot),
                }
            }
        }

        alarm.arm(max_idle_time);

        // Short-circuit: done as soon as no rank is still unreported.
        if snapshot.values().all(|v| reported(&v)) {
            break Ok(snapshot);
        }
    }
}

impl GetRankStatus {
    pub async fn wait(
        rx: PortReceiver<crate::StatusMesh>,
        num_ranks: usize,
        max_idle_time: Duration,
        region: Region, // used only for fallback
    ) -> Result<crate::StatusMesh, crate::StatusMesh> {
        debug_assert_eq!(region.num_ranks(), num_ranks, "region/num_ranks mismatch");

        let fallback = crate::StatusMesh::from_single(region, crate::resource::Status::NotExist);
        wait_mesh(rx, max_idle_time, fallback, |s| {
            !matches!(s, crate::resource::Status::NotExist)
        })
        .await
    }
}

/// The state of a resource.
#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq, Handler)]
pub struct State<S> {
    /// The resource identifier.
    pub id: ResourceId,
    /// Its status.
    pub status: Status,
    /// Optionally, a resource-defined state.
    pub state: Option<S>,
    /// Monotonic generation counter for last-writer-wins ordering.
    pub generation: u64,
    /// Wall-clock timestamp for debugging context.
    pub timestamp: std::time::SystemTime,
}
wirevalue::register_type!(State<ActorState>);
wirevalue::register_type!(State<ProcState>);

impl<S: Serialize> fmt::Display for State<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Use serde_json to serialize the struct to a compact JSON string
        match serde_json::to_string(self) {
            Ok(json) => write!(f, "{}", json),
            Err(e) => write!(f, "<state: serde_json error: {}>", e),
        }
    }
}

/// Create or update a resource according to a spec.
#[derive(
    Debug,
    Clone,
    Serialize,
    Deserialize,
    Named,
    Handler,
    HandleClient,
    RefClient
)]
pub struct CreateOrUpdate<S> {
    /// The resource identifier.
    pub id: ResourceId,
    /// The rank of the resource, when available.
    pub rank: Rank,
    /// The specification of the resource.
    pub spec: S,
}
wirevalue::register_type!(CreateOrUpdate<ProcSpec>);
wirevalue::register_type!(CreateOrUpdate<ActorSpec>);

/// Stop a resource according to a spec.
#[derive(
    Debug,
    Clone,
    Serialize,
    Deserialize,
    Named,
    Handler,
    HandleClient,
    RefClient
)]
pub struct Stop {
    /// The resource identifier.
    pub id: ResourceId,
    /// The reason for stopping the resource.
    pub reason: String,
}
wirevalue::register_type!(Stop);

/// Stop all resources owned by the receiver of this message.
/// No reply, this just issues the stop command.
/// Use GetRankStatus to determine if it has successfully stopped.
#[derive(
    Debug,
    Clone,
    Serialize,
    Deserialize,
    Named,
    Handler,
    HandleClient,
    RefClient
)]
pub struct StopAll {
    /// The reason for stopping.
    pub reason: String,
}
wirevalue::register_type!(StopAll);

/// Retrieve the current state of the resource.
#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
pub struct GetState<S> {
    /// The resource identifier.
    pub id: ResourceId,
    /// A reply containing the state.
    #[reply]
    pub reply: PortRef<State<S>>,
}
wirevalue::register_type!(GetState<ProcState>);
wirevalue::register_type!(GetState<ActorState>);

impl<S> Clone for GetState<S>
where
    S: RemoteMessage,
{
    fn clone(&self) -> Self {
        Self {
            id: self.id.clone(),
            reply: self.reply.clone(),
        }
    }
}

/// Same as GetState, but additionally tells the receiver that the owner is still alive.
/// If the receiver does not receive this message for a while, it might assume the owner is dead.
#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
pub struct KeepaliveGetState<S> {
    /// The time at which the actor should be considered expired if no further
    /// keepalive is received.
    pub expires_after: std::time::SystemTime,
    pub get_state: GetState<S>,
}
wirevalue::register_type!(KeepaliveGetState<ProcState>);
wirevalue::register_type!(KeepaliveGetState<ActorState>);

impl<S> Clone for KeepaliveGetState<S>
where
    S: RemoteMessage,
{
    fn clone(&self) -> Self {
        Self {
            expires_after: self.expires_after,
            get_state: self.get_state.clone(),
        }
    }
}

/// Subscribe to streaming state updates for a named resource.
/// The subscriber port will receive `RankedState<S>` whenever the resource's
/// state changes. The current state is sent immediately upon subscription.
#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
#[serde(bound(serialize = "S: Named", deserialize = "S: Named"))]
pub struct StreamState<S> {
    /// The resource identifier.
    pub id: ResourceId,
    /// The recipient's rank in the subscriber's view. Cast callers leave this
    /// unset so the comm layer can fill it for each recipient; direct callers
    /// set it explicitly.
    pub subscriber_rank: Rank,
    /// A streaming port that will receive state updates.
    pub subscriber: PortRef<RankedState<S>>,
}
wirevalue::register_type!(StreamState<ActorState>);
wirevalue::register_type!(StreamState<ProcState>);

impl<S> Clone for StreamState<S>
where
    S: RemoteMessage,
{
    fn clone(&self) -> Self {
        Self {
            id: self.id.clone(),
            subscriber_rank: self.subscriber_rank.clone(),
            subscriber: self.subscriber.clone(),
        }
    }
}

/// A state update positioned in the subscriber's current view.
#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq, Handler)]
pub struct RankedState<S> {
    /// The resource's rank in the consuming view.
    pub rank: Rank,
    /// The state observed at that rank.
    pub state: State<S>,
}
wirevalue::register_type!(RankedState<ActorState>);
wirevalue::register_type!(RankedState<ProcState>);

/// List the set of resources managed by the controller.
#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
pub struct List {
    /// List of resource names managed by this controller.
    #[reply]
    pub reply: PortRef<Vec<ResourceId>>,
}
wirevalue::register_type!(List);

/// A trait that bundles a set of types that together define a resource.
pub trait Resource {
    /// The spec specification for this resource.
    type Spec: typeuri::Named
        + Serialize
        + for<'de> Deserialize<'de>
        + Send
        + Sync
        + std::fmt::Debug;

    /// The state for this resource.
    type State: typeuri::Named
        + Serialize
        + for<'de> Deserialize<'de>
        + Send
        + Sync
        + std::fmt::Debug;
}

// A behavior defining the interface for a mesh controller.
hyperactor::behavior!(
    Controller<R: Resource>,
    CreateOrUpdate<R::Spec>,
    GetState<R::State>,
    Stop,
);

/// RankedValues compactly represents rank-indexed values of type T.
/// It stores contiguous values in a set of intervals; thus it is
/// efficient and compact when the cardinality of T-typed values is
/// low.
#[derive(Debug, Clone, Named, Serialize, Deserialize)]
pub struct RankedValues<T> {
    intervals: Vec<(Range<usize>, T)>,
}

impl<T: PartialEq> PartialEq for RankedValues<T> {
    fn eq(&self, other: &Self) -> bool {
        self.intervals == other.intervals
    }
}

impl<T: Eq> Eq for RankedValues<T> {}

impl<T> Default for RankedValues<T> {
    fn default() -> Self {
        Self {
            intervals: Vec::new(),
        }
    }
}

impl<T> RankedValues<T> {
    /// Iterate over contiguous rank intervals of values.
    pub fn iter(&self) -> impl Iterator<Item = &(Range<usize>, T)> + '_ {
        self.intervals.iter()
    }

    /// The (set) rank of the RankedValues is the number of values stored with
    /// rank less than `value`.
    pub fn rank(&self, value: usize) -> usize {
        self.iter()
            .take_while(|(ranks, _)| ranks.start <= value)
            .map(|(ranks, _)| ranks.end.min(value) - ranks.start)
            .sum()
    }
}

impl<T: Clone> RankedValues<T> {
    pub fn materialized_iter(&self, until: usize) -> impl Iterator<Item = &T> + '_ {
        assert_eq!(self.rank(until), until, "insufficient rank");
        self.iter()
            .flat_map(|(range, value)| std::iter::repeat_n(value, range.end - range.start))
    }
}

impl<T: Hash + Eq + Clone> RankedValues<T> {
    /// Invert this ranked values into a [`ValuesByRank<T>`].
    pub fn invert(&self) -> ValuesByRank<T> {
        let mut inverted: HashMap<T, Vec<Range<usize>>> = HashMap::new();
        for (range, value) in self.iter() {
            inverted
                .entry(value.clone())
                .or_default()
                .push(range.clone());
        }
        ValuesByRank { values: inverted }
    }
}

impl<T: Eq + Clone> RankedValues<T> {
    /// Merge `other` into this set of ranked values. Values in `other` that overlap
    /// with `self` take prededence.
    ///
    /// This currently uses a simple algorithm that merges the full set of RankedValues.
    /// This remains efficient when the cardinality of T-typed values is low. However,
    /// it does not efficiently merge high cardinality value sets. Consider using interval
    /// trees or bitmap techniques like Roaring Bitmaps in these cases.
    pub fn merge_from(&mut self, other: Self) {
        let mut left_iter = take(&mut self.intervals).into_iter();
        let mut right_iter = other.intervals.into_iter();

        let mut left = left_iter.next();
        let mut right = right_iter.next();

        while left.is_some() && right.is_some() {
            let (left_ranks, left_value) = left.as_mut().unwrap();
            let (right_ranks, right_value) = right.as_mut().unwrap();

            if left_ranks.is_overlapping(right_ranks) {
                if left_value == right_value {
                    let ranks = left_ranks.start.min(right_ranks.start)..right_ranks.end;
                    let (_, value) = replace(&mut right, right_iter.next()).unwrap();
                    left_ranks.start = ranks.end;
                    if left_ranks.is_empty() {
                        left = left_iter.next();
                    }
                    self.append(ranks, value);
                } else if left_ranks.start < right_ranks.start {
                    let ranks = left_ranks.start..right_ranks.start;
                    left_ranks.start = ranks.end;
                    // TODO: get rid of clone
                    self.append(ranks, left_value.clone());
                } else {
                    let (ranks, value) = replace(&mut right, right_iter.next()).unwrap();
                    left_ranks.start = ranks.end;
                    if left_ranks.is_empty() {
                        left = left_iter.next();
                    }
                    self.append(ranks, value);
                }
            } else if left_ranks.start < right_ranks.start {
                let (ranks, value) = replace(&mut left, left_iter.next()).unwrap();
                self.append(ranks, value);
            } else {
                let (ranks, value) = replace(&mut right, right_iter.next()).unwrap();
                self.append(ranks, value);
            }
        }

        while let Some((left_ranks, left_value)) = left {
            self.append(left_ranks, left_value);
            left = left_iter.next();
        }
        while let Some((right_ranks, right_value)) = right {
            self.append(right_ranks, right_value);
            right = right_iter.next();
        }
    }

    /// Merge the contents of this RankedValues into another RankedValues.
    pub fn merge_into(self, other: &mut Self) {
        other.merge_from(self);
    }

    fn append(&mut self, range: Range<usize>, value: T) {
        if let Some(last) = self.intervals.last_mut()
            && last.0.end == range.start
            && last.1 == value
        {
            last.0.end = range.end;
        } else {
            self.intervals.push((range, value));
        }
    }
}

impl RankedValues<Status> {
    pub fn first_terminating(&self) -> Option<(usize, Status)> {
        self.intervals
            .iter()
            .find(|(_, status)| status.is_terminating())
            .map(|(range, status)| (range.start, status.clone()))
    }

    pub fn first_failed(&self) -> Option<(usize, Status)> {
        self.intervals
            .iter()
            .find(|(_, status)| matches!(status, Status::Failed(_) | Status::Timeout(_)))
            .map(|(range, status)| (range.start, status.clone()))
    }
}

impl<T> From<(usize, T)> for RankedValues<T> {
    fn from((rank, value): (usize, T)) -> Self {
        Self {
            intervals: vec![(rank..rank + 1, value)],
        }
    }
}

impl<T> From<(Range<usize>, T)> for RankedValues<T> {
    fn from((range, value): (Range<usize>, T)) -> Self {
        Self {
            intervals: vec![(range, value)],
        }
    }
}

/// An inverted index of RankedValues, providing all ranks for
/// which each unique T-typed value appears.
#[derive(Clone, Debug)]
pub struct ValuesByRank<T> {
    values: HashMap<T, Vec<Range<usize>>>,
}

impl<T: Eq + Hash> PartialEq for ValuesByRank<T> {
    fn eq(&self, other: &Self) -> bool {
        self.values == other.values
    }
}

impl<T: Eq + Hash> Eq for ValuesByRank<T> {}

impl<T: fmt::Display> fmt::Display for ValuesByRank<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut first_value = true;
        for (value, ranges) in self.iter() {
            if first_value {
                first_value = false;
            } else {
                write!(f, ";")?;
            }
            write!(f, "{}=", value)?;
            let mut first_range = true;
            for range in ranges.iter() {
                if first_range {
                    first_range = false;
                } else {
                    write!(f, ",")?;
                }
                write!(f, "{}..{}", range.start, range.end)?;
            }
        }
        Ok(())
    }
}

impl<T> Deref for ValuesByRank<T> {
    type Target = HashMap<T, Vec<Range<usize>>>;

    fn deref(&self) -> &Self::Target {
        &self.values
    }
}

impl<T> DerefMut for ValuesByRank<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.values
    }
}

/// Enabled for test only because we have to guarantee that the input
/// iterator is well-formed.
#[cfg(test)]
impl<T> FromIterator<(Range<usize>, T)> for RankedValues<T> {
    fn from_iter<I: IntoIterator<Item = (Range<usize>, T)>>(iter: I) -> Self {
        Self {
            intervals: iter.into_iter().collect(),
        }
    }
}

/// Spec for a host mesh agent to use when spawning a new proc.
#[derive(Clone, Debug, Serialize, Deserialize, Named, Default)]
pub(crate) struct ProcSpec {
    /// Config values to set on the spawned proc's global config,
    /// at the `ClientOverride` layer.
    pub(crate) client_config_override: Attrs,
    /// Optional per-process CPU/NUMA binding configuration.
    pub(crate) proc_bind: Option<ProcBind>,
    /// Optional bootstrap command override. When set, this command is used
    /// to spawn the proc instead of the host agent's default bootstrap command.
    pub(crate) bootstrap_command: Option<BootstrapCommand>,
    /// The id of the HostMesh that owns this proc. Used by
    /// `DrainHost` to selectively drain only procs belonging to a
    /// specific mesh.
    pub(crate) host_mesh_id: Option<crate::mesh_id::HostMeshId>,
    /// The id of the ProcMesh that owns this proc. A host agent can hold procs
    /// from several proc meshes, so this lets per-mesh queries (e.g.
    /// `StreamState`) scope to a single mesh rather than every proc on the host.
    pub(crate) proc_mesh_id: Option<crate::mesh_id::ProcMeshId>,
}
wirevalue::register_type!(ProcSpec);

#[cfg(test)]
mod tests {
    use hyperactor::port::Port;

    use super::*;

    #[test]
    fn handler_ports_are_distinct_for_resource_messages() {
        assert_ne!(
            Port::handler::<CreateOrUpdate<ProcSpec>>(),
            Port::handler::<Stop>(),
        );
    }

    #[test]
    fn test_ranked_values_merge() {
        #[derive(PartialEq, Debug, Eq, Clone)]
        enum Side {
            Left,
            Right,
            Both,
        }
        use Side::Both;
        use Side::Left;
        use Side::Right;

        let mut left: RankedValues<Side> = [
            (0..10, Left),
            (15..20, Left),
            (30..50, Both),
            (60..70, Both),
        ]
        .into_iter()
        .collect();

        let right: RankedValues<Side> = [
            (9..12, Right),
            (25..30, Right),
            (30..40, Both),
            (40..50, Right),
            (50..60, Both),
        ]
        .into_iter()
        .collect();

        left.merge_from(right);
        assert_eq!(
            left.iter().cloned().collect::<Vec<_>>(),
            vec![
                (0..9, Left),
                (9..12, Right),
                (15..20, Left),
                (25..30, Right),
                (30..40, Both),
                (40..50, Right),
                // Merge consecutive:
                (50..70, Both)
            ]
        );

        assert_eq!(left.rank(5), 5);
        assert_eq!(left.rank(10), 10);
        assert_eq!(left.rank(16), 13);
        assert_eq!(left.rank(70), 62);
        assert_eq!(left.rank(100), 62);
    }

    #[test]
    fn test_equality() {
        assert_eq!(
            RankedValues::from((0..10, 123)),
            RankedValues::from((0..10, 123))
        );
        assert_eq!(
            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
            RankedValues::from((0..10, Status::Failed("foo".to_string()))),
        );
    }

    #[test]
    fn test_default_through_merging() {
        let values: RankedValues<usize> =
            [(0..10, 1), (15..20, 1), (30..50, 1)].into_iter().collect();

        let mut default = RankedValues::from((0..50, 0));
        default.merge_from(values);

        assert_eq!(
            default.iter().cloned().collect::<Vec<_>>(),
            vec![
                (0..10, 1),
                (10..15, 0),
                (15..20, 1),
                (20..30, 0),
                (30..50, 1)
            ]
        );
    }
}