linkrs_core/metadata/
index_manager.rs

1use super::schema_events::{SchemaChangeCallback, SchemaChangeEvent};
2use crate::event_dispatch::{EventFilter, EventSubscriptions, SubscriptionId};
3use crate::types::Index;
4use crate::StorageError;
5use parking_lot::RwLock;
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::path::Path;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::Arc;
11
12const INDEX_FORMAT_VERSION: u32 = 1;
13
14#[derive(Serialize, Deserialize)]
15struct IndexSnapshot {
16    version: u32,
17    tag_indexes: Vec<(u64, String, Index)>,
18    edge_indexes: Vec<(u64, String, Index)>,
19}
20
21use crate::types::IndexStatus;
22
23pub trait IndexMetadataManager: Send + Sync + std::fmt::Debug {
24    fn create_tag_index(&self, space_id: u64, index: &Index) -> Result<bool, StorageError>;
25    fn drop_tag_index(&self, space_id: u64, index_name: &str) -> Result<bool, StorageError>;
26    fn get_tag_index(&self, space_id: u64, index_name: &str)
27        -> Result<Option<Index>, StorageError>;
28    fn list_tag_indexes(&self, space_id: u64) -> Result<Vec<Index>, StorageError>;
29    fn drop_tag_indexes_by_tag(
30        &self,
31        space_id: u64,
32        tag_name: &str,
33    ) -> Result<Vec<String>, StorageError>;
34
35    fn create_edge_index(&self, space_id: u64, index: &Index) -> Result<bool, StorageError>;
36    fn drop_edge_index(&self, space_id: u64, index_name: &str) -> Result<bool, StorageError>;
37    fn get_edge_index(
38        &self,
39        space_id: u64,
40        index_name: &str,
41    ) -> Result<Option<Index>, StorageError>;
42    fn list_edge_indexes(&self, space_id: u64) -> Result<Vec<Index>, StorageError>;
43    fn drop_edge_indexes_by_type(
44        &self,
45        space_id: u64,
46        edge_type: &str,
47    ) -> Result<Vec<String>, StorageError>;
48
49    /// Set the status of a tag index (used for generation rebuild lifecycle).
50    fn set_tag_index_status(
51        &self,
52        space_id: u64,
53        index_name: &str,
54        status: IndexStatus,
55    ) -> Result<bool, StorageError>;
56
57    /// Set the status of an edge index (used for generation rebuild lifecycle).
58    fn set_edge_index_status(
59        &self,
60        space_id: u64,
61        index_name: &str,
62        status: IndexStatus,
63    ) -> Result<bool, StorageError>;
64}
65
66pub struct IndexManager {
67    tag_indexes: Arc<RwLock<HashMap<(u64, String), Index>>>,
68    edge_indexes: Arc<RwLock<HashMap<(u64, String), Index>>>,
69    next_index_id: AtomicU64,
70    schema_callbacks: Arc<EventSubscriptions<SchemaChangeEvent>>,
71}
72
73impl std::fmt::Debug for IndexManager {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("IndexManager")
76            .field("tag_indexes_count", &self.tag_indexes.read().len())
77            .field("edge_indexes_count", &self.edge_indexes.read().len())
78            .finish()
79    }
80}
81
82impl IndexManager {
83    pub fn new() -> Self {
84        Self {
85            tag_indexes: Arc::new(RwLock::new(HashMap::new())),
86            edge_indexes: Arc::new(RwLock::new(HashMap::new())),
87            next_index_id: AtomicU64::new(1),
88            schema_callbacks: Arc::new(EventSubscriptions::new()),
89        }
90    }
91
92    /// Build a manager sharing one schema-event registry with `SchemaManager`.
93    ///
94    /// Single enum dual emission source: index DDL is emitted here, everything
95    /// else by `SchemaManager`; sharing one registry lets observers subscribe
96    /// once and receive both halves.
97    pub fn with_shared_schema_callbacks(
98        shared: Arc<EventSubscriptions<SchemaChangeEvent>>,
99    ) -> Self {
100        Self {
101            tag_indexes: Arc::new(RwLock::new(HashMap::new())),
102            edge_indexes: Arc::new(RwLock::new(HashMap::new())),
103            next_index_id: AtomicU64::new(1),
104            schema_callbacks: shared,
105        }
106    }
107
108    /// Shared schema-event registry behind this manager.
109    pub fn shared_schema_callbacks(&self) -> Arc<EventSubscriptions<SchemaChangeEvent>> {
110        Arc::clone(&self.schema_callbacks)
111    }
112
113    /// Register a runtime observer for index DDL events.
114    ///
115    /// Writes to the registry shared with `SchemaManager`, so one
116    /// subscription receives both index DDL and table/space DDL.
117    pub fn register_schema_callback(&self, callback: SchemaChangeCallback) -> SubscriptionId {
118        self.schema_callbacks.add(callback)
119    }
120
121    /// Register a filtered observer invoked only when `filter` returns true.
122    pub fn register_schema_callback_filtered(
123        &self,
124        callback: SchemaChangeCallback,
125        filter: EventFilter<SchemaChangeEvent>,
126    ) -> SubscriptionId {
127        self.schema_callbacks.add_filtered(callback, Some(filter))
128    }
129
130    /// Remove a previously registered observer. Returns true if present.
131    pub fn unregister_schema_callback(&self, id: SubscriptionId) -> bool {
132        self.schema_callbacks.remove(id)
133    }
134
135    /// Number of registered index observers.
136    pub fn schema_callback_count(&self) -> usize {
137        self.schema_callbacks.len()
138    }
139
140    fn emit_schema_event(&self, event: SchemaChangeEvent) {
141        if self.schema_callbacks.is_empty() {
142            return;
143        }
144        self.schema_callbacks.dispatch("schema", &event);
145    }
146
147    pub fn save_indexes(&self, path: &Path) -> Result<(), StorageError> {
148        use std::fs::{self, File};
149        use std::io::Write;
150
151        if let Some(parent) = path.parent() {
152            fs::create_dir_all(parent).map_err(|e| StorageError::io_error(e.to_string()))?;
153        }
154
155        let tag_indexes: Vec<(u64, String, Index)> = self
156            .tag_indexes
157            .read()
158            .iter()
159            .map(|((space_id, name), index)| (*space_id, name.clone(), index.clone()))
160            .collect();
161
162        let edge_indexes: Vec<(u64, String, Index)> = self
163            .edge_indexes
164            .read()
165            .iter()
166            .map(|((space_id, name), index)| (*space_id, name.clone(), index.clone()))
167            .collect();
168
169        let snapshot = IndexSnapshot {
170            version: INDEX_FORMAT_VERSION,
171            tag_indexes,
172            edge_indexes,
173        };
174
175        let json = serde_json::to_string_pretty(&snapshot)
176            .map_err(|e| StorageError::serialize_error(e.to_string()))?;
177
178        let mut file = File::create(path).map_err(|e| StorageError::io_error(e.to_string()))?;
179        file.write_all(json.as_bytes())
180            .map_err(|e| StorageError::io_error(e.to_string()))?;
181
182        Ok(())
183    }
184
185    pub fn load_indexes(&self, path: &Path) -> Result<(), StorageError> {
186        use std::fs::File;
187        use std::io::Read;
188
189        if !path.exists() {
190            return Ok(());
191        }
192
193        let mut file = File::open(path).map_err(|e| StorageError::io_error(e.to_string()))?;
194        let mut json = String::new();
195        file.read_to_string(&mut json)
196            .map_err(|e| StorageError::io_error(e.to_string()))?;
197
198        let snapshot: IndexSnapshot = serde_json::from_str(&json)
199            .map_err(|e| StorageError::deserialize_error(e.to_string()))?;
200
201        if snapshot.version > INDEX_FORMAT_VERSION {
202            return Err(StorageError::deserialize_error(format!(
203                "Index snapshot version {} is newer than supported version {}",
204                snapshot.version, INDEX_FORMAT_VERSION
205            )));
206        }
207
208        self.tag_indexes.write().clear();
209        self.edge_indexes.write().clear();
210
211        let mut max_id = 0;
212        for (space_id, name, index) in snapshot.tag_indexes {
213            if index.id > max_id {
214                max_id = index.id;
215            }
216            self.tag_indexes.write().insert((space_id, name), index);
217        }
218
219        for (space_id, name, index) in snapshot.edge_indexes {
220            if index.id > max_id {
221                max_id = index.id;
222            }
223            self.edge_indexes.write().insert((space_id, name), index);
224        }
225
226        self.next_index_id.store(max_id + 1, Ordering::SeqCst);
227
228        Ok(())
229    }
230}
231
232impl Default for IndexManager {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238impl IndexMetadataManager for IndexManager {
239    fn create_tag_index(&self, space_id: u64, index: &Index) -> Result<bool, StorageError> {
240        let created = {
241            let mut indexes = self.tag_indexes.write();
242            let key = (space_id, index.name.clone());
243            if indexes.contains_key(&key) {
244                return Ok(false);
245            }
246            let mut index_with_space_id = index.clone();
247            index_with_space_id.space_id = space_id;
248            if index_with_space_id.id == 0 {
249                index_with_space_id.id = self.next_index_id.fetch_add(1, Ordering::SeqCst);
250            }
251            indexes.insert(key, index_with_space_id);
252            true
253        };
254        if created {
255            self.emit_schema_event(SchemaChangeEvent::TagIndexCreated {
256                space_id,
257                index_name: index.name.clone(),
258            });
259        }
260        Ok(created)
261    }
262
263    fn drop_tag_index(&self, space_id: u64, index_name: &str) -> Result<bool, StorageError> {
264        let removed = {
265            let mut indexes = self.tag_indexes.write();
266            let key = (space_id, index_name.to_string());
267            indexes.remove(&key).is_some()
268        };
269        if removed {
270            self.emit_schema_event(SchemaChangeEvent::TagIndexDropped {
271                space_id,
272                index_name: index_name.to_string(),
273            });
274        }
275        Ok(removed)
276    }
277
278    fn get_tag_index(
279        &self,
280        space_id: u64,
281        index_name: &str,
282    ) -> Result<Option<Index>, StorageError> {
283        let indexes = self.tag_indexes.read();
284        Ok(indexes.get(&(space_id, index_name.to_string())).cloned())
285    }
286
287    fn list_tag_indexes(&self, space_id: u64) -> Result<Vec<Index>, StorageError> {
288        let indexes = self.tag_indexes.read();
289        Ok(indexes
290            .iter()
291            .filter(|((sid, _), _)| *sid == space_id)
292            .map(|(_, index)| index.clone())
293            .collect())
294    }
295
296    fn drop_tag_indexes_by_tag(
297        &self,
298        space_id: u64,
299        tag_name: &str,
300    ) -> Result<Vec<String>, StorageError> {
301        let removed: Vec<String> = {
302            let mut indexes = self.tag_indexes.write();
303            let mut removed = Vec::new();
304            indexes.retain(|_, index| {
305                let drop_it = index.space_id == space_id && index.schema_name == tag_name;
306                if drop_it {
307                    removed.push(index.name.clone());
308                }
309                !drop_it
310            });
311            removed
312        };
313        for index_name in &removed {
314            self.emit_schema_event(SchemaChangeEvent::TagIndexDropped {
315                space_id,
316                index_name: index_name.clone(),
317            });
318        }
319        Ok(removed)
320    }
321
322    fn create_edge_index(&self, space_id: u64, index: &Index) -> Result<bool, StorageError> {
323        let created = {
324            let mut indexes = self.edge_indexes.write();
325            let key = (space_id, index.name.clone());
326            if indexes.contains_key(&key) {
327                return Ok(false);
328            }
329            let mut index_with_space_id = index.clone();
330            index_with_space_id.space_id = space_id;
331            if index_with_space_id.id == 0 {
332                index_with_space_id.id = self.next_index_id.fetch_add(1, Ordering::SeqCst);
333            }
334            indexes.insert(key, index_with_space_id);
335            true
336        };
337        if created {
338            self.emit_schema_event(SchemaChangeEvent::EdgeIndexCreated {
339                space_id,
340                index_name: index.name.clone(),
341            });
342        }
343        Ok(created)
344    }
345
346    fn drop_edge_index(&self, space_id: u64, index_name: &str) -> Result<bool, StorageError> {
347        let removed = {
348            let mut indexes = self.edge_indexes.write();
349            let key = (space_id, index_name.to_string());
350            indexes.remove(&key).is_some()
351        };
352        if removed {
353            self.emit_schema_event(SchemaChangeEvent::EdgeIndexDropped {
354                space_id,
355                index_name: index_name.to_string(),
356            });
357        }
358        Ok(removed)
359    }
360
361    fn get_edge_index(
362        &self,
363        space_id: u64,
364        index_name: &str,
365    ) -> Result<Option<Index>, StorageError> {
366        let indexes = self.edge_indexes.read();
367        Ok(indexes.get(&(space_id, index_name.to_string())).cloned())
368    }
369
370    fn list_edge_indexes(&self, space_id: u64) -> Result<Vec<Index>, StorageError> {
371        let indexes = self.edge_indexes.read();
372        Ok(indexes
373            .iter()
374            .filter(|((sid, _), _)| *sid == space_id)
375            .map(|(_, index)| index.clone())
376            .collect())
377    }
378
379    fn drop_edge_indexes_by_type(
380        &self,
381        space_id: u64,
382        edge_type: &str,
383    ) -> Result<Vec<String>, StorageError> {
384        let removed: Vec<String> = {
385            let mut indexes = self.edge_indexes.write();
386            let mut removed = Vec::new();
387            indexes.retain(|_, index| {
388                let drop_it = index.space_id == space_id && index.schema_name == edge_type;
389                if drop_it {
390                    removed.push(index.name.clone());
391                }
392                !drop_it
393            });
394            removed
395        };
396        for index_name in &removed {
397            self.emit_schema_event(SchemaChangeEvent::EdgeIndexDropped {
398                space_id,
399                index_name: index_name.clone(),
400            });
401        }
402        Ok(removed)
403    }
404
405    fn set_tag_index_status(
406        &self,
407        space_id: u64,
408        index_name: &str,
409        status: IndexStatus,
410    ) -> Result<bool, StorageError> {
411        let mut indexes = self.tag_indexes.write();
412        let key = (space_id, index_name.to_string());
413        if let Some(index) = indexes.get_mut(&key) {
414            index.status = status;
415            Ok(true)
416        } else {
417            Ok(false)
418        }
419    }
420
421    fn set_edge_index_status(
422        &self,
423        space_id: u64,
424        index_name: &str,
425        status: IndexStatus,
426    ) -> Result<bool, StorageError> {
427        let mut indexes = self.edge_indexes.write();
428        let key = (space_id, index_name.to_string());
429        if let Some(index) = indexes.get_mut(&key) {
430            index.status = status;
431            Ok(true)
432        } else {
433            Ok(false)
434        }
435    }
436}
437
438#[cfg(test)]
439mod tests {
440    use super::*;
441    use crate::types::{IndexConfig, IndexType};
442    use std::sync::atomic::{AtomicUsize, Ordering};
443
444    fn tag_index(name: &str, schema: &str) -> Index {
445        Index::new(IndexConfig {
446            id: 0,
447            name: name.to_string(),
448            space_id: 1,
449            schema_name: schema.to_string(),
450            fields: Vec::new(),
451            properties: Vec::new(),
452            index_type: IndexType::TagIndex,
453            is_unique: false,
454            covering: false,
455            partial_condition: None,
456        })
457    }
458
459    #[test]
460    fn batch_drop_emits_one_event_per_index() {
461        let manager = IndexManager::new();
462        manager
463            .create_tag_index(1, &tag_index("idx_a", "Person"))
464            .unwrap();
465        manager
466            .create_tag_index(1, &tag_index("idx_b", "Person"))
467            .unwrap();
468
469        let hits = Arc::new(AtomicUsize::new(0));
470        let probe = Arc::clone(&hits);
471        manager.register_schema_callback(Arc::new(move |event| {
472            if matches!(event, SchemaChangeEvent::TagIndexDropped { .. }) {
473                probe.fetch_add(1, Ordering::SeqCst);
474            }
475        }));
476
477        manager.drop_tag_indexes_by_tag(1, "Person").unwrap();
478        assert_eq!(hits.load(Ordering::SeqCst), 2);
479        assert!(manager.list_tag_indexes(1).unwrap().is_empty());
480        let removed = manager.drop_tag_indexes_by_tag(1, "Person").unwrap();
481        assert!(removed.is_empty());
482    }
483
484    #[test]
485    fn unregister_stops_delivery() {
486        let manager = IndexManager::new();
487        let hits = Arc::new(AtomicUsize::new(0));
488        let probe = Arc::clone(&hits);
489        let id = manager.register_schema_callback(Arc::new(move |_| {
490            probe.fetch_add(1, Ordering::SeqCst);
491        }));
492        assert!(manager.unregister_schema_callback(id));
493        manager
494            .create_tag_index(1, &tag_index("idx_c", "Person"))
495            .unwrap();
496        assert_eq!(hits.load(Ordering::SeqCst), 0);
497    }
498
499    #[test]
500    fn shared_registry_receives_schema_and_index_halves_once() {
501        use super::super::schema_events::SchemaChangeEvent as SharedSchemaEvent;
502        use crate::event_dispatch::EventSubscriptions;
503        use crate::metadata::SchemaManager;
504        use crate::types::SpaceInfo;
505
506        let shared = Arc::new(EventSubscriptions::<SharedSchemaEvent>::new());
507        let schema = SchemaManager::with_shared_schema_callbacks(Arc::clone(&shared));
508        let indexes = IndexManager::with_shared_schema_callbacks(shared);
509        assert!(Arc::ptr_eq(
510            &schema.shared_schema_callbacks(),
511            &indexes.shared_schema_callbacks()
512        ));
513
514        let spaces = Arc::new(AtomicUsize::new(0));
515        let idx_events = Arc::new(AtomicUsize::new(0));
516        let space_probe = Arc::clone(&spaces);
517        let idx_probe = Arc::clone(&idx_events);
518        // Single subscription receives both halves.
519        schema.register_schema_callback(Arc::new(move |event| match event {
520            SharedSchemaEvent::SpaceCreated { .. } => {
521                space_probe.fetch_add(1, Ordering::SeqCst);
522            }
523            SharedSchemaEvent::TagIndexCreated { .. } => {
524                idx_probe.fetch_add(1, Ordering::SeqCst);
525            }
526            _ => {}
527        }));
528
529        let mut space = SpaceInfo::new("shared_space".to_string());
530        assert!(schema.create_space(&mut space).unwrap());
531        indexes
532            .create_tag_index(space.space_id, &tag_index("idx_shared", "Person"))
533            .unwrap();
534
535        assert_eq!(spaces.load(Ordering::SeqCst), 1);
536        assert_eq!(idx_events.load(Ordering::SeqCst), 1);
537        // Registered once: the shared registry holds exactly one observer.
538        assert_eq!(schema.schema_callback_count(), 1);
539        assert_eq!(indexes.schema_callback_count(), 1);
540    }
541}
linkrs_core/metadata/index_manager.rs

linkrs_core/metadata/
index_manager.rs
