Skip to main content

nodedb_cluster/
vshard_handler.rs

1// SPDX-License-Identifier: BUSL-1.1
2
3//! Shard-side handler for incoming VShardEnvelope messages.
4//!
5//! When a remote node sends a VShardEnvelope via QUIC (wrapped in
6//! `RaftRpc::VShardEnvelope`), the transport server routes it here.
7//! This handler dispatches based on `VShardMessageType` to the
8//! appropriate engine handler:
9//!
10//! - Timeseries scatter → local scan + partial aggregate response
11//! - Retention command → local retention enforcement
12//! - S3 archive command → local archive execution
13
14use crate::wire::{VShardEnvelope, VShardMessageType};
15
16/// Result of handling a VShardEnvelope.
17pub enum HandleResult {
18    /// Send a response envelope back to the coordinator.
19    Response(VShardEnvelope),
20    /// No response needed (fire-and-forget message).
21    NoResponse,
22    /// Handler error.
23    Error(String),
24}
25
26/// Trait for handling incoming VShardEnvelope messages on a shard.
27///
28/// Implemented by the main binary, which has access to the Data Plane
29/// engines (timeseries scan, graph traversal, etc.).
30pub trait VShardHandler: Send + Sync + 'static {
31    /// Handle an incoming VShardEnvelope and optionally produce a response.
32    fn handle_vshard_envelope(
33        &self,
34        envelope: VShardEnvelope,
35    ) -> impl std::future::Future<Output = HandleResult> + Send;
36}
37
38/// Default handler that dispatches based on message type.
39///
40/// The main binary implements `VShardHandler` and uses this function
41/// as the dispatch core, delegating to engine-specific handlers.
42pub fn dispatch_by_type(envelope: &VShardEnvelope) -> DispatchTarget {
43    match envelope.msg_type {
44        // Timeseries distributed
45        VShardMessageType::TsScatterRequest => DispatchTarget::TimeseriesScan,
46        VShardMessageType::TsScatterResponse => DispatchTarget::TimeseriesCoordinator,
47        VShardMessageType::TsRetentionCommand => DispatchTarget::TimeseriesRetention,
48        VShardMessageType::TsRetentionAck => DispatchTarget::TimeseriesCoordinator,
49        VShardMessageType::TsArchiveCommand => DispatchTarget::TimeseriesArchive,
50        VShardMessageType::TsArchiveAck => DispatchTarget::TimeseriesCoordinator,
51
52        // Migration / infrastructure
53        VShardMessageType::SegmentChunk
54        | VShardMessageType::SegmentComplete
55        | VShardMessageType::WalTail
56        | VShardMessageType::RoutingUpdate
57        | VShardMessageType::RoutingAck => DispatchTarget::Migration,
58
59        VShardMessageType::GhostCreate
60        | VShardMessageType::GhostDelete
61        | VShardMessageType::GhostVerifyRequest
62        | VShardMessageType::GhostVerifyResponse => DispatchTarget::Ghost,
63
64        VShardMessageType::MigrationBaseCopy => DispatchTarget::Migration,
65        VShardMessageType::GsiForward => DispatchTarget::Forward,
66        VShardMessageType::EdgeValidation => DispatchTarget::GraphValidation,
67
68        // Vector distributed search
69        VShardMessageType::VectorScatterRequest => DispatchTarget::VectorSearch,
70        VShardMessageType::VectorScatterResponse => DispatchTarget::VectorCoordinator,
71
72        // Compass coarse-routing phase
73        VShardMessageType::VectorCoarseRouteRequest => DispatchTarget::VectorCoarseRoute,
74        VShardMessageType::VectorCoarseRouteResponse => DispatchTarget::VectorCoordinator,
75
76        // SPIRE build-time centroid exchange
77        VShardMessageType::VectorBuildExchangeRequest => DispatchTarget::VectorBuildExchange,
78        VShardMessageType::VectorBuildExchangeResponse => DispatchTarget::VectorBuildExchange,
79
80        // CoTra-RDMA memory region registration
81        VShardMessageType::VectorMemRegionRequest => DispatchTarget::VectorMemRegion,
82        VShardMessageType::VectorMemRegionResponse => DispatchTarget::VectorMemRegion,
83
84        // Spatial distributed queries
85        VShardMessageType::SpatialScatterRequest => DispatchTarget::SpatialSearch,
86        VShardMessageType::SpatialScatterResponse => DispatchTarget::SpatialCoordinator,
87
88        // Event Plane cross-shard delivery
89        VShardMessageType::CrossShardEvent => DispatchTarget::EventPlane,
90        VShardMessageType::CrossShardEventAck => DispatchTarget::EventPlane,
91        VShardMessageType::NotifyBroadcast => DispatchTarget::EventPlane,
92        VShardMessageType::NotifyBroadcastAck => DispatchTarget::EventPlane,
93
94        // Distributed array (Hilbert-sharded sparse arrays)
95        VShardMessageType::ArrayShardSliceReq => DispatchTarget::ArrayShard,
96        VShardMessageType::ArrayShardSliceResp => DispatchTarget::ArrayCoordinator,
97        VShardMessageType::ArrayShardAggReq => DispatchTarget::ArrayShard,
98        VShardMessageType::ArrayShardAggResp => DispatchTarget::ArrayCoordinator,
99        VShardMessageType::ArrayShardPutReq => DispatchTarget::ArrayShard,
100        VShardMessageType::ArrayShardPutResp => DispatchTarget::ArrayCoordinator,
101        VShardMessageType::ArrayShardDeleteReq => DispatchTarget::ArrayShard,
102        VShardMessageType::ArrayShardDeleteResp => DispatchTarget::ArrayCoordinator,
103        VShardMessageType::ArrayShardSurrogateBitmapReq => DispatchTarget::ArrayShard,
104        VShardMessageType::ArrayShardSurrogateBitmapResp => DispatchTarget::ArrayCoordinator,
105    }
106}
107
108/// Which engine subsystem should handle this envelope.
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub enum DispatchTarget {
111    /// Graph edge validation.
112    GraphValidation,
113    /// Timeseries local scan (shard executes scan, returns partials).
114    TimeseriesScan,
115    /// Timeseries coordinator (receives shard responses).
116    TimeseriesCoordinator,
117    /// Timeseries retention enforcement on local shard.
118    TimeseriesRetention,
119    /// Timeseries S3 archive execution on local shard.
120    TimeseriesArchive,
121    /// Vector local k-NN search (shard executes search, returns top-K hits).
122    VectorSearch,
123    /// Vector coordinator (receives shard search responses).
124    VectorCoordinator,
125    /// Compass coarse-route handler: shard returns its coarse routing descriptor.
126    VectorCoarseRoute,
127    /// SPIRE build-time exchange: shard sends or receives IVF centroid tables.
128    VectorBuildExchange,
129    /// CoTra-RDMA memory region handler: shard registers or exposes a pinned
130    /// memory region for one-sided reads by a remote peer.
131    VectorMemRegion,
132    /// Migration infrastructure.
133    Migration,
134    /// Ghost stub management.
135    Ghost,
136    /// Query/transaction forwarding.
137    Forward,
138    /// Spatial local R-tree search (shard executes predicate, returns matching docs).
139    SpatialSearch,
140    /// Spatial coordinator (receives shard search responses).
141    SpatialCoordinator,
142    /// Event Plane cross-shard delivery (trigger DML, CDC events).
143    EventPlane,
144    /// Array shard: executes slice/agg/put/delete/surrogate-scan locally.
145    ArrayShard,
146    /// Array coordinator: receives shard responses and merges them.
147    ArrayCoordinator,
148}
149
150/// Build a response envelope for a timeseries scatter response.
151pub fn build_ts_scatter_response(
152    source_node: u64,
153    target_node: u64,
154    vshard_id: u32,
155    partials_json: &[u8],
156) -> VShardEnvelope {
157    VShardEnvelope::new(
158        VShardMessageType::TsScatterResponse,
159        source_node,
160        target_node,
161        vshard_id,
162        partials_json.to_vec(),
163    )
164}
165
166/// Build a response envelope for a retention acknowledgement.
167pub fn build_ts_retention_ack(
168    source_node: u64,
169    target_node: u64,
170    vshard_id: u32,
171    result_json: &[u8],
172) -> VShardEnvelope {
173    VShardEnvelope::new(
174        VShardMessageType::TsRetentionAck,
175        source_node,
176        target_node,
177        vshard_id,
178        result_json.to_vec(),
179    )
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn dispatch_ts_scatter() {
188        let env = VShardEnvelope::new(VShardMessageType::TsScatterRequest, 1, 2, 42, vec![]);
189        assert_eq!(dispatch_by_type(&env), DispatchTarget::TimeseriesScan);
190    }
191
192    #[test]
193    fn dispatch_ts_retention() {
194        let env = VShardEnvelope::new(VShardMessageType::TsRetentionCommand, 1, 2, 42, vec![]);
195        assert_eq!(dispatch_by_type(&env), DispatchTarget::TimeseriesRetention);
196    }
197
198    #[test]
199    fn dispatch_migration() {
200        let env = VShardEnvelope::new(VShardMessageType::SegmentChunk, 1, 2, 42, vec![]);
201        assert_eq!(dispatch_by_type(&env), DispatchTarget::Migration);
202    }
203
204    #[test]
205    fn all_message_types_dispatched() {
206        // Every VShardMessageType must map to a DispatchTarget —
207        // this test ensures no new types are added without a handler.
208        let all_types = [
209            VShardMessageType::SegmentChunk,
210            VShardMessageType::SegmentComplete,
211            VShardMessageType::WalTail,
212            VShardMessageType::RoutingUpdate,
213            VShardMessageType::RoutingAck,
214            VShardMessageType::GhostCreate,
215            VShardMessageType::GhostDelete,
216            VShardMessageType::GhostVerifyRequest,
217            VShardMessageType::GhostVerifyResponse,
218            VShardMessageType::MigrationBaseCopy,
219            VShardMessageType::GsiForward,
220            VShardMessageType::EdgeValidation,
221            VShardMessageType::TsScatterRequest,
222            VShardMessageType::TsScatterResponse,
223            VShardMessageType::TsRetentionCommand,
224            VShardMessageType::TsRetentionAck,
225            VShardMessageType::TsArchiveCommand,
226            VShardMessageType::TsArchiveAck,
227            VShardMessageType::VectorScatterRequest,
228            VShardMessageType::VectorScatterResponse,
229            VShardMessageType::VectorCoarseRouteRequest,
230            VShardMessageType::VectorCoarseRouteResponse,
231            VShardMessageType::VectorBuildExchangeRequest,
232            VShardMessageType::VectorBuildExchangeResponse,
233            VShardMessageType::VectorMemRegionRequest,
234            VShardMessageType::VectorMemRegionResponse,
235            VShardMessageType::SpatialScatterRequest,
236            VShardMessageType::SpatialScatterResponse,
237            VShardMessageType::CrossShardEvent,
238            VShardMessageType::CrossShardEventAck,
239            VShardMessageType::NotifyBroadcast,
240            VShardMessageType::NotifyBroadcastAck,
241            VShardMessageType::ArrayShardSliceReq,
242            VShardMessageType::ArrayShardSliceResp,
243            VShardMessageType::ArrayShardAggReq,
244            VShardMessageType::ArrayShardAggResp,
245            VShardMessageType::ArrayShardPutReq,
246            VShardMessageType::ArrayShardPutResp,
247            VShardMessageType::ArrayShardDeleteReq,
248            VShardMessageType::ArrayShardDeleteResp,
249            VShardMessageType::ArrayShardSurrogateBitmapReq,
250            VShardMessageType::ArrayShardSurrogateBitmapResp,
251        ];
252        for msg_type in all_types {
253            let env = VShardEnvelope::new(msg_type, 1, 2, 0, vec![]);
254            let _ = dispatch_by_type(&env); // Must not panic.
255        }
256    }
257}