Skip to main content

kvbm_engine/leader/session/
state.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Session state types for the unified session model.
5//!
6//! This module provides the core state machine types:
7//! - [`ControlRole`]: Whether this session is controller, controllee, or neutral
8//! - [`AttachmentState`]: Whether a peer is attached
9//! - [`SessionPhase`]: The current operational phase of the session
10
11use serde::{Deserialize, Serialize};
12
13use crate::InstanceId;
14
15/// Control role in a session relationship.
16///
17/// Sessions can dynamically transition between roles:
18/// - Start as `Neutral` (independent, can initiate in either direction)
19/// - Become `Controller` when issuing commands to a peer
20/// - Become `Controllee` when executing commands from a peer
21///
22/// Control can be transferred bidirectionally via `YieldControl`/`AcquireControl`.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
24pub enum ControlRole {
25    /// Independent - can initiate control in either direction.
26    /// This is the initial state and the state after yielding control.
27    #[default]
28    Neutral,
29    /// Currently controlling a peer session (issues commands).
30    Controller,
31    /// Currently being controlled by a peer session (executes commands).
32    Controllee,
33}
34
35impl ControlRole {
36    /// Check if this role can issue control commands.
37    pub fn can_command(&self) -> bool {
38        matches!(self, ControlRole::Controller)
39    }
40
41    /// Check if this role should respond to control commands.
42    pub fn responds_to_commands(&self) -> bool {
43        matches!(self, ControlRole::Controllee)
44    }
45
46    /// Check if this role is neutral (can transition either way).
47    pub fn is_neutral(&self) -> bool {
48        matches!(self, ControlRole::Neutral)
49    }
50
51    /// Get the opposite role.
52    ///
53    /// - `Controller` ↔ `Controllee`
54    /// - `Neutral` → `Neutral` (no opposite)
55    pub fn opposite(&self) -> ControlRole {
56        match self {
57            ControlRole::Controller => ControlRole::Controllee,
58            ControlRole::Controllee => ControlRole::Controller,
59            ControlRole::Neutral => ControlRole::Neutral,
60        }
61    }
62}
63
64/// Attachment state - whether a peer is connected.
65///
66/// Valid state combinations:
67/// - `Neutral + Unattached`: Initial state, waiting for connection
68/// - `Neutral + Attached`: Post-yield state, peer still connected
69/// - `Controllee + Unattached`: Waiting for controller to attach
70/// - `Controllee + Attached`: Being actively controlled
71/// - `Controller + Attached`: Actively controlling (Controller + Unattached is invalid)
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
73pub enum AttachmentState {
74    /// No peer attached.
75    #[default]
76    Unattached,
77    /// Peer attached with the given instance ID.
78    Attached { peer: InstanceId },
79}
80
81impl AttachmentState {
82    /// Check if a peer is attached.
83    pub fn is_attached(&self) -> bool {
84        matches!(self, AttachmentState::Attached { .. })
85    }
86
87    /// Get the attached peer's instance ID if attached.
88    pub fn peer(&self) -> Option<InstanceId> {
89        match self {
90            AttachmentState::Attached { peer } => Some(*peer),
91            AttachmentState::Unattached => None,
92        }
93    }
94}
95
96/// Operational phase of a session.
97///
98/// Represents the lifecycle of block operations within a session:
99/// 1. `Searching` - Initial discovery/search phase
100/// 2. `Holding` - Blocks found and held, no staging yet
101/// 3. `Staging` - Transfer in progress (G3→G2, G4→G2, etc.)
102/// 4. `Ready` - All blocks in target tier, ready for transfer
103/// 5. `Complete` - Session completed successfully
104/// 6. `Failed` - Session failed or cancelled
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
106pub enum SessionPhase {
107    /// Initial search/discovery phase.
108    #[default]
109    Searching,
110    /// Blocks found and held, no staging started.
111    Holding,
112    /// Transfer/staging in progress (any direction).
113    Staging,
114    /// All blocks in target tier, ready for RDMA pull.
115    Ready,
116    /// Session completed successfully.
117    Complete,
118    /// Session failed or was cancelled.
119    Failed,
120}
121
122impl SessionPhase {
123    /// Check if the session is in a terminal state.
124    pub fn is_terminal(&self) -> bool {
125        matches!(self, SessionPhase::Complete | SessionPhase::Failed)
126    }
127
128    /// Check if the session is active (not terminal).
129    pub fn is_active(&self) -> bool {
130        !self.is_terminal()
131    }
132
133    /// Check if blocks are ready for transfer.
134    pub fn is_ready(&self) -> bool {
135        matches!(self, SessionPhase::Ready)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[test]
144    fn test_control_role_transitions() {
145        let role = ControlRole::Neutral;
146        assert!(role.is_neutral());
147        assert!(!role.can_command());
148        assert!(!role.responds_to_commands());
149
150        let role = ControlRole::Controller;
151        assert!(!role.is_neutral());
152        assert!(role.can_command());
153        assert!(!role.responds_to_commands());
154
155        let role = ControlRole::Controllee;
156        assert!(!role.is_neutral());
157        assert!(!role.can_command());
158        assert!(role.responds_to_commands());
159    }
160
161    #[test]
162    fn test_attachment_state() {
163        let state = AttachmentState::Unattached;
164        assert!(!state.is_attached());
165        assert!(state.peer().is_none());
166
167        let peer_id = InstanceId::new_v4();
168        let state = AttachmentState::Attached { peer: peer_id };
169        assert!(state.is_attached());
170        assert_eq!(state.peer(), Some(peer_id));
171    }
172
173    #[test]
174    fn test_session_phase() {
175        assert!(!SessionPhase::Searching.is_terminal());
176        assert!(!SessionPhase::Holding.is_terminal());
177        assert!(!SessionPhase::Staging.is_terminal());
178        assert!(!SessionPhase::Ready.is_terminal());
179        assert!(SessionPhase::Complete.is_terminal());
180        assert!(SessionPhase::Failed.is_terminal());
181
182        assert!(SessionPhase::Ready.is_ready());
183        assert!(!SessionPhase::Staging.is_ready());
184    }
185}