use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use super::{
extract_keyspace_and_table_name, extract_table_name, is_apple_double_sidecar, reader,
SSTableId, SSTableManager, MAX_SSTABLE_SCAN_DEPTH,
};
use crate::Result;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RefreshReport {
pub tables_scanned: usize,
pub readers_added: usize,
pub readers_removed: usize,
}
#[derive(Debug, Clone)]
pub(crate) enum DiscoverySource {
BasePath,
TableDirs(Vec<PathBuf>),
}
pub(crate) fn base_path_table_key(
path: &Path,
base_dir_name: &str,
header_table_name: &str,
) -> Option<String> {
if let Some((keyspace, table_name)) = extract_keyspace_and_table_name(path) {
if table_name.as_str() != base_dir_name {
return Some(format!("{}.{}", keyspace, table_name));
}
}
if let Some(name) = extract_table_name(path) {
if name.as_str() != base_dir_name {
return Some(name);
}
}
if header_table_name != "test_table" && !header_table_name.is_empty() {
return Some(header_table_name.to_string());
}
None
}
pub(crate) fn table_dir_table_key(path: &Path) -> Option<String> {
if let Some((keyspace, table_name)) = extract_keyspace_and_table_name(path) {
Some(format!("{}.{}", keyspace, table_name))
} else {
extract_table_name(path)
}
}
fn canon(p: &Path) -> PathBuf {
std::fs::canonicalize(p).unwrap_or_else(|_| p.to_path_buf())
}
impl SSTableManager {
pub(crate) async fn open_reader_with_schema(
&self,
path: &Path,
) -> Result<Arc<reader::SSTableReader>> {
#[cfg_attr(not(feature = "state_machine"), allow(unused_mut))]
let mut reader = reader::SSTableReader::open_with_cache(
path,
&self.config,
self.platform.clone(),
self.chunk_cache.clone(),
)
.await?;
#[cfg(feature = "state_machine")]
{
let schema_reg_guard = self.schema_registry.read().await;
if let Some(ref registry_rwlock) = *schema_reg_guard {
#[allow(deprecated)]
reader.set_schema_registry(Arc::clone(registry_rwlock));
reader.resolve_registry_schema().await;
let schema_registry = registry_rwlock.read().await;
let udt_registry_lock = schema_registry.get_udt_registry();
let udt_registry = udt_registry_lock.read().await.clone();
reader.set_udt_registry(udt_registry);
}
}
Ok(Arc::new(reader))
}
async fn discover_data_file_paths(&self) -> Result<Vec<PathBuf>> {
match &self.discovery_source {
DiscoverySource::BasePath => {
if !self.platform.fs().exists(&self.base_path).await? {
return Ok(Vec::new());
}
SSTableManager::find_data_files(
&self.platform,
&self.base_path,
MAX_SSTABLE_SCAN_DEPTH,
)
.await
}
DiscoverySource::TableDirs(dirs) => {
let mut out = Vec::new();
for dir in dirs {
if !self.platform.fs().exists(dir).await? {
continue;
}
let mut entries = match self.platform.fs().read_dir(dir).await {
Ok(entries) => entries,
Err(_) => continue,
};
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if let Some(fname) = path.file_name().and_then(|n| n.to_str()) {
if fname.ends_with("-Data.db") && !is_apple_double_sidecar(fname) {
out.push(path);
}
}
}
}
Ok(out)
}
}
}
fn table_key_for(&self, path: &Path, reader: &reader::SSTableReader) -> Option<String> {
match &self.discovery_source {
DiscoverySource::BasePath => {
let base_dir_name = self
.base_path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("");
base_path_table_key(path, base_dir_name, &reader.header().table_name)
}
DiscoverySource::TableDirs(_) => table_dir_table_key(path),
}
}
pub async fn refresh_tables(&self) -> Result<RefreshReport> {
let _refresh_guard = self.refresh_lock.lock().await;
let discovered_paths = self.discover_data_file_paths().await?;
let mut canon_cache: HashMap<PathBuf, PathBuf> =
HashMap::with_capacity(discovered_paths.len());
for p in &discovered_paths {
canon_cache.entry(p.clone()).or_insert_with(|| canon(p));
}
let mut held_canon: HashSet<PathBuf> = HashSet::new();
{
let readers = self.readers.read().await;
let table_readers = self.table_readers.read().await;
for r in readers.values().chain(table_readers.values().flatten()) {
let fp = r.file_path();
let c = canon_cache
.entry(fp.clone())
.or_insert_with(|| canon(&fp))
.clone();
held_canon.insert(c);
}
}
let canon_of = |p: &Path| -> PathBuf {
canon_cache
.get(p)
.cloned()
.unwrap_or_else(|| p.to_path_buf())
};
let discovered_canon: HashSet<PathBuf> =
discovered_paths.iter().map(|p| canon_of(p)).collect();
let added_paths: Vec<PathBuf> = discovered_paths
.iter()
.filter(|p| !held_canon.contains(&canon_of(p)))
.cloned()
.collect();
let mut opened: Vec<(
PathBuf,
SSTableId,
Option<String>,
Arc<reader::SSTableReader>,
)> = Vec::with_capacity(added_paths.len());
for path in &added_paths {
let Some(filename) = path.file_name().and_then(|n| n.to_str()) else {
continue;
};
let sstable_id = SSTableId::from_filename(filename);
let reader_arc = self.open_reader_with_schema(path).await?;
let key = self.table_key_for(path, &reader_arc);
opened.push((canon_of(path), sstable_id, key, reader_arc));
}
let mut readers = self.readers.write().await;
let mut table_readers = self.table_readers.write().await;
let before_canon: HashSet<PathBuf> = readers
.values()
.map(|r| canon_of(&r.file_path()))
.chain(
table_readers
.values()
.flatten()
.map(|r| canon_of(&r.file_path())),
)
.collect();
{
let mut seen: HashSet<*const reader::SSTableReader> = HashSet::new();
for r in readers.values().chain(table_readers.values().flatten()) {
if discovered_canon.contains(&canon_of(&r.file_path())) {
continue; }
if seen.insert(Arc::as_ptr(r)) {
r.invalidate_key_cache_entries();
}
}
}
readers.retain(|_id, r| discovered_canon.contains(&canon_of(&r.file_path())));
for list in table_readers.values_mut() {
list.retain(|r| discovered_canon.contains(&canon_of(&r.file_path())));
}
table_readers.retain(|_key, list| !list.is_empty());
let after_removal_canon: HashSet<PathBuf> = readers
.values()
.map(|r| canon_of(&r.file_path()))
.chain(
table_readers
.values()
.flatten()
.map(|r| canon_of(&r.file_path())),
)
.collect();
let readers_removed = before_canon.difference(&after_removal_canon).count();
let mut readers_added = 0usize;
for (cpath, sstable_id, key, reader_arc) in opened {
if after_removal_canon.contains(&cpath) {
continue;
}
if readers.values().any(|r| canon_of(&r.file_path()) == cpath) {
continue;
}
readers.insert(sstable_id, Arc::clone(&reader_arc));
if let Some(key) = key {
table_readers.entry(key).or_default().push(reader_arc);
}
readers_added += 1;
}
let tables_scanned = table_readers.len();
Ok(RefreshReport {
tables_scanned,
readers_added,
readers_removed,
})
}
}
#[cfg(all(test, feature = "write-support", feature = "state_machine"))]
mod tests {
use super::*;
use crate::schema::{Column, KeyColumn, TableSchema};
use crate::storage::write_engine::{
CellOperation, Mutation, PartitionKey, TableId, WriteEngine, WriteEngineConfig,
};
use crate::types::Value;
use crate::{Config, Platform};
use std::collections::HashMap;
use tempfile::TempDir;
fn users_schema() -> TableSchema {
TableSchema {
keyspace: "ks_ptr".to_string(),
table: "users".to_string(),
partition_keys: vec![KeyColumn {
name: "id".to_string(),
data_type: "int".to_string(),
position: 0,
}],
clustering_keys: vec![],
columns: vec![
Column {
name: "id".to_string(),
data_type: "int".to_string(),
nullable: false,
default: None,
is_static: false,
},
Column {
name: "value".to_string(),
data_type: "text".to_string(),
nullable: true,
default: None,
is_static: false,
},
],
comments: HashMap::new(),
dropped_columns: HashMap::new(),
}
}
async fn flush_one_partition(data_dir: &Path, wal_dir: &Path, id: i32) {
let config = WriteEngineConfig::new(
data_dir.to_path_buf(),
wal_dir.to_path_buf(),
users_schema(),
);
let mut engine = WriteEngine::new(config).expect("write engine");
let table_id = TableId::new("ks_ptr", "users");
let pk = PartitionKey::single("id", Value::Integer(id));
let ops = vec![CellOperation::Write {
column: "value".to_string(),
value: Value::text(format!("v{}", id)),
}];
engine
.write_async(Mutation::new(
table_id,
pk,
None,
ops,
1_000 + id as i64,
None,
))
.await
.expect("write");
engine.flush().await.expect("flush");
}
async fn build_two_generations(root: &Path) -> PathBuf {
let data_dir = root.join("data");
let wal_dir = root.join("wal");
let config = WriteEngineConfig::new(data_dir.clone(), wal_dir, users_schema());
let mut engine = WriteEngine::new(config).expect("write engine");
for id in [1_i32, 2_i32] {
let table_id = TableId::new("ks_ptr", "users");
let pk = PartitionKey::single("id", Value::Integer(id));
let ops = vec![CellOperation::Write {
column: "value".to_string(),
value: Value::text(format!("v{}", id)),
}];
engine
.write_async(Mutation::new(
table_id,
pk,
None,
ops,
1_000 + id as i64,
None,
))
.await
.expect("write");
engine.flush().await.expect("flush");
}
data_dir.join("ks_ptr").join("users")
}
fn copy_generation(src_table_dir: &Path, dst_table_dir: &Path, gen: u32) {
std::fs::create_dir_all(dst_table_dir).expect("mkdir dst");
let prefix = format!("nb-{}-big-", gen);
let mut copied = 0;
for entry in std::fs::read_dir(src_table_dir).expect("read src") {
let entry = entry.expect("entry");
let name = entry.file_name();
let name = name.to_str().expect("utf8");
if name.starts_with(&prefix) {
std::fs::copy(entry.path(), dst_table_dir.join(name)).expect("copy");
copied += 1;
}
}
assert!(copied > 0, "copied generation {} components", gen);
}
fn scan_partition_ids(
rows: &[(crate::RowKey, crate::ScanRow)],
) -> std::collections::BTreeSet<i32> {
rows.iter()
.map(|(key, _)| {
let b = key.as_bytes();
assert_eq!(b.len(), 4, "int pk is 4 bytes");
i32::from_be_bytes([b[0], b[1], b[2], b[3]])
})
.collect()
}
#[tokio::test]
async fn in_flight_scan_unaffected_by_concurrent_refresh() {
let src = TempDir::new().expect("tmp src");
let src_table_dir = build_two_generations(src.path()).await;
let live = TempDir::new().expect("tmp live");
let live_table_dir = live.path().join("ks_ptr").join("users");
copy_generation(&src_table_dir, &live_table_dir, 1);
let config = Config::default();
let platform = Arc::new(Platform::new(&config).await.expect("platform"));
let manager = SSTableManager::new(
live.path(),
&config,
platform,
#[cfg(feature = "state_machine")]
None,
)
.await
.expect("manager");
let schema = users_schema();
let table_id = crate::types::TableId::new("ks_ptr.users");
let pre_rows = manager
.scan(&table_id, None, None, None, Some(&schema))
.await
.expect("pre scan");
assert_eq!(
scan_partition_ids(&pre_rows),
std::collections::BTreeSet::from([1]),
"pre-refresh scan content is gen-1 only"
);
let held_snapshot: Vec<Arc<reader::SSTableReader>> = {
let table_readers = manager.table_readers.read().await;
table_readers.values().flatten().map(Arc::clone).collect()
};
assert_eq!(
held_snapshot.len(),
1,
"in-flight scan holds one gen-1 reader"
);
copy_generation(&src_table_dir, &live_table_dir, 2);
let report = manager.refresh_tables().await.expect("refresh");
assert_eq!(report.readers_added, 1, "gen-2 added");
let post_snapshot: Vec<Arc<reader::SSTableReader>> = {
let table_readers = manager.table_readers.read().await;
table_readers.values().flatten().map(Arc::clone).collect()
};
assert_eq!(post_snapshot.len(), 2, "manager now holds both generations");
assert_eq!(
held_snapshot.len(),
1,
"the in-flight snapshot is untouched by refresh"
);
assert!(
post_snapshot
.iter()
.any(|r| Arc::ptr_eq(r, &held_snapshot[0])),
"the pre-refresh gen-1 reader Arc survives the refresh unchanged"
);
let gen2_reader = post_snapshot
.iter()
.find(|r| !Arc::ptr_eq(r, &held_snapshot[0]))
.expect("gen-2 reader present post-refresh");
assert!(
!Arc::ptr_eq(gen2_reader, &held_snapshot[0]),
"the added gen-2 reader is NOT part of the in-flight snapshot (isolation)"
);
let post_rows = manager
.scan(&table_id, None, None, None, Some(&schema))
.await
.expect("post scan");
assert_eq!(
scan_partition_ids(&post_rows),
std::collections::BTreeSet::from([1, 2]),
"post-refresh scan sees the new generation"
);
}
#[tokio::test]
async fn refresh_noop_preserves_reader_arc_identity() {
let tmp = TempDir::new().expect("tmp");
let data_dir = tmp.path().join("data");
let wal_dir = tmp.path().join("wal");
flush_one_partition(&data_dir, &wal_dir, 1).await;
let config = Config::default();
let platform = Arc::new(Platform::new(&config).await.expect("platform"));
let manager = SSTableManager::new(
&data_dir,
&config,
platform,
#[cfg(feature = "state_machine")]
None,
)
.await
.expect("manager");
let before: Vec<Arc<reader::SSTableReader>> = {
let table_readers = manager.table_readers.read().await;
table_readers.values().flatten().map(Arc::clone).collect()
};
assert_eq!(before.len(), 1, "one generation expected pre-refresh");
let report = manager.refresh_tables().await.expect("refresh");
assert_eq!(report.readers_added, 0, "no-op: nothing added");
assert_eq!(report.readers_removed, 0, "no-op: nothing removed");
let after: Vec<Arc<reader::SSTableReader>> = {
let table_readers = manager.table_readers.read().await;
table_readers.values().flatten().map(Arc::clone).collect()
};
assert_eq!(after.len(), 1, "still one generation after no-op refresh");
assert!(
Arc::ptr_eq(&before[0], &after[0]),
"no-op refresh must keep the SAME reader Arc (warm state preserved)"
);
}
#[tokio::test]
async fn sequential_refreshes_each_apply_correctly() {
let src = TempDir::new().expect("tmp src");
let src_table_dir = build_two_generations(src.path()).await;
let live = TempDir::new().expect("tmp live");
let live_table_dir = live.path().join("ks_ptr").join("users");
copy_generation(&src_table_dir, &live_table_dir, 1);
let config = Config::default();
let platform = Arc::new(Platform::new(&config).await.expect("platform"));
let manager = SSTableManager::new(
live.path(),
&config,
platform,
#[cfg(feature = "state_machine")]
None,
)
.await
.expect("manager");
let first = manager.refresh_tables().await.expect("first refresh");
assert_eq!(first.readers_added, 0, "first refresh adds nothing");
assert_eq!(first.readers_removed, 0, "first refresh removes nothing");
copy_generation(&src_table_dir, &live_table_dir, 2);
let second = manager.refresh_tables().await.expect("second refresh");
assert_eq!(second.readers_added, 1, "second refresh adds gen-2");
assert_eq!(second.readers_removed, 0, "second refresh removes nothing");
let held = {
let table_readers = manager.table_readers.read().await;
table_readers
.values()
.flatten()
.map(Arc::clone)
.collect::<Vec<_>>()
};
assert_eq!(
held.len(),
2,
"both generations held after the two refreshes"
);
}
#[tokio::test]
async fn concurrent_refreshes_serialize_no_reader_vanishes() {
let src = TempDir::new().expect("tmp src");
let src_table_dir = build_two_generations(src.path()).await;
let live = TempDir::new().expect("tmp live");
let live_table_dir = live.path().join("ks_ptr").join("users");
copy_generation(&src_table_dir, &live_table_dir, 1);
let config = Config::default();
let platform = Arc::new(Platform::new(&config).await.expect("platform"));
let manager = SSTableManager::new(
live.path(),
&config,
platform,
#[cfg(feature = "state_machine")]
None,
)
.await
.expect("manager");
copy_generation(&src_table_dir, &live_table_dir, 2);
let (r0, r1, r2, r3) = tokio::join!(
manager.refresh_tables(),
manager.refresh_tables(),
manager.refresh_tables(),
manager.refresh_tables(),
);
let reports = [
r0.expect("refresh 0"),
r1.expect("refresh 1"),
r2.expect("refresh 2"),
r3.expect("refresh 3"),
];
let total_added: usize = reports.iter().map(|r| r.readers_added).sum();
let total_removed: usize = reports.iter().map(|r| r.readers_removed).sum();
assert_eq!(
total_added, 1,
"exactly one serialized refresh adds gen-2 (others are no-ops)"
);
assert_eq!(
total_removed, 0,
"no refresh may remove a reader a concurrent refresh just added"
);
let held = {
let table_readers = manager.table_readers.read().await;
table_readers
.values()
.flatten()
.map(Arc::clone)
.collect::<Vec<_>>()
};
assert_eq!(
held.len(),
2,
"final reader set equals on-disk truth ({{gen-1, gen-2}})"
);
let scanned = manager
.scan(
&crate::types::TableId::new("ks_ptr.users"),
None,
None,
None,
Some(&users_schema()),
)
.await
.expect("post scan");
assert_eq!(
scan_partition_ids(&scanned),
std::collections::BTreeSet::from([1, 2]),
"both partitions queryable after concurrent refreshes"
);
}
}