use futures::{StreamExt, stream};
use tracing::debug;
use super::ReadOnlyCatalog;
use crate::{
catalog::{CatalogSnapshot, TableId},
error::Result,
store::{
handle::{ReadHandle, ScanShape},
key::{
IndexKind, InlineOperationKind, index_index_prefix, inline_chunk_range_table_prefix,
inline_live_table_prefix, inline_schema_table_prefix,
},
},
};
const WARM_RANGE_CONCURRENCY: usize = 32;
fn probe_prefixes(view: &CatalogSnapshot, table: TableId) -> Vec<Vec<u8>> {
let table_id = table.get();
let mut prefixes = view
.indexes_of(table)
.into_iter()
.map(|index| {
let kind = if index.unique {
IndexKind::Unique
} else {
IndexKind::Multi
};
index_index_prefix(kind, index.id.get())
})
.collect::<Vec<_>>();
prefixes.extend([
inline_live_table_prefix(InlineOperationKind::Insert, table_id),
inline_live_table_prefix(InlineOperationKind::InlineDelete, table_id),
inline_live_table_prefix(InlineOperationKind::FileDelete, table_id),
inline_schema_table_prefix(table_id),
inline_chunk_range_table_prefix(table_id),
]);
prefixes
}
async fn warm_prefix(handle: ReadHandle<'_>, prefix: Vec<u8>) -> Result<()> {
let mut iterator = handle.scan_prefix(prefix, .., ScanShape::Probe).await?;
iterator.next().await?;
Ok(())
}
impl ReadOnlyCatalog {
pub async fn warm_tables(&self, tables: &[TableId]) -> Result<()> {
let session = self.begin_read().await?;
let handle = session.handle();
let view = self.head_view(handle).await?;
let prefixes = tables
.iter()
.flat_map(|table| probe_prefixes(&view, *table))
.collect::<Vec<_>>();
let ranges = prefixes.len();
let failed = stream::iter(prefixes)
.map(|prefix| warm_prefix(handle, prefix))
.buffer_unordered(WARM_RANGE_CONCURRENCY)
.filter(|result| std::future::ready(result.is_err()))
.count()
.await;
debug!(
tables = tables.len(),
ranges, failed, "warmed table probe ranges"
);
session.finish();
Ok(())
}
}