kvbm_engine/leader/
onboarding.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::Result;
5use tokio::sync::mpsc;
6
7use super::session::SessionId;
8use super::types::StagingMode;
9
10/// Status of an onboarding operation.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum OnboardingStatus {
13    /// Searching for blocks (local or remote).
14    Searching,
15
16    /// Holding blocks without staging (StagingMode::Hold).
17    /// Provides location breakdown for cost analysis.
18    /// - `local_g2`: number of blocks in local G2 (ready to use)
19    /// - `local_g3`: number of blocks in local G3 (needs local staging)
20    /// - `remote_g2`: number of blocks in remote G2 (needs RDMA pull)
21    /// - `remote_g3`: number of blocks in remote G3 (needs remote staging + RDMA)
22    /// - `pending_g4`: number of blocks with G4 load in progress
23    /// - `loaded_g4`: number of blocks successfully loaded from G4 (included in local_g2)
24    /// - `failed_g4`: number of blocks that failed to load from G4
25    Holding {
26        local_g2: usize,
27        local_g3: usize,
28        remote_g2: usize,
29        remote_g3: usize,
30        pending_g4: usize,
31        loaded_g4: usize,
32        failed_g4: usize,
33    },
34
35    /// Preparing: staging G3→G2 (StagingMode::Prepare or Full).
36    /// - `matched`: total number of blocks matched during search
37    /// - `staging_local`: number of local G3→G2 transfers in progress
38    /// - `staging_remote`: number of remote G3→G2 transfers in progress
39    Preparing {
40        matched: usize,
41        staging_local: usize,
42        staging_remote: usize,
43    },
44
45    /// Prepared: all blocks in G2, session still alive (StagingMode::Prepare).
46    /// - `local_g2`: number of blocks in local G2
47    /// - `remote_g2`: number of blocks in remote G2 instances
48    Prepared { local_g2: usize, remote_g2: usize },
49
50    /// Staging: full mode with RDMA pulls (StagingMode::Full).
51    /// - `matched`: total number of blocks matched
52    /// - `staging_local`: local G3→G2 in progress
53    /// - `staging_remote`: remote G3→G2 in progress
54    /// - `pulling`: remote G2→local G2 (RDMA) in progress
55    Staging {
56        matched: usize,
57        staging_local: usize,
58        staging_remote: usize,
59        pulling: usize,
60    },
61
62    /// Operation complete - all blocks are in initiator's G2 (StagingMode::Full).
63    /// Or terminal state for Hold/Prepare modes.
64    /// - `matched`: total number of blocks in local G2
65    Complete { matched_blocks: usize },
66}
67
68/// Control commands for managing live sessions.
69#[derive(Debug)]
70pub(crate) enum SessionControl {
71    /// Trigger prepare operation (Hold → Prepare): stage all G3→G2
72    Prepare,
73
74    /// Trigger pull operation (Prepare → Full): RDMA pull remote G2→local G2
75    Pull,
76
77    /// Cancel session and release all blocks
78    Cancel,
79
80    /// Shutdown session (normal completion)
81    Shutdown,
82}
83
84/// Handle to a live onboarding session for deferred operations.
85///
86/// Only available for StagingMode::Hold and StagingMode::Prepare.
87#[derive(Debug)]
88pub struct SessionHandle {
89    session_id: SessionId,
90    mode: StagingMode,
91    control_tx: mpsc::Sender<SessionControl>,
92}
93
94impl SessionHandle {
95    pub(crate) fn new(
96        session_id: SessionId,
97        mode: StagingMode,
98        control_tx: mpsc::Sender<SessionControl>,
99    ) -> Self {
100        Self {
101            session_id,
102            mode,
103            control_tx,
104        }
105    }
106
107    /// Get the session ID.
108    pub fn session_id(&self) -> SessionId {
109        self.session_id
110    }
111
112    /// Get the current staging mode.
113    pub fn mode(&self) -> StagingMode {
114        self.mode
115    }
116
117    /// Trigger G3→G2 staging on all instances (Hold → Prepare).
118    ///
119    /// The server validates that the session is in Hold mode before processing.
120    /// After this completes, the session transitions to Prepare mode internally.
121    pub async fn prepare(&self) -> Result<()> {
122        self.control_tx
123            .send(SessionControl::Prepare)
124            .await
125            .map_err(|_| anyhow::anyhow!("session task has exited"))
126    }
127
128    /// Trigger RDMA pull from remote G2→local G2 (Prepare → Complete).
129    ///
130    /// The server validates that the session is in Prepare mode before processing.
131    /// After this completes, the session transitions to Complete status.
132    pub async fn pull(&self) -> Result<()> {
133        self.control_tx
134            .send(SessionControl::Pull)
135            .await
136            .map_err(|_| anyhow::anyhow!("session task has exited"))
137    }
138
139    /// Cancel session and release all held blocks.
140    pub async fn cancel(&self) -> Result<()> {
141        self.control_tx
142            .send(SessionControl::Cancel)
143            .await
144            .map_err(|_| anyhow::anyhow!("session task has exited"))
145    }
146
147    /// Shutdown session (used internally).
148    #[expect(dead_code)]
149    pub(crate) async fn shutdown(&self) -> Result<()> {
150        self.control_tx
151            .send(SessionControl::Shutdown)
152            .await
153            .map_err(|_| anyhow::anyhow!("session task has exited"))
154    }
155}
kvbm_engine/leader/onboarding.rs

kvbm_engine/leader/
onboarding.rs
