mongreldb_core/
ai_generation.rs

1//! AI index generations (spec section 13.3, Stage 4C).
2//!
3//! AI indexes remain local derived state. Each generation records definition
4//! version, applied_through watermark, preprocessing/model versions, and
5//! base/delta generation ids. Replicas build local graphs (no byte-identical
6//! HNSW requirement). A replica may serve an indexed read only when
7//! `applied_through >= requested read timestamp`; otherwise the caller waits,
8//! routes elsewhere, uses exact fallback, or returns [`IndexNotReady`].
9
10use mongreldb_types::hlc::HlcTimestamp;
11use mongreldb_types::ids::SchemaVersion;
12use serde::{Deserialize, Serialize};
13
14/// Opaque index identifier (stable within a database).
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
16#[repr(transparent)]
17pub struct IndexId(pub u64);
18
19impl IndexId {
20    /// Wrap a raw id.
21    pub const fn new(id: u64) -> Self {
22        Self(id)
23    }
24
25    /// Raw id.
26    pub const fn get(self) -> u64 {
27        self.0
28    }
29}
30
31impl std::fmt::Display for IndexId {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "{}", self.0)
34    }
35}
36
37/// One AI index generation (spec §13.3).
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct AiIndexGeneration {
40    /// Index identity.
41    pub index_id: IndexId,
42    /// Definition version (catalog/schema of the index).
43    pub definition_version: u64,
44    /// Highest commit timestamp whose rows are reflected in this generation.
45    pub applied_through: HlcTimestamp,
46    /// Source table schema version at build time.
47    pub source_schema_version: SchemaVersion,
48    /// Preprocessing pipeline version string.
49    pub preprocessing_version: String,
50    /// Optional embedding/model version.
51    pub model_version: Option<String>,
52    /// Base (full rebuild) generation id.
53    pub base_generation: u64,
54    /// Delta generation ids applied on top of the base.
55    pub delta_generations: Vec<u64>,
56}
57
58/// Why an indexed AI read cannot be served from this replica.
59#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
60pub enum IndexReadinessError {
61    /// Index generation has not caught up to the requested read timestamp.
62    #[error("index {index_id} not ready: applied_through {applied_through} < read_ts {read_ts}")]
63    IndexNotReady {
64        /// Index id.
65        index_id: IndexId,
66        /// Generation watermark.
67        applied_through: HlcTimestamp,
68        /// Requested read timestamp.
69        read_ts: HlcTimestamp,
70    },
71    /// No generation is registered for the index.
72    #[error("index {0} has no generation")]
73    Missing(IndexId),
74}
75
76/// How the caller should react to an unreadiness result.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum ReadinessAction {
79    /// Wait for local catch-up.
80    Wait,
81    /// Route the read to another replica.
82    RouteElsewhere,
83    /// Fall back to exact (non-ANN) retrieval.
84    ExactFallback,
85    /// Return IndexNotReady to the client.
86    FailClosed,
87}
88
89/// Evaluate readiness of a generation against a requested read timestamp.
90pub fn evaluate_readiness(
91    generation: &AiIndexGeneration,
92    read_ts: HlcTimestamp,
93) -> Result<(), IndexReadinessError> {
94    if generation.applied_through >= read_ts {
95        Ok(())
96    } else {
97        Err(IndexReadinessError::IndexNotReady {
98            index_id: generation.index_id,
99            applied_through: generation.applied_through,
100            read_ts,
101        })
102    }
103}
104
105/// Choose an action when readiness fails (policy helper for the gateway).
106pub fn readiness_action(
107    prefer_exact_fallback: bool,
108    has_alternate_replica: bool,
109    deadline_exhausted: bool,
110) -> ReadinessAction {
111    if deadline_exhausted {
112        if prefer_exact_fallback {
113            ReadinessAction::ExactFallback
114        } else {
115            ReadinessAction::FailClosed
116        }
117    } else if has_alternate_replica {
118        ReadinessAction::RouteElsewhere
119    } else if prefer_exact_fallback {
120        ReadinessAction::ExactFallback
121    } else {
122        ReadinessAction::Wait
123    }
124}
125
126/// Registry of AI index generations on one replica (local derived state).
127#[derive(Debug, Default, Clone)]
128pub struct AiIndexGenerationRegistry {
129    by_index: std::collections::BTreeMap<IndexId, AiIndexGeneration>,
130}
131
132impl AiIndexGenerationRegistry {
133    /// Empty registry.
134    pub fn new() -> Self {
135        Self::default()
136    }
137
138    /// Publish (or replace) a generation.
139    pub fn publish(&mut self, generation: AiIndexGeneration) {
140        self.by_index.insert(generation.index_id, generation);
141    }
142
143    /// Lookup.
144    pub fn get(&self, index_id: IndexId) -> Option<&AiIndexGeneration> {
145        self.by_index.get(&index_id)
146    }
147
148    /// Number of registered index generations.
149    pub fn len(&self) -> usize {
150        self.by_index.len()
151    }
152
153    /// Whether no generations are registered.
154    pub fn is_empty(&self) -> bool {
155        self.by_index.is_empty()
156    }
157
158    /// Serve check: ready for `read_ts` or typed error.
159    pub fn require_ready(
160        &self,
161        index_id: IndexId,
162        read_ts: HlcTimestamp,
163    ) -> Result<&AiIndexGeneration, IndexReadinessError> {
164        let gen = self
165            .by_index
166            .get(&index_id)
167            .ok_or(IndexReadinessError::Missing(index_id))?;
168        evaluate_readiness(gen, read_ts)?;
169        Ok(gen)
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    fn ts(micros: u64) -> HlcTimestamp {
178        HlcTimestamp {
179            physical_micros: micros,
180            logical: 0,
181            node_tiebreaker: 1,
182        }
183    }
184
185    fn gen(applied: u64) -> AiIndexGeneration {
186        AiIndexGeneration {
187            index_id: IndexId::new(7),
188            definition_version: 1,
189            applied_through: ts(applied),
190            source_schema_version: SchemaVersion::new(1),
191            preprocessing_version: "prep-1".into(),
192            model_version: Some("embed-v2".into()),
193            base_generation: 1,
194            delta_generations: vec![2, 3],
195        }
196    }
197
198    #[test]
199    fn readiness_requires_applied_through_ge_read_ts() {
200        let g = gen(100);
201        assert!(evaluate_readiness(&g, ts(100)).is_ok());
202        assert!(evaluate_readiness(&g, ts(50)).is_ok());
203        let err = evaluate_readiness(&g, ts(101)).unwrap_err();
204        assert!(matches!(err, IndexReadinessError::IndexNotReady { .. }));
205    }
206
207    #[test]
208    fn registry_serve_and_missing() {
209        let mut reg = AiIndexGenerationRegistry::new();
210        reg.publish(gen(200));
211        assert!(reg.require_ready(IndexId::new(7), ts(150)).is_ok());
212        assert!(matches!(
213            reg.require_ready(IndexId::new(9), ts(1)).unwrap_err(),
214            IndexReadinessError::Missing(_)
215        ));
216    }
217
218    #[test]
219    fn readiness_action_policy() {
220        assert_eq!(
221            readiness_action(true, false, true),
222            ReadinessAction::ExactFallback
223        );
224        assert_eq!(
225            readiness_action(false, true, false),
226            ReadinessAction::RouteElsewhere
227        );
228        assert_eq!(readiness_action(false, false, false), ReadinessAction::Wait);
229        assert_eq!(
230            readiness_action(false, false, true),
231            ReadinessAction::FailClosed
232        );
233    }
234}
mongreldb_core/ai_generation.rs

mongreldb_core/
ai_generation.rs
