kvbm_engine/worker/
protocol.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 futures::future::{Either, Ready, ready};
6use serde::{Deserialize, Serialize};
7use std::{
8    pin::Pin,
9    task::{Context, Poll},
10};
11
12pub use crate::worker::{ImportMetadataResponseAwaiter, SerializedResponseAwaiter};
13pub use crate::{BlockId, SequenceHash};
14pub use kvbm_common::LogicalLayoutHandle;
15pub use kvbm_physical::manager::{LayoutHandle, SerializedLayout};
16
17pub struct SerializedLayoutResponse {
18    awaiter: Either<Ready<Result<SerializedLayout>>, SerializedResponseAwaiter>,
19}
20
21impl SerializedLayoutResponse {
22    pub fn ready(layout: SerializedLayout) -> Self {
23        Self {
24            awaiter: Either::Left(ready(Ok(layout))),
25        }
26    }
27
28    pub fn from_boxed(awaiter: SerializedResponseAwaiter) -> Self {
29        Self {
30            awaiter: Either::Right(awaiter),
31        }
32    }
33
34    pub fn could_yield(&self) -> bool {
35        matches!(self.awaiter, Either::Right(_))
36    }
37}
38
39impl std::future::IntoFuture for SerializedLayoutResponse {
40    type Output = Result<SerializedLayout>;
41    type IntoFuture = Either<Ready<Result<SerializedLayout>>, SerializedResponseAwaiter>;
42
43    fn into_future(self) -> Self::IntoFuture {
44        self.awaiter
45    }
46}
47
48pub struct ImportMetadataResponse {
49    awaiter: Either<Ready<Result<Vec<LayoutHandle>>>, ImportMetadataResponseAwaiter>,
50}
51
52impl ImportMetadataResponse {
53    pub fn ready(handles: Vec<LayoutHandle>) -> Self {
54        Self {
55            awaiter: Either::Left(ready(Ok(handles))),
56        }
57    }
58
59    pub fn from_boxed(awaiter: ImportMetadataResponseAwaiter) -> Self {
60        Self {
61            awaiter: Either::Right(awaiter),
62        }
63    }
64
65    pub fn could_yield(&self) -> bool {
66        matches!(self.awaiter, Either::Right(_))
67    }
68}
69
70impl std::future::IntoFuture for ImportMetadataResponse {
71    type Output = Result<Vec<LayoutHandle>>;
72    type IntoFuture = Either<Ready<Result<Vec<LayoutHandle>>>, ImportMetadataResponseAwaiter>;
73
74    fn into_future(self) -> Self::IntoFuture {
75        self.awaiter
76    }
77}
78
79/// Response type for `connect_remote` operations.
80///
81/// This type represents the completion state of a remote metadata import
82/// with handle mapping storage. Like other response types, it can be awaited.
83///
84/// For direct workers, this is typically ready immediately.
85/// For replicated workers, this aggregates multiple underlying imports.
86pub struct ConnectRemoteResponse {
87    awaiter: ConnectRemoteAwaiter,
88}
89
90pub enum ConnectRemoteAwaiter {
91    Ready(Ready<Result<()>>),
92    Event(::velo::EventAwaiter),
93}
94
95impl std::future::Future for ConnectRemoteAwaiter {
96    type Output = Result<()>;
97
98    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
99        match self.get_mut() {
100            Self::Ready(ready) => Pin::new(ready).poll(cx),
101            Self::Event(waiter) => Pin::new(waiter).poll(cx),
102        }
103    }
104}
105
106impl ConnectRemoteResponse {
107    /// Create a response that is already completed.
108    ///
109    /// This is used when the connect operation completes synchronously,
110    /// such as for DirectWorker with local metadata import.
111    pub fn ready() -> Self {
112        Self {
113            awaiter: ConnectRemoteAwaiter::Ready(ready(Ok(()))),
114        }
115    }
116
117    /// Create a response from an event waiter.
118    ///
119    /// This is used when the connect operation requires waiting for
120    /// multiple underlying operations to complete (e.g., ReplicatedWorker).
121    pub fn from_awaiter(awaiter: ::velo::EventAwaiter) -> Self {
122        Self {
123            awaiter: ConnectRemoteAwaiter::Event(awaiter),
124        }
125    }
126
127    /// Check if the response can yield the current task.
128    pub fn could_yield(&self) -> bool {
129        matches!(self.awaiter, ConnectRemoteAwaiter::Event(_))
130    }
131}
132
133impl std::future::IntoFuture for ConnectRemoteResponse {
134    type Output = Result<()>;
135    type IntoFuture = ConnectRemoteAwaiter;
136
137    fn into_future(self) -> Self::IntoFuture {
138        self.awaiter
139    }
140}
141
142/// Remote descriptor for transfer operations.
143#[derive(Serialize, Deserialize, Clone)]
144pub enum RemoteDescriptor {
145    Layout {
146        handle: LayoutHandle,
147        block_ids: Vec<BlockId>,
148    },
149    Object {
150        keys: Vec<SequenceHash>,
151    },
152}
153
154/// Configuration sent from leader to workers for G2/G3/G4 layout creation.
155///
156/// This message is sent via Nova RPC during Phase 3 coordination.
157/// Workers use this to create additional cache tiers beyond G1 (GPU KV).
158#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct LeaderLayoutConfig {
160    /// Leader provided rank of this worker
161    ///
162    /// The Connector framework provides us with an ordered list of workers. To ensure
163    /// leaders and workers are all in-sync on this information, the leader will send
164    /// each worker the rank provided by the Connector framework.
165    pub rank: usize,
166
167    /// Number of host/pinned blocks for G2 tier.
168    pub host_block_count: usize,
169
170    /// Number of disk blocks for G3 tier (None = no disk tier).
171    pub disk_block_count: Option<usize>,
172
173    /// Object storage configuration for G4 tier (None = no object tier).
174    ///
175    /// When present, workers should instantiate object clients for storing
176    /// blocks in external object storage (S3/MinIO).
177    #[serde(default)]
178    pub object: Option<kvbm_config::ObjectConfig>,
179
180    /// Parallelism mode for this worker.
181    ///
182    /// When `ReplicatedData` and rank > 0, the worker skips G2/G3 creation
183    /// since only rank 0 has host/disk storage in replicated mode.
184    #[serde(default)]
185    pub parallelism: kvbm_config::ParallelismMode,
186}
187
188/// Worker's response after configuring additional layouts (G2, G3).
189///
190/// Returned in response to a `LeaderLayoutConfig` request.
191#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct WorkerLayoutResponse {
193    /// Full exported metadata including all registered layouts (G1, G2, G3).
194    pub metadata: SerializedLayout,
195
196    /// Which logical layouts were successfully created in this operation.
197    pub created_layouts: Vec<LogicalLayoutHandle>,
198}
kvbm_engine/worker/protocol.rs

kvbm_engine/worker/
protocol.rs
