use std::{
collections::{HashMap, HashSet},
ops::Bound,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::{Duration, Instant},
};
use object_store::{ObjectStore, path::Path};
use slatedb::{CloseReason, Db, DbReader, DbStatus, DbTransaction, IsolationLevel};
use tokio::sync::watch;
use tracing::{info, warn};
use crate::{
catalog::{
CatalogSnapshot, ColumnId, ColumnOrder, DataFileId, DataFileInfo, FileIndexEntry, IndexDef,
IndexEntry, IndexId, IndexInfo, IndexState, RecentRow, RowHolder, RowLocation, SnapshotId,
TableId,
census::{
CensusRequest, CompactStoreReport, CompactStoreRequest, CompactionTarget, LiveCount,
MergeOutcome, StoreCensus, StoreObjects, SubspaceCensus, SubspaceMerge, SubspaceName,
},
inline::{InlineScanKind, materialize_inline_rows},
projection::{
ProjectionCache, cache_epoch, cached_head_view, held_head_view, install_head_view_at,
install_shared_current_entities, shared_current_entities,
},
scoped_read,
},
error::{Error, Result},
store::{
census::{self as store_census, SegmentSize},
compaction::{self as store_compaction, MergeEnd},
handle::{ReadHandle, ReadSession, ScanShape},
index_encoding::{
CanonicalKey, Direction, IndexKeyValue, NullOrder, encode_ordered_values,
},
inline as store_inline,
key::{
IndexKey, IndexKind, InlineOperation, Key, Subspace, index_index_prefix,
index_kind_prefix, subspace_prefix,
},
open::{self, StoreBuilder},
proto::HeadValue,
},
transaction::{MigrationReport, Transaction, commit, index_maintenance, migration},
};
fn measure(subspace: SubspaceName, segment: Option<&SegmentSize>) -> SubspaceCensus {
SubspaceCensus {
subspace,
bytes: segment.map_or(0, |segment| segment.bytes),
l0_ssts: segment.map_or(0, |segment| segment.l0_ssts),
sorted_runs: segment.map_or(0, |segment| segment.sorted_runs),
sorted_run_ssts: segment.map_or(0, |segment| segment.sorted_run_ssts),
live: None,
}
}
async fn count_live_entries(
handle: ReadHandle<'_>,
subspaces: &mut [SubspaceCensus],
) -> Result<()> {
for measured in subspaces {
let Some(subspace) = measured.subspace.subspace() else {
continue;
};
let tally = store_census::scan_live(handle, subspace).await?;
measured.live = Some(LiveCount {
keys: tally.keys,
key_bytes: tally.key_bytes,
value_bytes: tally.value_bytes,
scheduled_files: tally.scheduled_files,
});
}
Ok(())
}
fn bytes_of(census: &StoreCensus, subspace: &SubspaceName) -> u64 {
census
.subspaces
.iter()
.find(|measured| &measured.subspace == subspace)
.map_or(0, |measured| measured.bytes)
}
const BUILD_STEP_ENTRIES: usize = 1_000_000;
const BUILD_DERIVATION_ATTEMPTS: usize = 8;
fn requested_orders(orders: &[ColumnOrder], columns: usize) -> (Vec<Direction>, Vec<NullOrder>) {
(0..columns)
.map(|position| {
orders
.get(position)
.map_or((Direction::Ascending, NullOrder::Last), |order| {
(order.direction, order.nulls)
})
})
.unzip()
}
enum BuildProgress {
Ready,
Conflicted,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct MaintenanceRequest {
pub sweep_orphaned_index_entries: bool,
pub batch_size: usize,
}
impl Default for MaintenanceRequest {
fn default() -> Self {
Self {
sweep_orphaned_index_entries: true,
batch_size: 1024,
}
}
}
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct MigrationRequest {
pub checkpoint: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct MaintenanceReport {
pub indexes_swept: u64,
pub index_entries_reclaimed: u64,
}
#[derive(Debug, Default)]
struct ReadTally {
materializations: AtomicU64,
refreshes: AtomicU64,
cache_hits: AtomicU64,
head_reads: AtomicU64,
materialize_micros: AtomicU64,
current_scans: AtomicU64,
history_scans: AtomicU64,
}
impl ReadTally {
fn materialized(&self, elapsed: Duration) {
let count = self.materializations.fetch_add(1, Ordering::Relaxed) + 1;
let micros = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX);
let total = self.materialize_micros.fetch_add(micros, Ordering::Relaxed) + micros;
info!(
elapsed_ms = elapsed.as_secs_f64() * 1_000.0,
materializations = count,
total_micros = total,
cache_hits = self.cache_hits.load(Ordering::Relaxed),
refreshes = self.refreshes.load(Ordering::Relaxed),
head_reads = self.head_reads.load(Ordering::Relaxed),
"materialized the catalog from `current`"
);
}
}
struct StoreLocation {
path: String,
object_store: Arc<dyn ObjectStore>,
}
enum Store {
Writer(Db),
Reader(Arc<DbReader>),
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct CatalogOptions {
pub path: String,
pub encrypted: bool,
pub flush_interval: Duration,
pub cache_dir: Option<std::path::PathBuf>,
pub cache_size: Option<u64>,
pub cache_memory: Option<u64>,
pub cache_preload: Option<CachePreload>,
pub cache_puts: bool,
pub data_path: Option<String>,
pub reader_poll_interval: Duration,
pub checkpoint: Option<String>,
}
impl Default for CatalogOptions {
fn default() -> Self {
Self {
path: String::new(),
encrypted: false,
flush_interval: Duration::from_millis(100),
cache_dir: None,
cache_size: None,
cache_memory: None,
cache_preload: None,
cache_puts: false,
data_path: None,
reader_poll_interval: Duration::from_secs(10),
checkpoint: None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CachePreload {
L0,
All,
}
async fn warn_if_preload_cannot_fit(options: &CatalogOptions, object_store: Arc<dyn ObjectStore>) {
if options.cache_preload != Some(CachePreload::All) || options.cache_dir.is_none() {
return;
}
let Ok(store_bytes) = store_census::manifest_bytes(&options.path, object_store).await else {
return;
};
if let Some(shortfall) = open::preload_shortfall(store_bytes, options.cache_size) {
warn!(
path = options.path,
store_bytes,
cache_size = options.cache_size,
shortfall,
"preload cannot hold this store: the load stops once the cache is full, leaving \
the rest to be fetched on demand"
);
}
}
fn parse_checkpoint(checkpoint: Option<&str>) -> Result<Option<uuid::Uuid>> {
checkpoint
.map(|id| {
uuid::Uuid::parse_str(id).map_err(|err| {
Error::Configuration(format!("checkpoint `{id}` is not a valid id: {err}"))
})
})
.transpose()
}
pub type CommitMember<'a> = &'a (dyn Fn(&mut Transaction) -> Result<()> + Sync);
#[derive(Clone)]
pub struct ReadOnlyCatalog {
store: Arc<Store>,
reads: Arc<ReadTally>,
location: Arc<StoreLocation>,
projections: Arc<std::sync::RwLock<ProjectionCache>>,
commits: Arc<commit::Coalescer>,
writer_status: Option<watch::Receiver<DbStatus>>,
}
impl std::fmt::Debug for ReadOnlyCatalog {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ReadOnlyCatalog").finish_non_exhaustive()
}
}
#[derive(Clone)]
pub struct Catalog {
inner: ReadOnlyCatalog,
}
impl std::fmt::Debug for Catalog {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Catalog").finish_non_exhaustive()
}
}
impl std::ops::Deref for Catalog {
type Target = ReadOnlyCatalog;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl ReadOnlyCatalog {
pub(crate) fn projections(&self) -> &Arc<std::sync::RwLock<ProjectionCache>> {
&self.projections
}
pub(crate) fn tally_current_scan(&self) {
self.reads.current_scans.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn tally_history_scan(&self) {
self.reads.history_scans.fetch_add(1, Ordering::Relaxed);
}
#[cfg(test)]
pub(crate) fn entity_scan_tallies(&self) -> (u64, u64) {
(
self.reads.current_scans.load(Ordering::Relaxed),
self.reads.history_scans.load(Ordering::Relaxed),
)
}
pub(crate) fn note_head_read(&self) {
self.reads.head_reads.fetch_add(1, Ordering::Relaxed);
}
#[cfg(test)]
pub(crate) fn head_reads(&self) -> u64 {
self.reads.head_reads.load(Ordering::Relaxed)
}
pub(crate) fn refuse_if_closed(&self) -> Result<()> {
let Some(status) = &self.writer_status else {
return Ok(());
};
match status.borrow().close_reason {
None => Ok(()),
Some(CloseReason::Fenced) => Err(crate::error::fenced()),
Some(reason) => Err(Error::Interrupted(format!(
"this catalog's store handle has closed ({reason:?}); re-attach to use it"
))),
}
}
pub(crate) fn writer_head_view(&self) -> Option<Arc<CatalogSnapshot>> {
if !self.holds_the_writer() {
return None;
}
held_head_view(&self.projections)
}
pub(crate) fn holds_the_writer(&self) -> bool {
matches!(self.store.as_ref(), Store::Writer(_))
}
pub async fn snapshot(&self) -> Result<Arc<CatalogSnapshot>> {
self.view(None).await
}
pub async fn snapshot_at(&self, snapshot: SnapshotId) -> Result<Arc<CatalogSnapshot>> {
self.view(Some(snapshot.get())).await
}
pub async fn recent_rows(&self, table: TableId) -> Result<Vec<RecentRow>> {
let session = self.begin_read().await?;
let outcome = self.scan_recent_rows(&session, table, None).await;
session.finish();
outcome
}
pub async fn recent_rows_at(
&self,
table: TableId,
snapshot: SnapshotId,
) -> Result<Vec<RecentRow>> {
let session = self.begin_read().await?;
let outcome = self
.scan_recent_rows(&session, table, Some(snapshot.get()))
.await;
session.finish();
outcome
}
pub async fn recent_row(&self, table: TableId, row_id: u64) -> Result<Option<RecentRow>> {
Ok(self
.recent_rows(table)
.await?
.into_iter()
.find(|row| row.row_id == row_id))
}
async fn scan_recent_rows(
&self,
session: &ReadSession,
table: TableId,
at: Option<u64>,
) -> Result<Vec<RecentRow>> {
let handle = session.handle();
commit::refuse_mid_migration(handle).await?;
let read_at = match at {
Some(_) => commit::resolve_read_snapshot(handle, at).await?.0,
None => commit::read_head_id(handle).await?,
};
let chunks = store_inline::scan_inline_chunks(handle, table.get()).await?;
let tombstones = store_inline::scan_inline_inline_deletes(handle, table.get()).await?;
let live = InlineScanKind::Table.select(
&materialize_inline_rows(&chunks, &tombstones),
read_at,
0,
);
let mut bodies: HashMap<usize, Arc<Vec<u8>>> = HashMap::new();
let mut schemas: HashMap<u64, Arc<Vec<u8>>> = HashMap::new();
let mut rows = Vec::with_capacity(live.len());
for row in live {
let (operation, chunk) = &chunks[row.chunk];
let InlineOperation::Insert { schema_version, .. } = operation else {
return Err(Error::Corruption(format!(
"inline row {} of table {table} references a non-insert chunk",
row.row_id
)));
};
let arrow_schema = if let Some(schema) = schemas.get(schema_version) {
Arc::clone(schema)
} else {
let schema = store_inline::read_inline_schema(handle, table.get(), *schema_version)
.await?
.ok_or_else(|| {
Error::Corruption(format!(
"no inline schema for table {table} version {schema_version}"
))
})?;
let schema = Arc::new(schema.arrow_schema);
schemas.insert(*schema_version, Arc::clone(&schema));
schema
};
let chunk_body = Arc::clone(
bodies
.entry(row.chunk)
.or_insert_with(|| Arc::new(chunk.body.clone())),
);
rows.push(RecentRow {
row_id: row.row_id,
begin_snapshot: SnapshotId::new(row.begin_snapshot),
schema_version: *schema_version,
offset_in_chunk: row.offset_in_chunk,
chunk_body,
arrow_schema,
});
}
Ok(rows)
}
async fn view(&self, at: Option<u64>) -> Result<Arc<CatalogSnapshot>> {
if at.is_some() {
let session = self.begin_read().await?;
let started = Instant::now();
let view = commit::materialize(session.handle(), at).await;
session.finish();
if view.is_ok() {
self.reads.materialized(started.elapsed());
}
return view.map(Arc::new);
}
if let Some(view) = self.writer_head_view() {
self.refuse_if_closed()?;
self.reads.cache_hits.fetch_add(1, Ordering::Relaxed);
return Ok(view);
}
let epoch = cache_epoch(&self.projections);
let session = self.begin_read().await?;
let view = self.head_view(session.handle()).await;
session.finish();
let view = view?;
install_head_view_at(&self.projections, epoch, Arc::clone(&view));
Ok(view)
}
async fn head_view(&self, handle: ReadHandle<'_>) -> Result<Arc<CatalogSnapshot>> {
self.note_head_read();
let head = commit::read_head_value(handle).await?;
if let Some(cached) = cached_head_view(&self.projections, &head) {
self.reads.cache_hits.fetch_add(1, Ordering::Relaxed);
return Ok(cached);
}
if let Some(behind) = held_head_view(&self.projections)
&& let Some(refreshed) = commit::refresh(handle, &behind).await?
{
self.reads.refreshes.fetch_add(1, Ordering::Relaxed);
return Ok(Arc::new(refreshed));
}
if let Some(current) = shared_current_entities(&self.projections, &head)
&& let Some(view) = commit::materialize_from(handle, &head, ¤t).await?
{
self.reads.cache_hits.fetch_add(1, Ordering::Relaxed);
return Ok(Arc::new(view));
}
let started = Instant::now();
let (view, records) = commit::materialize_capturing(handle).await?;
self.reads.materialized(started.elapsed());
self.tally_current_scan();
install_shared_current_entities(
&self.projections,
HeadValue {
snapshot_id: view.snapshot.snapshot_id,
batch_seq: view.batch_seq,
},
records,
);
Ok(Arc::new(view))
}
pub async fn index_lookup(
&self,
table: TableId,
index: IndexId,
values: &[IndexKeyValue],
) -> Result<Vec<RowLocation>> {
let session = self.begin_read().await?;
let handle = session.handle();
let outcome = async {
let view = self.head_view(handle).await?;
let info = view
.index_by_id(table, index)
.ok_or_else(|| Error::NotFound(format!("index {index} on table {table}")))?;
match info.state {
IndexState::Ready => {}
IndexState::Building => {
return Err(Error::IndexBuilding(format!(
"index {index} is still building"
)));
}
IndexState::Poisoned => {
return Err(Error::NotFound(format!("index {index} was poisoned")));
}
}
let key = encode_ordered_values(
&values.iter().cloned().map(Some).collect::<Vec<_>>(),
&info.directions,
&info.nulls,
)?;
let row_ids =
index_maintenance::lookup_row_ids(handle, index.get(), info.unique, &key).await?;
let holders = RowHolders::of(&view.data_files_of(table));
Ok(row_ids
.into_iter()
.map(|row_id| RowLocation {
row_id,
holder: holders.holder(row_id),
})
.collect())
}
.await;
session.finish();
outcome
}
pub async fn index_range(
&self,
table: TableId,
index: IndexId,
lower: Bound<Vec<IndexKeyValue>>,
upper: Bound<Vec<IndexKeyValue>>,
reverse: bool,
) -> Result<Vec<RowLocation>> {
let session = self.begin_read().await?;
let handle = session.handle();
let outcome = async {
let view = self.head_view(handle).await?;
let info = view
.index_by_id(table, index)
.ok_or_else(|| Error::NotFound(format!("index {index} on table {table}")))?;
match info.state {
IndexState::Ready => {}
IndexState::Building => {
return Err(Error::IndexBuilding(format!(
"index {index} is still building"
)));
}
IndexState::Poisoned => {
return Err(Error::NotFound(format!("index {index} was poisoned")));
}
}
let (byte_lower, byte_upper) = encode_range_bounds(&info, index, lower, upper)?;
let leading_nulls = info.nulls.first().copied().unwrap_or(NullOrder::Last);
let mut row_ids = index_maintenance::range_row_ids(
handle,
index.get(),
info.unique,
leading_nulls,
byte_lower,
byte_upper,
)
.await?;
if reverse {
row_ids.reverse();
}
let holders = RowHolders::of(&view.data_files_of(table));
Ok(row_ids
.into_iter()
.map(|row_id| RowLocation {
row_id,
holder: holders.holder(row_id),
})
.collect())
}
.await;
session.finish();
outcome
}
pub async fn index_nulls(
&self,
table: TableId,
index: IndexId,
prefix: Vec<Option<IndexKeyValue>>,
reverse: bool,
) -> Result<Vec<RowLocation>> {
let session = self.begin_read().await?;
let handle = session.handle();
let outcome = async {
let view = self.head_view(handle).await?;
let info = view
.index_by_id(table, index)
.ok_or_else(|| Error::NotFound(format!("index {index} on table {table}")))?;
match info.state {
IndexState::Ready => {}
IndexState::Building => {
return Err(Error::IndexBuilding(format!(
"index {index} is still building"
)));
}
IndexState::Poisoned => {
return Err(Error::NotFound(format!("index {index} was poisoned")));
}
}
if prefix.is_empty() || prefix.len() > info.columns.len() {
return Err(Error::Constraint(format!(
"index_nulls: a prefix of {} predicates does not fit the {}-column index \
{index}",
prefix.len(),
info.columns.len()
)));
}
if prefix.iter().all(Option::is_some) {
return Err(Error::Constraint(
"index_nulls: the prefix names no IS NULL; use index_lookup for pure equality"
.to_owned(),
));
}
let key = encode_ordered_values(&prefix, &info.directions, &info.nulls)?;
let mut row_ids =
index_maintenance::null_prefix_row_ids(handle, index.get(), &key).await?;
if reverse {
row_ids.reverse();
}
let holders = RowHolders::of(&view.data_files_of(table));
Ok(row_ids
.into_iter()
.map(|row_id| RowLocation {
row_id,
holder: holders.holder(row_id),
})
.collect())
}
.await;
session.finish();
outcome
}
pub(crate) async fn begin_read(&self) -> Result<ReadSession> {
let session = match self.store.as_ref() {
Store::Writer(db) => ReadSession::Tx(
db.begin(IsolationLevel::Snapshot)
.await
.map_err(Error::from)?,
),
Store::Reader(reader) => ReadSession::Reader(reader.clone()),
};
if let Err(error) = commit::refuse_mid_migration(session.handle()).await {
session.finish();
return Err(error);
}
Ok(session)
}
pub async fn scoped_file_index_entries(
&self,
object_store: Arc<dyn ObjectStore>,
path: &Path,
index: IndexId,
indexed_positions: &[usize],
) -> Result<Vec<FileIndexEntry>> {
let entries = scoped_read::scoped_read_entries(
object_store,
path,
indexed_positions,
scoped_read::RowIdSource::Ordinal,
None,
)
.await?;
Ok(entries
.into_iter()
.map(|entry| FileIndexEntry {
index,
ordinal: entry.row_id,
values: entry.values,
})
.collect())
}
pub async fn scoped_backfill_entries(
&self,
object_store: Arc<dyn ObjectStore>,
data_prefix: &str,
table: TableId,
columns: &[ColumnId],
) -> Result<Vec<IndexEntry>> {
let session = self.begin_read().await?;
let outcome = async {
let snapshot = commit::materialize(session.handle(), None).await?;
let live_columns = snapshot.columns_of(table);
let positions = columns
.iter()
.map(|column| {
live_columns
.iter()
.position(|c| c.id == *column)
.ok_or_else(|| Error::NotFound(format!("column {column} of table {table}")))
})
.collect::<Result<Vec<_>>>()?;
let table_prefix = snapshot.table_data_prefix(table)?;
let resolve = |path: &str, is_relative: bool| {
let relative = match (is_relative, data_prefix.is_empty()) {
(false, _) => path.to_owned(),
(true, true) => format!("{table_prefix}{path}"),
(true, false) => format!("{data_prefix}/{table_prefix}{path}"),
};
object_store::path::Path::from(relative.as_str())
};
let mut killed_positions: HashMap<u64, HashSet<u64>> = HashMap::new();
let mut killed_row_ids: HashMap<u64, HashSet<u64>> = HashMap::new();
for (data_file_id, row_id, _) in
store_inline::scan_inline_file_deletes(session.handle(), table.get()).await?
{
killed_row_ids
.entry(data_file_id)
.or_default()
.insert(row_id);
}
for delete in snapshot.delete_files_of(table) {
let path = resolve(&delete.path, delete.path_is_relative);
let positions =
scoped_read::delete_file_positions(object_store.as_ref(), &path).await?;
killed_positions
.entry(delete.data_file_id.get())
.or_default()
.extend(positions);
}
let mut entries = Vec::new();
for file in snapshot.data_files_of(table) {
let path = resolve(&file.path, file.path_is_relative);
let scoped = scoped_read::scoped_read_entries(
Arc::clone(&object_store),
&path,
&positions,
scoped_read::RowIdSource::Resolve {
row_id_start: file.row_id_start,
},
Some(file.file_size_bytes),
)
.await?;
let dead_positions = killed_positions.get(&file.id.get());
let dead_row_ids = killed_row_ids.get(&file.id.get());
entries.extend(
scoped
.into_iter()
.enumerate()
.filter_map(|(ordinal, entry)| {
let ordinal = u64::try_from(ordinal).unwrap_or(u64::MAX);
let dead = dead_positions.is_some_and(|dead| dead.contains(&ordinal))
|| dead_row_ids.is_some_and(|dead| dead.contains(&entry.row_id));
(!dead).then_some(IndexEntry {
row_id: entry.row_id,
values: entry.values,
})
}),
);
}
Ok(entries)
}
.await;
session.finish();
outcome
}
async fn staged_build_cursor(&self, table: TableId, index: IndexId) -> Result<Option<u64>> {
let info = self
.snapshot()
.await?
.indexes_of(table)
.into_iter()
.find(|info| info.id == index)
.ok_or_else(|| Error::NotFound(format!("index {index}")))?;
match info.state {
IndexState::Building => Ok(info.build_cursor),
IndexState::Ready => Err(Error::Constraint(format!(
"index {index} finished building under this build"
))),
IndexState::Poisoned => Err(Error::Constraint(format!(
"index {index} was poisoned by a duplicate value"
))),
}
}
pub async fn inline_backfill_entries(
&self,
table: TableId,
columns: &[ColumnId],
) -> Result<Vec<IndexEntry>> {
let session = self.begin_read().await?;
let outcome = async {
let snapshot = commit::materialize(session.handle(), None).await?;
let live_columns = snapshot.columns_of(table);
let positions = columns
.iter()
.map(|column| {
live_columns
.iter()
.position(|c| c.id == *column)
.ok_or_else(|| Error::NotFound(format!("column {column} of table {table}")))
})
.collect::<Result<Vec<_>>>()?;
let dead: std::collections::HashSet<u64> =
store_inline::scan_inline_inline_deletes(session.handle(), table.get())
.await?
.into_iter()
.map(|(row_id, _)| row_id)
.collect();
let mut entries = Vec::new();
for (op, chunk) in
store_inline::scan_inline_chunks(session.handle(), table.get()).await?
{
let InlineOperation::Insert { schema_version, .. } = op else {
continue;
};
let schema =
store_inline::read_inline_schema(session.handle(), table.get(), schema_version)
.await?
.ok_or_else(|| {
Error::Corruption(format!(
"no inline schema for table {table} version {schema_version}"
))
})?;
let scoped = scoped_read::inline_batch_entries(
&schema.arrow_schema,
&chunk.body,
&positions,
chunk.row_id_start,
)?;
entries.extend(
scoped
.into_iter()
.filter(|entry| !dead.contains(&entry.row_id))
.map(|entry| IndexEntry {
row_id: entry.row_id,
values: entry.values,
}),
);
}
Ok(entries)
}
.await;
session.finish();
outcome
}
pub async fn store_census(&self, request: CensusRequest) -> Result<StoreCensus> {
let physical = store_census::read_manifest_census(
&self.location.path,
Arc::clone(&self.location.object_store),
)
.await?;
let mut subspaces: Vec<SubspaceCensus> = Subspace::ALL
.into_iter()
.map(|subspace| {
let prefix = subspace_prefix(subspace);
measure(SubspaceName::from(subspace), physical.segment(&prefix))
})
.collect();
subspaces.extend(
physical
.segments
.iter()
.filter(|segment| {
matches!(
SubspaceName::of_prefix(&segment.prefix),
SubspaceName::Unknown(_)
)
})
.map(|segment| measure(SubspaceName::of_prefix(&segment.prefix), Some(segment))),
);
if request.count_live_entries {
let session = self.begin_read().await?;
let counted = count_live_entries(session.handle(), &mut subspaces).await;
session.finish();
counted?;
}
Ok(StoreCensus {
manifest_id: physical.manifest_id,
subspaces,
objects: physical.objects.map(|totals| StoreObjects {
total_objects: totals.total_objects,
total_bytes: totals.total_bytes,
wal_objects: totals.wal_objects,
wal_bytes: totals.wal_bytes,
manifest_objects: totals.manifest_objects,
manifest_bytes: totals.manifest_bytes,
sst_objects: totals.sst_objects,
sst_bytes: totals.sst_bytes,
other_objects: totals.other_objects,
other_bytes: totals.other_bytes,
}),
})
}
pub(crate) async fn first_index_id_from(
&self,
kind: IndexKind,
from: u64,
) -> Result<Option<u64>> {
let kind_prefix = index_kind_prefix(kind);
let start = index_index_prefix(kind, from);
let suffix = start[kind_prefix.len()..].to_vec();
let session = self.begin_read().await?;
let first = session
.handle()
.scan_prefix(kind_prefix, suffix.., ScanShape::Probe)
.await
.map_err(Error::from)?
.next()
.await
.map_err(Error::from)?;
session.finish();
let Some(entry) = first else {
return Ok(None);
};
match Key::decode(&entry.key)? {
Key::Index(IndexKey::Unique { index_id, .. } | IndexKey::Multi { index_id, .. }) => {
Ok(Some(index_id))
}
other => Err(Error::Corruption(format!(
"key in the index subspace decoded as {other:?}"
))),
}
}
pub async fn close(&self) -> Result<()> {
match self.store.as_ref() {
Store::Writer(db) => db.close().await.map_err(Error::from),
Store::Reader(reader) => reader.close().await.map_err(Error::from),
}
}
}
fn encode_range_bounds(
info: &IndexInfo,
index: IndexId,
lower: Bound<Vec<IndexKeyValue>>,
upper: Bound<Vec<IndexKeyValue>>,
) -> Result<(Bound<CanonicalKey>, Bound<CanonicalKey>)> {
let named_len = |bound: &Bound<Vec<IndexKeyValue>>| match bound {
Bound::Included(values) | Bound::Excluded(values) => Some(values.len()),
Bound::Unbounded => None,
};
let (lower_len, upper_len) = (named_len(&lower), named_len(&upper));
if lower_len == Some(0) || upper_len == Some(0) {
return Err(Error::Constraint(format!(
"index_range: a bound of index {index} must name at least one column"
)));
}
let widest = lower_len.unwrap_or(0).max(upper_len.unwrap_or(0));
if widest > info.columns.len() {
return Err(Error::Constraint(format!(
"index_range: a bound of {widest} values does not fit the {}-column index {index}",
info.columns.len()
)));
}
let range_column = widest.saturating_sub(1);
let mixed_directions = info
.directions
.iter()
.take(widest)
.any(|direction| Some(direction) != info.directions.first());
let pinned_prefix = match (&lower, &upper) {
(
Bound::Included(low) | Bound::Excluded(low),
Bound::Included(high) | Bound::Excluded(high),
) => low.len() == high.len() && low[..range_column] == high[..range_column],
_ => false,
};
if mixed_directions && !pinned_prefix {
return Err(Error::Constraint(format!(
"index_range: index {index} names columns of differing sort directions, so both \
bounds must pin the same leading values and compare on the last column"
)));
}
let encode_bound = |bound: Bound<Vec<IndexKeyValue>>| -> Result<Bound<_>> {
let encode = |values: Vec<IndexKeyValue>| {
encode_ordered_values(
&values.into_iter().map(Some).collect::<Vec<_>>(),
&info.directions,
&info.nulls,
)
};
Ok(match bound {
Bound::Included(values) => Bound::Included(encode(values)?),
Bound::Excluded(values) => Bound::Excluded(encode(values)?),
Bound::Unbounded => Bound::Unbounded,
})
};
let descending = info.directions.get(range_column).copied() == Some(Direction::Descending);
let (byte_lower, byte_upper) = if descending {
(upper, lower)
} else {
(lower, upper)
};
Ok((encode_bound(byte_lower)?, encode_bound(byte_upper)?))
}
struct RowHolders {
ranges: Vec<(u64, u64, DataFileId)>,
}
impl RowHolders {
fn of(files: &[DataFileInfo]) -> Self {
let mut ranges: Vec<_> = files
.iter()
.filter_map(|file| {
let start = file.row_id_start?;
Some((start, start.saturating_add(file.record_count), file.id))
})
.collect();
ranges.sort_unstable();
Self { ranges }
}
fn holder(&self, row_id: u64) -> RowHolder {
let above = self.ranges.partition_point(|(start, ..)| *start <= row_id);
match self.ranges[..above].last() {
Some(&(_, end, file)) if row_id < end => RowHolder::DataFile(file),
_ => RowHolder::Inline,
}
}
}
impl Catalog {
fn writer(&self) -> Result<&Db> {
match self.store.as_ref() {
Store::Writer(db) => Ok(db),
Store::Reader(_) => Err(Error::Constraint(
"catalog opened read-only; writes are unavailable".to_string(),
)),
}
}
pub async fn open(object_store: Arc<dyn ObjectStore>, options: CatalogOptions) -> Result<Self> {
if options.checkpoint.is_some() {
return Err(Error::Configuration(
"a checkpoint pins a read-only catalog to a fixed cut; a writer commits new \
state and cannot be opened against one"
.to_string(),
));
}
warn_if_preload_cannot_fit(&options, Arc::clone(&object_store)).await;
let located = Arc::clone(&object_store);
let store = StoreBuilder::new(&options.path, object_store)
.flush_interval(options.flush_interval)
.cache_dir(options.cache_dir.clone())
.cache_size(options.cache_size)
.cache_memory(options.cache_memory)
.cache_preload(options.cache_preload)
.cache_puts(options.cache_puts);
let db = commit::open_initialized(store, options.encrypted, options.data_path.as_deref())
.await?;
info!(
path = options.path,
flush_interval_ms = options.flush_interval.as_millis(),
"opened catalog read-write"
);
let projections = Arc::new(std::sync::RwLock::new(ProjectionCache::empty()));
Ok(Self {
inner: ReadOnlyCatalog {
writer_status: Some(db.subscribe()),
store: Arc::new(Store::Writer(db)),
location: Arc::new(StoreLocation {
path: options.path,
object_store: located,
}),
reads: Arc::new(ReadTally::default()),
commits: Arc::new(commit::Coalescer::new(Arc::clone(&projections))),
projections,
},
})
}
pub async fn open_read_only(
object_store: Arc<dyn ObjectStore>,
options: CatalogOptions,
) -> Result<ReadOnlyCatalog> {
let checkpoint = parse_checkpoint(options.checkpoint.as_deref())?;
warn_if_preload_cannot_fit(&options, Arc::clone(&object_store)).await;
let located = Arc::clone(&object_store);
let store = StoreBuilder::new(&options.path, object_store)
.cache_dir(options.cache_dir.clone())
.cache_size(options.cache_size)
.cache_memory(options.cache_memory)
.cache_preload(options.cache_preload)
.cache_puts(options.cache_puts)
.poll_interval(options.reader_poll_interval)
.checkpoint(checkpoint);
let reader = commit::open_reader_initialized(store).await?;
info!(
path = options.path,
checkpoint = options.checkpoint,
"opened catalog read-only"
);
let projections = Arc::new(std::sync::RwLock::new(ProjectionCache::empty()));
Ok(ReadOnlyCatalog {
writer_status: None,
store: Arc::new(Store::Reader(Arc::new(reader))),
location: Arc::new(StoreLocation {
path: options.path,
object_store: located,
}),
reads: Arc::new(ReadTally::default()),
commits: Arc::new(commit::Coalescer::new(Arc::clone(&projections))),
projections,
})
}
pub async fn migrate(
object_store: Arc<dyn ObjectStore>,
options: CatalogOptions,
request: MigrationRequest,
) -> Result<MigrationReport> {
let db = StoreBuilder::new(&options.path, object_store.clone())
.flush_interval(options.flush_interval)
.cache_dir(options.cache_dir.clone())
.cache_size(options.cache_size)
.cache_memory(options.cache_memory)
.cache_preload(options.cache_preload)
.cache_puts(options.cache_puts)
.open_writer()
.await?;
let checkpoint = if request.checkpoint {
let taken = open::create_checkpoint(&db, None).await?;
info!(checkpoint = %taken, "took a pre-migration checkpoint");
Some(taken)
} else {
None
};
let report = migration::run(&db).await;
let closed = db.close().await.map_err(Error::from);
let report = report.and_then(|report| closed.map(|()| report))?;
if let Some(checkpoint) = checkpoint {
StoreBuilder::new(&options.path, object_store)
.delete_checkpoint(checkpoint)
.await?;
}
Ok(report)
}
pub async fn create_checkpoint(&self, lifetime: Option<Duration>) -> Result<String> {
let id = open::create_checkpoint(self.writer()?, lifetime).await?;
info!(checkpoint = %id, "created a checkpoint");
Ok(id.to_string())
}
pub async fn delete_checkpoint(
object_store: Arc<dyn ObjectStore>,
options: CatalogOptions,
checkpoint: &str,
) -> Result<()> {
let id = parse_checkpoint(Some(checkpoint))?
.ok_or_else(|| Error::Configuration("no checkpoint given".to_string()))?;
StoreBuilder::new(&options.path, object_store)
.delete_checkpoint(id)
.await
}
pub async fn checkpoints(
object_store: Arc<dyn ObjectStore>,
options: CatalogOptions,
) -> Result<Vec<String>> {
let ids = StoreBuilder::new(&options.path, object_store)
.list_checkpoints()
.await?;
Ok(ids.iter().map(uuid::Uuid::to_string).collect())
}
pub(crate) async fn begin_write_tx(&self) -> Result<DbTransaction> {
self.writer()?
.begin(IsolationLevel::Snapshot)
.await
.map_err(Error::from)
}
pub async fn create_index_staged(
&self,
table: TableId,
def: &IndexDef,
orders: &[ColumnOrder],
data_store: Option<Arc<dyn ObjectStore>>,
data_prefix: &str,
step_entries: Option<usize>,
) -> Result<IndexId> {
let step_entries = step_entries.unwrap_or(BUILD_STEP_ENTRIES);
if step_entries == 0 {
return Err(Error::Constraint(
"a staged build's step size must be at least one entry".to_owned(),
));
}
let index = self.begin_staged_index(table, def, orders).await?;
let outcome = self
.drive_staged_build(table, def, index, data_store, data_prefix, step_entries)
.await;
if outcome.is_err()
&& let Err(cleanup) = self.commit(|tx| tx.drop_index(index)).await
{
warn!(
index = index.get(),
error = %cleanup,
"could not drop the definition of a failed staged build"
);
}
outcome.map(|()| index)
}
async fn commit_build_steps(
&self,
table: TableId,
index: IndexId,
entries: &[IndexEntry],
step_entries: usize,
) -> Result<BuildProgress> {
loop {
let cursor = self.staged_build_cursor(table, index).await?;
let pending = match cursor {
Some(covered) => entries.partition_point(|entry| entry.row_id <= covered),
None => 0,
};
let remaining = &entries[pending..];
let step = &remaining[..remaining.len().min(step_entries)];
let is_final = step.len() == remaining.len();
match self
.commit(|tx| tx.build_index_step(index, step, is_final).map(|_| ()))
.await
{
Ok(_) => {
if is_final {
return Ok(BuildProgress::Ready);
}
}
Err(Error::CommitConflict(_)) => return Ok(BuildProgress::Conflicted),
Err(other) => return Err(other),
}
}
}
pub async fn reclaim_index_entries(&self, index: IndexId, limit: usize) -> Result<usize> {
let head = self.snapshot().await?;
if head
.indexes
.values()
.any(|per_table| per_table.contains_key(&index.get()))
{
return Err(Error::Constraint(format!(
"index {index} is still live; drop it before reclaiming its entries"
)));
}
let tx = self.begin_write_tx().await?;
let deleted = index_maintenance::reclaim_entries(&tx, index.get(), limit).await?;
commit::commit_durable(tx, "entry reclamation")
.await
.map_err(Error::from)?;
Ok(deleted)
}
pub async fn maintain(&self, request: MaintenanceRequest) -> Result<MaintenanceReport> {
self.writer()?;
if request.batch_size == 0 {
return Err(Error::Configuration(
"batch_size must be nonzero; zero would reclaim nothing and never terminate"
.to_string(),
));
}
let mut report = MaintenanceReport::default();
if !request.sweep_orphaned_index_entries {
return Ok(report);
}
let live: HashSet<u64> = self
.snapshot()
.await?
.indexes
.values()
.flat_map(|per_table| per_table.keys().copied())
.collect();
for kind in [IndexKind::Unique, IndexKind::Multi] {
let mut from = 0u64;
while let Some(index_id) = self.first_index_id_from(kind, from).await? {
if !live.contains(&index_id) {
let reclaimed = self
.reclaim_dead_range(kind, index_id, request.batch_size)
.await?;
if reclaimed > 0 {
report.indexes_swept += 1;
report.index_entries_reclaimed += reclaimed;
}
}
match index_id.checked_add(1) {
Some(next) => from = next,
None => break,
}
}
}
Ok(report)
}
pub async fn compact_store(&self, request: CompactStoreRequest) -> Result<CompactStoreReport> {
self.writer()?;
let target = match &request.target {
CompactionTarget::WholeStore => None,
CompactionTarget::Subspace(name) => match name.subspace() {
Some(subspace) => Some(subspace_prefix(subspace)),
None => {
return Err(Error::Configuration(format!(
"{name} is not a subspace this build can merge"
)));
}
},
};
let before = self.store_census(CensusRequest::default()).await?;
let submitted = store_compaction::submit_full_merge(
&self.location.path,
Arc::clone(&self.location.object_store),
target.as_deref(),
)
.await?;
let mut merges = Vec::new();
for merge in &submitted {
let subspace = SubspaceName::of_prefix(&merge.segment);
let outcome = match request.wait {
None => MergeOutcome::Pending,
Some(budget) => match store_compaction::await_merge(
&self.location.path,
Arc::clone(&self.location.object_store),
&merge.compaction,
budget,
)
.await?
{
MergeEnd::Completed => MergeOutcome::Completed,
MergeEnd::Failed => {
MergeOutcome::Failed("the merge ended without committing".to_string())
}
MergeEnd::Pending => MergeOutcome::Pending,
},
};
merges.push(SubspaceMerge {
subspace: subspace.clone(),
outcome,
bytes_before: bytes_of(&before, &subspace),
bytes_after: None,
});
}
for measured in &before.subspaces {
let covered = target
.as_ref()
.is_none_or(|prefix| SubspaceName::of_prefix(prefix) == measured.subspace);
if !covered || merges.iter().any(|m| m.subspace == measured.subspace) {
continue;
}
merges.push(SubspaceMerge {
subspace: measured.subspace.clone(),
outcome: MergeOutcome::Skipped("no sorted runs to merge"),
bytes_before: measured.bytes,
bytes_after: None,
});
}
if merges
.iter()
.any(|merge| merge.outcome == MergeOutcome::Completed)
{
let after = self.store_census(CensusRequest::default()).await?;
for merge in &mut merges {
if merge.outcome == MergeOutcome::Completed {
merge.bytes_after = Some(bytes_of(&after, &merge.subspace));
}
}
}
Ok(CompactStoreReport { merges })
}
pub async fn commit<F>(&self, f: F) -> Result<SnapshotId>
where
F: Fn(&mut Transaction) -> Result<()>,
{
let ids =
commit::commit_cycle(self.writer()?, std::slice::from_ref(&f), &self.commits).await?;
ids.first().copied().ok_or_else(|| {
Error::Corruption("a commit of one member reported no snapshot".to_string())
})
}
pub async fn commit_group(&self, members: &[CommitMember<'_>]) -> Result<Vec<SnapshotId>> {
commit::commit_cycle(self.writer()?, members, &self.commits).await
}
async fn begin_staged_index(
&self,
table: TableId,
def: &IndexDef,
orders: &[ColumnOrder],
) -> Result<IndexId> {
if let Some(existing) = self.snapshot().await?.index_by_name(table, &def.name) {
return match existing.state {
IndexState::Ready => Err(Error::AlreadyExists(format!(
"index {} on table {table}",
def.name
))),
IndexState::Building | IndexState::Poisoned => {
let (directions, nulls) = requested_orders(orders, def.columns.len());
if existing.columns != def.columns
|| existing.unique != def.unique
|| existing.directions != directions
|| existing.nulls != nulls
{
return Err(Error::Constraint(format!(
"index {} on table {table} is already building over a different \
definition; drop it to rebuild",
def.name
)));
}
Ok(existing.id)
}
};
}
let index = std::cell::Cell::new(None);
self.commit(|tx| {
let id = tx.create_index_staged_ordered(table, def, orders)?;
index.set(Some(id));
Ok(())
})
.await?;
index
.get()
.ok_or_else(|| Error::Corruption("staged create returned no index id".to_owned()))
}
async fn drive_staged_build(
&self,
table: TableId,
def: &IndexDef,
index: IndexId,
data_store: Option<Arc<dyn ObjectStore>>,
data_prefix: &str,
step_entries: usize,
) -> Result<()> {
for _ in 0..BUILD_DERIVATION_ATTEMPTS {
let mut entries = match &data_store {
Some(store) => {
self.scoped_backfill_entries(
Arc::clone(store),
data_prefix,
table,
&def.columns,
)
.await?
}
None => Vec::new(),
};
entries.extend(self.inline_backfill_entries(table, &def.columns).await?);
entries.sort_unstable_by_key(|entry| entry.row_id);
if let BuildProgress::Ready = self
.commit_build_steps(table, index, &entries, step_entries)
.await?
{
return Ok(());
}
}
Err(Error::CommitConflict(format!(
"staged build of index {index} lost its race {BUILD_DERIVATION_ATTEMPTS} times; \
the table is under concurrent write"
)))
}
async fn reclaim_dead_range(
&self,
kind: IndexKind,
index_id: u64,
batch_size: usize,
) -> Result<u64> {
let mut total = 0u64;
let mut cursor: Option<Vec<u8>> = None;
loop {
let tx = self.begin_write_tx().await?;
let (deleted, last) = index_maintenance::reclaim_entries_from(
&tx,
kind,
index_id,
batch_size,
cursor.as_deref(),
)
.await?;
if deleted == 0 {
tx.rollback();
return Ok(total);
}
tx.commit_with_options(&commit::non_durable())
.await
.map_err(Error::from)?;
total += deleted as u64;
cursor = last;
}
}
}