1use std::collections::{HashMap, HashSet};
40use std::path::{Path, PathBuf};
41use std::sync::Arc;
42
43use super::{
44 extract_keyspace_and_table_name, extract_table_name, is_apple_double_sidecar, reader,
45 SSTableId, SSTableManager, MAX_SSTABLE_SCAN_DEPTH,
46};
47use crate::Result;
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub struct RefreshReport {
56 pub tables_scanned: usize,
58 pub readers_added: usize,
60 pub readers_removed: usize,
62}
63
64#[derive(Debug, Clone)]
67pub(crate) enum DiscoverySource {
68 BasePath,
71 TableDirs(Vec<PathBuf>),
74}
75
76pub(crate) fn base_path_table_key(
86 path: &Path,
87 base_dir_name: &str,
88 header_table_name: &str,
89) -> Option<String> {
90 if let Some((keyspace, table_name)) = extract_keyspace_and_table_name(path) {
91 if table_name.as_str() != base_dir_name {
92 return Some(format!("{}.{}", keyspace, table_name));
93 }
94 }
95 if let Some(name) = extract_table_name(path) {
96 if name.as_str() != base_dir_name {
97 return Some(name);
98 }
99 }
100 if header_table_name != "test_table" && !header_table_name.is_empty() {
101 return Some(header_table_name.to_string());
102 }
103 None
104}
105
106pub(crate) fn table_dir_table_key(path: &Path) -> Option<String> {
110 if let Some((keyspace, table_name)) = extract_keyspace_and_table_name(path) {
111 Some(format!("{}.{}", keyspace, table_name))
112 } else {
113 extract_table_name(path)
114 }
115}
116
117fn canon(p: &Path) -> PathBuf {
121 std::fs::canonicalize(p).unwrap_or_else(|_| p.to_path_buf())
122}
123
124impl SSTableManager {
125 pub(crate) async fn open_reader_with_schema(
133 &self,
134 path: &Path,
135 ) -> Result<Arc<reader::SSTableReader>> {
136 #[cfg_attr(not(feature = "state_machine"), allow(unused_mut))]
137 let mut reader = reader::SSTableReader::open_with_cache(
138 path,
139 &self.config,
140 self.platform.clone(),
141 self.chunk_cache.clone(),
142 )
143 .await?;
144
145 #[cfg(feature = "state_machine")]
146 {
147 let schema_reg_guard = self.schema_registry.read().await;
148 if let Some(ref registry_rwlock) = *schema_reg_guard {
149 #[allow(deprecated)]
152 reader.set_schema_registry(Arc::clone(registry_rwlock));
153
154 reader.resolve_registry_schema().await;
159
160 let schema_registry = registry_rwlock.read().await;
162 let udt_registry_lock = schema_registry.get_udt_registry();
163 let udt_registry = udt_registry_lock.read().await.clone();
164 reader.set_udt_registry(udt_registry);
165 }
166 }
167
168 Ok(Arc::new(reader))
169 }
170
171 async fn discover_data_file_paths(&self) -> Result<Vec<PathBuf>> {
174 match &self.discovery_source {
175 DiscoverySource::BasePath => {
176 if !self.platform.fs().exists(&self.base_path).await? {
177 return Ok(Vec::new());
178 }
179 SSTableManager::find_data_files(
180 &self.platform,
181 &self.base_path,
182 MAX_SSTABLE_SCAN_DEPTH,
183 )
184 .await
185 }
186 DiscoverySource::TableDirs(dirs) => {
187 let mut out = Vec::new();
188 for dir in dirs {
189 if !self.platform.fs().exists(dir).await? {
190 continue;
191 }
192 let mut entries = match self.platform.fs().read_dir(dir).await {
193 Ok(entries) => entries,
194 Err(_) => continue,
195 };
196 while let Some(entry) = entries.next_entry().await? {
197 let path = entry.path();
198 if let Some(fname) = path.file_name().and_then(|n| n.to_str()) {
199 if fname.ends_with("-Data.db") && !is_apple_double_sidecar(fname) {
200 out.push(path);
201 }
202 }
203 }
204 }
205 Ok(out)
206 }
207 }
208 }
209
210 fn table_key_for(&self, path: &Path, reader: &reader::SSTableReader) -> Option<String> {
213 match &self.discovery_source {
214 DiscoverySource::BasePath => {
215 let base_dir_name = self
216 .base_path
217 .file_name()
218 .and_then(|n| n.to_str())
219 .unwrap_or("");
220 base_path_table_key(path, base_dir_name, &reader.header().table_name)
221 }
222 DiscoverySource::TableDirs(_) => table_dir_table_key(path),
223 }
224 }
225
226 pub async fn refresh_tables(&self) -> Result<RefreshReport> {
231 let _refresh_guard = self.refresh_lock.lock().await;
237
238 let discovered_paths = self.discover_data_file_paths().await?;
240
241 let mut canon_cache: HashMap<PathBuf, PathBuf> =
249 HashMap::with_capacity(discovered_paths.len());
250 for p in &discovered_paths {
251 canon_cache.entry(p.clone()).or_insert_with(|| canon(p));
252 }
253
254 let mut held_canon: HashSet<PathBuf> = HashSet::new();
260 {
261 let readers = self.readers.read().await;
262 let table_readers = self.table_readers.read().await;
263 for r in readers.values().chain(table_readers.values().flatten()) {
264 let fp = r.file_path();
265 let c = canon_cache
266 .entry(fp.clone())
267 .or_insert_with(|| canon(&fp))
268 .clone();
269 held_canon.insert(c);
270 }
271 }
272
273 let canon_of = |p: &Path| -> PathBuf {
277 canon_cache
278 .get(p)
279 .cloned()
280 .unwrap_or_else(|| p.to_path_buf())
281 };
282
283 let discovered_canon: HashSet<PathBuf> =
284 discovered_paths.iter().map(|p| canon_of(p)).collect();
285
286 let added_paths: Vec<PathBuf> = discovered_paths
288 .iter()
289 .filter(|p| !held_canon.contains(&canon_of(p)))
290 .cloned()
291 .collect();
292
293 let mut opened: Vec<(
296 PathBuf,
297 SSTableId,
298 Option<String>,
299 Arc<reader::SSTableReader>,
300 )> = Vec::with_capacity(added_paths.len());
301 for path in &added_paths {
302 let Some(filename) = path.file_name().and_then(|n| n.to_str()) else {
303 continue;
304 };
305 let sstable_id = SSTableId::from_filename(filename);
306 let reader_arc = self.open_reader_with_schema(path).await?;
307 let key = self.table_key_for(path, &reader_arc);
308 opened.push((canon_of(path), sstable_id, key, reader_arc));
309 }
310
311 let mut readers = self.readers.write().await;
313 let mut table_readers = self.table_readers.write().await;
314
315 let before_canon: HashSet<PathBuf> = readers
319 .values()
320 .map(|r| canon_of(&r.file_path()))
321 .chain(
322 table_readers
323 .values()
324 .flatten()
325 .map(|r| canon_of(&r.file_path())),
326 )
327 .collect();
328
329 readers.retain(|_id, r| discovered_canon.contains(&canon_of(&r.file_path())));
330 for list in table_readers.values_mut() {
331 list.retain(|r| discovered_canon.contains(&canon_of(&r.file_path())));
332 }
333 table_readers.retain(|_key, list| !list.is_empty());
334
335 let after_removal_canon: HashSet<PathBuf> = readers
336 .values()
337 .map(|r| canon_of(&r.file_path()))
338 .chain(
339 table_readers
340 .values()
341 .flatten()
342 .map(|r| canon_of(&r.file_path())),
343 )
344 .collect();
345 let readers_removed = before_canon.difference(&after_removal_canon).count();
346
347 let mut readers_added = 0usize;
350 for (cpath, sstable_id, key, reader_arc) in opened {
351 if after_removal_canon.contains(&cpath) {
352 continue;
353 }
354 if readers.values().any(|r| canon_of(&r.file_path()) == cpath) {
357 continue;
358 }
359 readers.insert(sstable_id, Arc::clone(&reader_arc));
360 if let Some(key) = key {
361 table_readers.entry(key).or_default().push(reader_arc);
362 }
363 readers_added += 1;
364 }
365
366 let tables_scanned = table_readers.len();
367
368 Ok(RefreshReport {
369 tables_scanned,
370 readers_added,
371 readers_removed,
372 })
373 }
374}
375
376#[cfg(all(test, feature = "write-support", feature = "state_machine"))]
377mod tests {
378 use super::*;
379 use crate::schema::{Column, KeyColumn, TableSchema};
380 use crate::storage::write_engine::{
381 CellOperation, Mutation, PartitionKey, TableId, WriteEngine, WriteEngineConfig,
382 };
383 use crate::types::Value;
384 use crate::{Config, Platform};
385 use std::collections::HashMap;
386 use tempfile::TempDir;
387
388 fn users_schema() -> TableSchema {
389 TableSchema {
390 keyspace: "ks_ptr".to_string(),
391 table: "users".to_string(),
392 partition_keys: vec![KeyColumn {
393 name: "id".to_string(),
394 data_type: "int".to_string(),
395 position: 0,
396 }],
397 clustering_keys: vec![],
398 columns: vec![
399 Column {
400 name: "id".to_string(),
401 data_type: "int".to_string(),
402 nullable: false,
403 default: None,
404 is_static: false,
405 },
406 Column {
407 name: "value".to_string(),
408 data_type: "text".to_string(),
409 nullable: true,
410 default: None,
411 is_static: false,
412 },
413 ],
414 comments: HashMap::new(),
415 dropped_columns: HashMap::new(),
416 }
417 }
418
419 async fn flush_one_partition(data_dir: &Path, wal_dir: &Path, id: i32) {
420 let config = WriteEngineConfig::new(
421 data_dir.to_path_buf(),
422 wal_dir.to_path_buf(),
423 users_schema(),
424 );
425 let mut engine = WriteEngine::new(config).expect("write engine");
426 let table_id = TableId::new("ks_ptr", "users");
427 let pk = PartitionKey::single("id", Value::Integer(id));
428 let ops = vec![CellOperation::Write {
429 column: "value".to_string(),
430 value: Value::Text(format!("v{}", id)),
431 }];
432 engine
433 .write_async(Mutation::new(
434 table_id,
435 pk,
436 None,
437 ops,
438 1_000 + id as i64,
439 None,
440 ))
441 .await
442 .expect("write");
443 engine.flush().await.expect("flush");
444 }
445
446 async fn build_two_generations(root: &Path) -> PathBuf {
449 let data_dir = root.join("data");
450 let wal_dir = root.join("wal");
451 let config = WriteEngineConfig::new(data_dir.clone(), wal_dir, users_schema());
452 let mut engine = WriteEngine::new(config).expect("write engine");
453 for id in [1_i32, 2_i32] {
454 let table_id = TableId::new("ks_ptr", "users");
455 let pk = PartitionKey::single("id", Value::Integer(id));
456 let ops = vec![CellOperation::Write {
457 column: "value".to_string(),
458 value: Value::Text(format!("v{}", id)),
459 }];
460 engine
461 .write_async(Mutation::new(
462 table_id,
463 pk,
464 None,
465 ops,
466 1_000 + id as i64,
467 None,
468 ))
469 .await
470 .expect("write");
471 engine.flush().await.expect("flush");
472 }
473 data_dir.join("ks_ptr").join("users")
474 }
475
476 fn copy_generation(src_table_dir: &Path, dst_table_dir: &Path, gen: u32) {
477 std::fs::create_dir_all(dst_table_dir).expect("mkdir dst");
478 let prefix = format!("nb-{}-big-", gen);
479 let mut copied = 0;
480 for entry in std::fs::read_dir(src_table_dir).expect("read src") {
481 let entry = entry.expect("entry");
482 let name = entry.file_name();
483 let name = name.to_str().expect("utf8");
484 if name.starts_with(&prefix) {
485 std::fs::copy(entry.path(), dst_table_dir.join(name)).expect("copy");
486 copied += 1;
487 }
488 }
489 assert!(copied > 0, "copied generation {} components", gen);
490 }
491
492 fn scan_partition_ids(
493 rows: &[(crate::RowKey, crate::ScanRow)],
494 ) -> std::collections::BTreeSet<i32> {
495 rows.iter()
496 .map(|(key, _)| {
497 let b = key.as_bytes();
498 assert_eq!(b.len(), 4, "int pk is 4 bytes");
499 i32::from_be_bytes([b[0], b[1], b[2], b[3]])
500 })
501 .collect()
502 }
503
504 #[tokio::test]
515 async fn in_flight_scan_unaffected_by_concurrent_refresh() {
516 let src = TempDir::new().expect("tmp src");
517 let src_table_dir = build_two_generations(src.path()).await;
518
519 let live = TempDir::new().expect("tmp live");
520 let live_table_dir = live.path().join("ks_ptr").join("users");
521 copy_generation(&src_table_dir, &live_table_dir, 1);
522
523 let config = Config::default();
524 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
525 let manager = SSTableManager::new(
526 live.path(),
527 &config,
528 platform,
529 #[cfg(feature = "state_machine")]
530 None,
531 )
532 .await
533 .expect("manager");
534
535 let schema = users_schema();
536 let table_id = crate::types::TableId::new("ks_ptr.users");
537
538 let pre_rows = manager
541 .scan(&table_id, None, None, None, Some(&schema))
542 .await
543 .expect("pre scan");
544 assert_eq!(
545 scan_partition_ids(&pre_rows),
546 std::collections::BTreeSet::from([1]),
547 "pre-refresh scan content is gen-1 only"
548 );
549 let held_snapshot: Vec<Arc<reader::SSTableReader>> = {
550 let table_readers = manager.table_readers.read().await;
551 table_readers.values().flatten().map(Arc::clone).collect()
552 };
553 assert_eq!(
554 held_snapshot.len(),
555 1,
556 "in-flight scan holds one gen-1 reader"
557 );
558
559 copy_generation(&src_table_dir, &live_table_dir, 2);
561 let report = manager.refresh_tables().await.expect("refresh");
562 assert_eq!(report.readers_added, 1, "gen-2 added");
563
564 let post_snapshot: Vec<Arc<reader::SSTableReader>> = {
568 let table_readers = manager.table_readers.read().await;
569 table_readers.values().flatten().map(Arc::clone).collect()
570 };
571 assert_eq!(post_snapshot.len(), 2, "manager now holds both generations");
572 assert_eq!(
573 held_snapshot.len(),
574 1,
575 "the in-flight snapshot is untouched by refresh"
576 );
577 assert!(
578 post_snapshot
579 .iter()
580 .any(|r| Arc::ptr_eq(r, &held_snapshot[0])),
581 "the pre-refresh gen-1 reader Arc survives the refresh unchanged"
582 );
583 let gen2_reader = post_snapshot
584 .iter()
585 .find(|r| !Arc::ptr_eq(r, &held_snapshot[0]))
586 .expect("gen-2 reader present post-refresh");
587 assert!(
588 !Arc::ptr_eq(gen2_reader, &held_snapshot[0]),
589 "the added gen-2 reader is NOT part of the in-flight snapshot (isolation)"
590 );
591
592 let post_rows = manager
594 .scan(&table_id, None, None, None, Some(&schema))
595 .await
596 .expect("post scan");
597 assert_eq!(
598 scan_partition_ids(&post_rows),
599 std::collections::BTreeSet::from([1, 2]),
600 "post-refresh scan sees the new generation"
601 );
602 }
603
604 #[tokio::test]
608 async fn refresh_noop_preserves_reader_arc_identity() {
609 let tmp = TempDir::new().expect("tmp");
610 let data_dir = tmp.path().join("data");
611 let wal_dir = tmp.path().join("wal");
612 flush_one_partition(&data_dir, &wal_dir, 1).await;
613
614 let config = Config::default();
615 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
616 let manager = SSTableManager::new(
617 &data_dir,
618 &config,
619 platform,
620 #[cfg(feature = "state_machine")]
621 None,
622 )
623 .await
624 .expect("manager");
625
626 let before: Vec<Arc<reader::SSTableReader>> = {
628 let table_readers = manager.table_readers.read().await;
629 table_readers.values().flatten().map(Arc::clone).collect()
630 };
631 assert_eq!(before.len(), 1, "one generation expected pre-refresh");
632
633 let report = manager.refresh_tables().await.expect("refresh");
634 assert_eq!(report.readers_added, 0, "no-op: nothing added");
635 assert_eq!(report.readers_removed, 0, "no-op: nothing removed");
636
637 let after: Vec<Arc<reader::SSTableReader>> = {
638 let table_readers = manager.table_readers.read().await;
639 table_readers.values().flatten().map(Arc::clone).collect()
640 };
641 assert_eq!(after.len(), 1, "still one generation after no-op refresh");
642 assert!(
643 Arc::ptr_eq(&before[0], &after[0]),
644 "no-op refresh must keep the SAME reader Arc (warm state preserved)"
645 );
646 }
647
648 #[tokio::test]
652 async fn sequential_refreshes_each_apply_correctly() {
653 let src = TempDir::new().expect("tmp src");
654 let src_table_dir = build_two_generations(src.path()).await;
655
656 let live = TempDir::new().expect("tmp live");
657 let live_table_dir = live.path().join("ks_ptr").join("users");
658 copy_generation(&src_table_dir, &live_table_dir, 1);
659
660 let config = Config::default();
661 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
662 let manager = SSTableManager::new(
663 live.path(),
664 &config,
665 platform,
666 #[cfg(feature = "state_machine")]
667 None,
668 )
669 .await
670 .expect("manager");
671
672 let first = manager.refresh_tables().await.expect("first refresh");
674 assert_eq!(first.readers_added, 0, "first refresh adds nothing");
675 assert_eq!(first.readers_removed, 0, "first refresh removes nothing");
676
677 copy_generation(&src_table_dir, &live_table_dir, 2);
679 let second = manager.refresh_tables().await.expect("second refresh");
680 assert_eq!(second.readers_added, 1, "second refresh adds gen-2");
681 assert_eq!(second.readers_removed, 0, "second refresh removes nothing");
682
683 let held = {
684 let table_readers = manager.table_readers.read().await;
685 table_readers
686 .values()
687 .flatten()
688 .map(Arc::clone)
689 .collect::<Vec<_>>()
690 };
691 assert_eq!(
692 held.len(),
693 2,
694 "both generations held after the two refreshes"
695 );
696 }
697
698 #[tokio::test]
706 async fn concurrent_refreshes_serialize_no_reader_vanishes() {
707 let src = TempDir::new().expect("tmp src");
708 let src_table_dir = build_two_generations(src.path()).await;
709
710 let live = TempDir::new().expect("tmp live");
711 let live_table_dir = live.path().join("ks_ptr").join("users");
712 copy_generation(&src_table_dir, &live_table_dir, 1);
713
714 let config = Config::default();
715 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
716 let manager = SSTableManager::new(
717 live.path(),
718 &config,
719 platform,
720 #[cfg(feature = "state_machine")]
721 None,
722 )
723 .await
724 .expect("manager");
725
726 copy_generation(&src_table_dir, &live_table_dir, 2);
728 let (r0, r1, r2, r3) = tokio::join!(
729 manager.refresh_tables(),
730 manager.refresh_tables(),
731 manager.refresh_tables(),
732 manager.refresh_tables(),
733 );
734 let reports = [
735 r0.expect("refresh 0"),
736 r1.expect("refresh 1"),
737 r2.expect("refresh 2"),
738 r3.expect("refresh 3"),
739 ];
740
741 let total_added: usize = reports.iter().map(|r| r.readers_added).sum();
744 let total_removed: usize = reports.iter().map(|r| r.readers_removed).sum();
745 assert_eq!(
746 total_added, 1,
747 "exactly one serialized refresh adds gen-2 (others are no-ops)"
748 );
749 assert_eq!(
750 total_removed, 0,
751 "no refresh may remove a reader a concurrent refresh just added"
752 );
753
754 let held = {
756 let table_readers = manager.table_readers.read().await;
757 table_readers
758 .values()
759 .flatten()
760 .map(Arc::clone)
761 .collect::<Vec<_>>()
762 };
763 assert_eq!(
764 held.len(),
765 2,
766 "final reader set equals on-disk truth ({{gen-1, gen-2}})"
767 );
768 let scanned = manager
769 .scan(
770 &crate::types::TableId::new("ks_ptr.users"),
771 None,
772 None,
773 None,
774 Some(&users_schema()),
775 )
776 .await
777 .expect("post scan");
778 assert_eq!(
779 scan_partition_ids(&scanned),
780 std::collections::BTreeSet::from([1, 2]),
781 "both partitions queryable after concurrent refreshes"
782 );
783 }
784}