1pub mod bm25_embed;
2mod builders;
3pub mod config;
4mod count;
5mod facet;
6mod info;
7mod optimize;
8mod query;
9mod reexports;
10mod retrieve;
11mod scroll;
12mod search;
13mod snapshots;
14mod types;
15pub use types::*;
16mod update;
17
18use std::path::{Path, PathBuf};
19use std::sync::Arc;
20use std::sync::atomic::AtomicBool;
21use std::time::Duration;
22
23pub use builders::{EdgeConfigBuilder, EdgeSparseVectorParamsBuilder, EdgeVectorParamsBuilder};
24use crate::common::save_on_disk::SaveOnDisk;
25pub use config::optimizers::EdgeOptimizersConfig;
26pub use config::shard::EdgeConfig;
27pub use config::vectors::{EdgeSparseVectorParams, EdgeVectorParams};
28use fs_err as fs;
29pub use info::ShardInfo;
30use parking_lot::Mutex;
31pub use reexports::*;
32use crate::segment::entry::ReadSegmentEntry as _;
33use crate::segment::segment_constructor::{load_segment, normalize_segment_dir};
34use crate::shard::files::{PAYLOAD_INDEX_CONFIG_FILE, SEGMENTS_PATH};
35use crate::shard::operations::CollectionUpdateOperations;
36use crate::shard::segment_holder::SegmentHolder;
37use crate::shard::segment_holder::locked::LockedSegmentHolder;
38use crate::shard::wal::SerdeWal;
39
40use crate::edge::config::shard::EDGE_CONFIG_FILE;
41
42#[derive(Debug)]
43pub struct EdgeShard {
44 path: PathBuf,
45 config: SaveOnDisk<EdgeConfig>,
46 wal: Mutex<SerdeWal<CollectionUpdateOperations>>,
47 segments: LockedSegmentHolder,
48}
49
50const WAL_PATH: &str = "wal";
51impl EdgeShard {
52 pub fn new(path: &Path, config: EdgeConfig) -> OperationResult<Self> {
59 if has_existing_segments(path) {
60 return Err(OperationError::service_error(
61 "cannot create edge shard: path already contains segment data",
62 ));
63 }
64
65 let wal_options = config.wal_options.clone().unwrap_or_default();
66 let (wal, segments_path) = ensure_dirs_and_open_wal(path, wal_options)?;
67 config.save(path)?;
68
69 let mut segments = SegmentHolder::default();
70 ensure_appendable_segment(&mut segments, path, &segments_path, &config)?;
71
72 let config_path = path.join(EDGE_CONFIG_FILE);
73 let config = SaveOnDisk::new(&config_path, config)
74 .map_err(|e| OperationError::service_error(e.to_string()))?;
75
76 Ok(Self {
77 path: path.into(),
78 config,
79 wal: parking_lot::Mutex::new(wal),
80 segments: LockedSegmentHolder::new(segments),
81 })
82 }
83
84 pub fn load(path: &Path, config: Option<EdgeConfig>) -> OperationResult<Self> {
97 let mut config = resolve_initial_config(path, config)?;
98
99 let wal_options = config
100 .as_ref()
101 .and_then(|c| c.wal_options.clone())
102 .unwrap_or_default();
103 let (wal, segments_path) = ensure_dirs_and_open_wal(path, wal_options)?;
104
105 let mut segments = load_segments(path, &segments_path, &mut config)?;
106
107 ensure_appendable_segment(
108 &mut segments,
109 path,
110 &segments_path,
111 config.as_ref().ok_or_else(|| {
112 OperationError::service_error(
113 "edge config is not provided and no segments were loaded",
114 )
115 })?,
116 )?;
117
118 let config = config.ok_or_else(|| {
119 OperationError::service_error("edge config is not provided and no segments were loaded")
120 })?;
121
122 let config_path = path.join(EDGE_CONFIG_FILE);
123 let config = SaveOnDisk::new(&config_path, config)
124 .map_err(|e| OperationError::service_error(e.to_string()))?;
125
126 Ok(Self {
127 path: path.into(),
128 config,
129 wal: parking_lot::Mutex::new(wal),
130 segments: LockedSegmentHolder::new(segments),
131 })
132 }
133
134 pub fn config(&self) -> parking_lot::RwLockReadGuard<'_, EdgeConfig> {
135 self.config.read()
136 }
137
138 pub fn path(&self) -> &Path {
139 &self.path
140 }
141
142 pub fn set_hnsw_config(&self, hnsw_config: crate::segment::types::HnswConfig) -> OperationResult<()> {
144 self.config
145 .write(|cfg| cfg.set_hnsw_config(hnsw_config))
146 .map_err(|e| OperationError::service_error(e.to_string()))
147 }
148
149 pub fn set_vector_hnsw_config(
152 &self,
153 vector_name: &str,
154 hnsw_config: crate::segment::types::HnswConfig,
155 ) -> OperationResult<()> {
156 let mut cfg = self.config.read().clone();
157 cfg.set_vector_hnsw_config(vector_name, hnsw_config)?;
158 self.config
159 .write(|c| *c = cfg)
160 .map_err(|e| OperationError::service_error(e.to_string()))
161 }
162
163 pub fn set_optimizers_config(&self, optimizers: EdgeOptimizersConfig) -> OperationResult<()> {
165 self.config
166 .write(|cfg| cfg.set_optimizers_config(optimizers))
167 .map_err(|e| OperationError::service_error(e.to_string()))
168 }
169
170 pub fn flush(&self) {
171 self.wal
172 .try_lock()
173 .expect("WAL lock acquired")
174 .flush()
175 .expect("WAL flushed");
176
177 self.segments
178 .try_read()
179 .expect("segment holder lock acquired")
180 .flush_all(true, true)
181 .expect("segments flushed");
182 }
183}
184
185impl Drop for EdgeShard {
186 fn drop(&mut self) {
187 self.flush();
188 }
189}
190
191fn has_existing_segments(path: &Path) -> bool {
192 let segments_path = path.join(SEGMENTS_PATH);
193 let Ok(entries) = fs::read_dir(&segments_path) else {
194 return false;
195 };
196 for entry in entries.flatten() {
197 let p = entry.path();
198 if !p.is_dir() {
199 continue;
200 }
201 if p.file_name()
202 .and_then(|n| n.to_str())
203 .is_some_and(|n| n.starts_with('.'))
204 {
205 continue;
206 }
207 if normalize_segment_dir(&p).ok().flatten().is_some() {
208 return true;
209 }
210 }
211 false
212}
213
214fn ensure_dirs_and_open_wal(
215 path: &Path,
216 wal_options: WalOptions,
217) -> OperationResult<(SerdeWal<CollectionUpdateOperations>, PathBuf)> {
218 let wal_path = path.join(WAL_PATH);
219 if !wal_path.exists() {
220 fs::create_dir(&wal_path).map_err(|err| {
221 OperationError::service_error(format!("failed to create WAL directory: {err}"))
222 })?;
223 }
224
225 let wal = SerdeWal::new(&wal_path, wal_options).map_err(|err| {
226 OperationError::service_error(format!("failed to open WAL {}: {err}", wal_path.display(),))
227 })?;
228
229 let segments_path = path.join(SEGMENTS_PATH);
230 if !segments_path.exists() {
231 fs::create_dir(&segments_path).map_err(|err| {
232 OperationError::service_error(format!("failed to create segments directory: {err}"))
233 })?;
234 }
235
236 Ok((wal, segments_path))
237}
238
239fn resolve_initial_config(
240 path: &Path,
241 config: Option<EdgeConfig>,
242) -> OperationResult<Option<EdgeConfig>> {
243 Ok(match config {
244 Some(c) => Some(c),
245 None => match EdgeConfig::load(path) {
246 Some(Ok(c)) => Some(c),
247 Some(Err(e)) => return Err(e),
248 None => None,
249 },
250 })
251}
252
253fn load_segments(
254 _path: &Path,
255 segments_path: &Path,
256 config: &mut Option<EdgeConfig>,
257) -> OperationResult<SegmentHolder> {
258 let segments_dir = fs::read_dir(segments_path).map_err(|err| {
259 OperationError::service_error(format!("failed to read segments directory: {err}"))
260 })?;
261
262 let mut segments = SegmentHolder::default();
263
264 for entry in segments_dir {
265 let entry = entry.map_err(|err| {
266 OperationError::service_error(format!(
267 "failed to read entry in segments directory: {err}",
268 ))
269 })?;
270
271 let segment_path = entry.path();
272
273 if !segment_path.is_dir() {
274 log::warn!(
275 "Skipping non-directory segment entry {}",
276 segment_path.display(),
277 );
278 continue;
279 }
280
281 if segment_path
282 .file_name()
283 .and_then(|n| n.to_str())
284 .is_some_and(|n| n.starts_with('.'))
285 {
286 log::warn!(
287 "Skipping hidden segment directory {}",
288 segment_path.display(),
289 );
290 continue;
291 }
292
293 let Some((segment_path, segment_uuid)) = normalize_segment_dir(&segment_path)? else {
294 continue;
295 };
296
297 let mut segment = load_segment(&segment_path, segment_uuid, None, &AtomicBool::new(false))
298 .map_err(|err| {
299 OperationError::service_error(format!(
300 "failed to load segment {}: {err}",
301 segment_path.display(),
302 ))
303 })?;
304
305 let segment_cfg = segment.config();
306 if let Some(cfg) = config.as_ref() {
307 cfg.check_compatible_with_segment_config(segment_cfg).map_err(
308 |err| OperationError::service_error(format!(
309 "segment {} is incompatible with provided config or previously loaded segments: {err}",
310 segment_path.display(),
311 ))
312 )?;
313 } else {
314 *config = Some(EdgeConfig::from_segment_config(segment_cfg));
315 }
316
317 segment.check_consistency_and_repair().map_err(|err| {
318 OperationError::service_error(format!(
319 "failed to repair segment {}: {err}",
320 segment_path.display(),
321 ))
322 })?;
323
324 segments.add_new(segment);
325 }
326
327 Ok(segments)
328}
329
330fn ensure_appendable_segment(
331 segments: &mut SegmentHolder,
332 path: &Path,
333 segments_path: &Path,
334 config: &EdgeConfig,
335) -> OperationResult<()> {
336 if segments.has_appendable_segment() {
337 return Ok(());
338 }
339
340 let payload_index_schema_path = path.join(PAYLOAD_INDEX_CONFIG_FILE);
341 let payload_index_schema = SaveOnDisk::load_or_init_default(&payload_index_schema_path)
342 .map_err(|err| {
343 OperationError::service_error(format!(
344 "failed to initialize payload index schema file {}: {err}",
345 payload_index_schema_path.display(),
346 ))
347 })?;
348
349 segments.create_appendable_segment(
350 segments_path,
351 config.plain_segment_config(),
352 Arc::new(payload_index_schema),
353 None,
354 )?;
355
356 debug_assert!(segments.has_appendable_segment());
357 Ok(())
358}
359
360pub(crate) const DEFAULT_EDGE_TIMEOUT: Duration = Duration::from_secs(3600);