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 {
337 let mut seen: HashSet<*const reader::SSTableReader> = HashSet::new();
338 for r in readers.values().chain(table_readers.values().flatten()) {
339 if discovered_canon.contains(&canon_of(&r.file_path())) {
340 continue; }
342 if seen.insert(Arc::as_ptr(r)) {
343 r.invalidate_key_cache_entries();
344 }
345 }
346 }
347
348 readers.retain(|_id, r| discovered_canon.contains(&canon_of(&r.file_path())));
349 for list in table_readers.values_mut() {
350 list.retain(|r| discovered_canon.contains(&canon_of(&r.file_path())));
351 }
352 table_readers.retain(|_key, list| !list.is_empty());
353
354 let after_removal_canon: HashSet<PathBuf> = readers
355 .values()
356 .map(|r| canon_of(&r.file_path()))
357 .chain(
358 table_readers
359 .values()
360 .flatten()
361 .map(|r| canon_of(&r.file_path())),
362 )
363 .collect();
364 let readers_removed = before_canon.difference(&after_removal_canon).count();
365
366 let mut readers_added = 0usize;
369 for (cpath, sstable_id, key, reader_arc) in opened {
370 if after_removal_canon.contains(&cpath) {
371 continue;
372 }
373 if readers.values().any(|r| canon_of(&r.file_path()) == cpath) {
376 continue;
377 }
378 readers.insert(sstable_id, Arc::clone(&reader_arc));
379 if let Some(key) = key {
380 table_readers.entry(key).or_default().push(reader_arc);
381 }
382 readers_added += 1;
383 }
384
385 let tables_scanned = table_readers.len();
386
387 Ok(RefreshReport {
388 tables_scanned,
389 readers_added,
390 readers_removed,
391 })
392 }
393}
394
395#[cfg(all(test, feature = "write-support", feature = "state_machine"))]
396mod tests {
397 use super::*;
398 use crate::schema::{Column, KeyColumn, TableSchema};
399 use crate::storage::write_engine::{
400 CellOperation, Mutation, PartitionKey, TableId, WriteEngine, WriteEngineConfig,
401 };
402 use crate::types::Value;
403 use crate::{Config, Platform};
404 use std::collections::HashMap;
405 use tempfile::TempDir;
406
407 fn users_schema() -> TableSchema {
408 TableSchema {
409 keyspace: "ks_ptr".to_string(),
410 table: "users".to_string(),
411 partition_keys: vec![KeyColumn {
412 name: "id".to_string(),
413 data_type: "int".to_string(),
414 position: 0,
415 }],
416 clustering_keys: vec![],
417 columns: vec![
418 Column {
419 name: "id".to_string(),
420 data_type: "int".to_string(),
421 nullable: false,
422 default: None,
423 is_static: false,
424 },
425 Column {
426 name: "value".to_string(),
427 data_type: "text".to_string(),
428 nullable: true,
429 default: None,
430 is_static: false,
431 },
432 ],
433 comments: HashMap::new(),
434 dropped_columns: HashMap::new(),
435 }
436 }
437
438 async fn flush_one_partition(data_dir: &Path, wal_dir: &Path, id: i32) {
439 let config = WriteEngineConfig::new(
440 data_dir.to_path_buf(),
441 wal_dir.to_path_buf(),
442 users_schema(),
443 );
444 let mut engine = WriteEngine::new(config).expect("write engine");
445 let table_id = TableId::new("ks_ptr", "users");
446 let pk = PartitionKey::single("id", Value::Integer(id));
447 let ops = vec![CellOperation::Write {
448 column: "value".to_string(),
449 value: Value::text(format!("v{}", id)),
450 }];
451 engine
452 .write_async(Mutation::new(
453 table_id,
454 pk,
455 None,
456 ops,
457 1_000 + id as i64,
458 None,
459 ))
460 .await
461 .expect("write");
462 engine.flush().await.expect("flush");
463 }
464
465 async fn build_two_generations(root: &Path) -> PathBuf {
468 let data_dir = root.join("data");
469 let wal_dir = root.join("wal");
470 let config = WriteEngineConfig::new(data_dir.clone(), wal_dir, users_schema());
471 let mut engine = WriteEngine::new(config).expect("write engine");
472 for id in [1_i32, 2_i32] {
473 let table_id = TableId::new("ks_ptr", "users");
474 let pk = PartitionKey::single("id", Value::Integer(id));
475 let ops = vec![CellOperation::Write {
476 column: "value".to_string(),
477 value: Value::text(format!("v{}", id)),
478 }];
479 engine
480 .write_async(Mutation::new(
481 table_id,
482 pk,
483 None,
484 ops,
485 1_000 + id as i64,
486 None,
487 ))
488 .await
489 .expect("write");
490 engine.flush().await.expect("flush");
491 }
492 data_dir.join("ks_ptr").join("users")
493 }
494
495 fn copy_generation(src_table_dir: &Path, dst_table_dir: &Path, gen: u32) {
496 std::fs::create_dir_all(dst_table_dir).expect("mkdir dst");
497 let prefix = format!("nb-{}-big-", gen);
498 let mut copied = 0;
499 for entry in std::fs::read_dir(src_table_dir).expect("read src") {
500 let entry = entry.expect("entry");
501 let name = entry.file_name();
502 let name = name.to_str().expect("utf8");
503 if name.starts_with(&prefix) {
504 std::fs::copy(entry.path(), dst_table_dir.join(name)).expect("copy");
505 copied += 1;
506 }
507 }
508 assert!(copied > 0, "copied generation {} components", gen);
509 }
510
511 fn scan_partition_ids(
512 rows: &[(crate::RowKey, crate::ScanRow)],
513 ) -> std::collections::BTreeSet<i32> {
514 rows.iter()
515 .map(|(key, _)| {
516 let b = key.as_bytes();
517 assert_eq!(b.len(), 4, "int pk is 4 bytes");
518 i32::from_be_bytes([b[0], b[1], b[2], b[3]])
519 })
520 .collect()
521 }
522
523 #[tokio::test]
534 async fn in_flight_scan_unaffected_by_concurrent_refresh() {
535 let src = TempDir::new().expect("tmp src");
536 let src_table_dir = build_two_generations(src.path()).await;
537
538 let live = TempDir::new().expect("tmp live");
539 let live_table_dir = live.path().join("ks_ptr").join("users");
540 copy_generation(&src_table_dir, &live_table_dir, 1);
541
542 let config = Config::default();
543 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
544 let manager = SSTableManager::new(
545 live.path(),
546 &config,
547 platform,
548 #[cfg(feature = "state_machine")]
549 None,
550 )
551 .await
552 .expect("manager");
553
554 let schema = users_schema();
555 let table_id = crate::types::TableId::new("ks_ptr.users");
556
557 let pre_rows = manager
560 .scan(&table_id, None, None, None, Some(&schema))
561 .await
562 .expect("pre scan");
563 assert_eq!(
564 scan_partition_ids(&pre_rows),
565 std::collections::BTreeSet::from([1]),
566 "pre-refresh scan content is gen-1 only"
567 );
568 let held_snapshot: Vec<Arc<reader::SSTableReader>> = {
569 let table_readers = manager.table_readers.read().await;
570 table_readers.values().flatten().map(Arc::clone).collect()
571 };
572 assert_eq!(
573 held_snapshot.len(),
574 1,
575 "in-flight scan holds one gen-1 reader"
576 );
577
578 copy_generation(&src_table_dir, &live_table_dir, 2);
580 let report = manager.refresh_tables().await.expect("refresh");
581 assert_eq!(report.readers_added, 1, "gen-2 added");
582
583 let post_snapshot: Vec<Arc<reader::SSTableReader>> = {
587 let table_readers = manager.table_readers.read().await;
588 table_readers.values().flatten().map(Arc::clone).collect()
589 };
590 assert_eq!(post_snapshot.len(), 2, "manager now holds both generations");
591 assert_eq!(
592 held_snapshot.len(),
593 1,
594 "the in-flight snapshot is untouched by refresh"
595 );
596 assert!(
597 post_snapshot
598 .iter()
599 .any(|r| Arc::ptr_eq(r, &held_snapshot[0])),
600 "the pre-refresh gen-1 reader Arc survives the refresh unchanged"
601 );
602 let gen2_reader = post_snapshot
603 .iter()
604 .find(|r| !Arc::ptr_eq(r, &held_snapshot[0]))
605 .expect("gen-2 reader present post-refresh");
606 assert!(
607 !Arc::ptr_eq(gen2_reader, &held_snapshot[0]),
608 "the added gen-2 reader is NOT part of the in-flight snapshot (isolation)"
609 );
610
611 let post_rows = manager
613 .scan(&table_id, None, None, None, Some(&schema))
614 .await
615 .expect("post scan");
616 assert_eq!(
617 scan_partition_ids(&post_rows),
618 std::collections::BTreeSet::from([1, 2]),
619 "post-refresh scan sees the new generation"
620 );
621 }
622
623 #[tokio::test]
627 async fn refresh_noop_preserves_reader_arc_identity() {
628 let tmp = TempDir::new().expect("tmp");
629 let data_dir = tmp.path().join("data");
630 let wal_dir = tmp.path().join("wal");
631 flush_one_partition(&data_dir, &wal_dir, 1).await;
632
633 let config = Config::default();
634 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
635 let manager = SSTableManager::new(
636 &data_dir,
637 &config,
638 platform,
639 #[cfg(feature = "state_machine")]
640 None,
641 )
642 .await
643 .expect("manager");
644
645 let before: Vec<Arc<reader::SSTableReader>> = {
647 let table_readers = manager.table_readers.read().await;
648 table_readers.values().flatten().map(Arc::clone).collect()
649 };
650 assert_eq!(before.len(), 1, "one generation expected pre-refresh");
651
652 let report = manager.refresh_tables().await.expect("refresh");
653 assert_eq!(report.readers_added, 0, "no-op: nothing added");
654 assert_eq!(report.readers_removed, 0, "no-op: nothing removed");
655
656 let after: Vec<Arc<reader::SSTableReader>> = {
657 let table_readers = manager.table_readers.read().await;
658 table_readers.values().flatten().map(Arc::clone).collect()
659 };
660 assert_eq!(after.len(), 1, "still one generation after no-op refresh");
661 assert!(
662 Arc::ptr_eq(&before[0], &after[0]),
663 "no-op refresh must keep the SAME reader Arc (warm state preserved)"
664 );
665 }
666
667 #[tokio::test]
671 async fn sequential_refreshes_each_apply_correctly() {
672 let src = TempDir::new().expect("tmp src");
673 let src_table_dir = build_two_generations(src.path()).await;
674
675 let live = TempDir::new().expect("tmp live");
676 let live_table_dir = live.path().join("ks_ptr").join("users");
677 copy_generation(&src_table_dir, &live_table_dir, 1);
678
679 let config = Config::default();
680 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
681 let manager = SSTableManager::new(
682 live.path(),
683 &config,
684 platform,
685 #[cfg(feature = "state_machine")]
686 None,
687 )
688 .await
689 .expect("manager");
690
691 let first = manager.refresh_tables().await.expect("first refresh");
693 assert_eq!(first.readers_added, 0, "first refresh adds nothing");
694 assert_eq!(first.readers_removed, 0, "first refresh removes nothing");
695
696 copy_generation(&src_table_dir, &live_table_dir, 2);
698 let second = manager.refresh_tables().await.expect("second refresh");
699 assert_eq!(second.readers_added, 1, "second refresh adds gen-2");
700 assert_eq!(second.readers_removed, 0, "second refresh removes nothing");
701
702 let held = {
703 let table_readers = manager.table_readers.read().await;
704 table_readers
705 .values()
706 .flatten()
707 .map(Arc::clone)
708 .collect::<Vec<_>>()
709 };
710 assert_eq!(
711 held.len(),
712 2,
713 "both generations held after the two refreshes"
714 );
715 }
716
717 #[tokio::test]
725 async fn concurrent_refreshes_serialize_no_reader_vanishes() {
726 let src = TempDir::new().expect("tmp src");
727 let src_table_dir = build_two_generations(src.path()).await;
728
729 let live = TempDir::new().expect("tmp live");
730 let live_table_dir = live.path().join("ks_ptr").join("users");
731 copy_generation(&src_table_dir, &live_table_dir, 1);
732
733 let config = Config::default();
734 let platform = Arc::new(Platform::new(&config).await.expect("platform"));
735 let manager = SSTableManager::new(
736 live.path(),
737 &config,
738 platform,
739 #[cfg(feature = "state_machine")]
740 None,
741 )
742 .await
743 .expect("manager");
744
745 copy_generation(&src_table_dir, &live_table_dir, 2);
747 let (r0, r1, r2, r3) = tokio::join!(
748 manager.refresh_tables(),
749 manager.refresh_tables(),
750 manager.refresh_tables(),
751 manager.refresh_tables(),
752 );
753 let reports = [
754 r0.expect("refresh 0"),
755 r1.expect("refresh 1"),
756 r2.expect("refresh 2"),
757 r3.expect("refresh 3"),
758 ];
759
760 let total_added: usize = reports.iter().map(|r| r.readers_added).sum();
763 let total_removed: usize = reports.iter().map(|r| r.readers_removed).sum();
764 assert_eq!(
765 total_added, 1,
766 "exactly one serialized refresh adds gen-2 (others are no-ops)"
767 );
768 assert_eq!(
769 total_removed, 0,
770 "no refresh may remove a reader a concurrent refresh just added"
771 );
772
773 let held = {
775 let table_readers = manager.table_readers.read().await;
776 table_readers
777 .values()
778 .flatten()
779 .map(Arc::clone)
780 .collect::<Vec<_>>()
781 };
782 assert_eq!(
783 held.len(),
784 2,
785 "final reader set equals on-disk truth ({{gen-1, gen-2}})"
786 );
787 let scanned = manager
788 .scan(
789 &crate::types::TableId::new("ks_ptr.users"),
790 None,
791 None,
792 None,
793 Some(&users_schema()),
794 )
795 .await
796 .expect("post scan");
797 assert_eq!(
798 scan_partition_ids(&scanned),
799 std::collections::BTreeSet::from([1, 2]),
800 "both partitions queryable after concurrent refreshes"
801 );
802 }
803}