use crate::durability::durable_replace_file;
use crate::file_io::{self as fs, File};
use hawdb_core::{HawDBError, PropertyId, Result, Value};
use hawdb_integrity::crc32c;
use std::collections::BTreeMap;
use std::io::Write;
use std::path::Path;
pub const SHADOW_KEY_DICTIONARY_FILE: &str = "property-keys.hawdb";
const SHADOW_KEY_DICTIONARY_MAGIC: &[u8; 9] = b"SKNSHKEY1";
const SHADOW_KEY_DICTIONARY_VERSION: u32 = 1;
const MAX_SHADOW_KEY_DICTIONARY_BYTES: u64 = 64 * 1024 * 1024;
pub const FIRST_DICTIONARY_COLUMN: u32 = 4;
pub const DEFAULT_SHADOW_METADATA_BUDGET_BYTES: u64 = 8 * 1024 * 1024;
const SHADOW_KEY_ENTRY_OVERHEAD_BYTES: u64 = 48;
pub const SHADOW_PASS1_TABLE_OVERHEAD_BYTES: u64 = 64;
const SHADOW_PASS1_PROPERTY_OVERHEAD_BYTES: u64 = SHADOW_KEY_ENTRY_OVERHEAD_BYTES;
const SHADOW_LAYOUT_TABLE_OVERHEAD_BYTES: u64 = 64;
const SHADOW_LAYOUT_PROPERTY_OVERHEAD_BYTES: u64 = 2 * SHADOW_KEY_ENTRY_OVERHEAD_BYTES;
#[derive(Debug, Default)]
pub struct ShadowKeyDictionary {
ids: BTreeMap<String, u32>,
keys: Vec<String>,
}
#[derive(Debug, Clone, Copy)]
pub struct ShadowMetadataBudget {
limit_bytes: u64,
used_bytes: u64,
peak_bytes: u64,
}
impl ShadowMetadataBudget {
pub fn new(limit_bytes: u64, used_bytes: u64) -> Result<Self> {
if used_bytes > limit_bytes {
return Err(Self::exceeded(used_bytes, limit_bytes));
}
Ok(Self {
limit_bytes,
used_bytes,
peak_bytes: used_bytes,
})
}
pub fn charge(&mut self, bytes: u64, allocation: &str) -> Result<()> {
let required = self.used_bytes.checked_add(bytes).ok_or_else(|| {
HawDBError::Storage(format!(
"columnar shadow metadata accounting overflows while reserving {allocation}"
))
})?;
if required > self.limit_bytes {
return Err(HawDBError::Storage(format!(
"columnar shadow {allocation} needs {required} metadata bytes, exceeding its \
enforced {} byte budget",
self.limit_bytes
)));
}
self.used_bytes = required;
self.peak_bytes = self.peak_bytes.max(required);
Ok(())
}
pub fn release(&mut self, bytes: u64) {
debug_assert!(bytes <= self.used_bytes);
self.used_bytes = self.used_bytes.saturating_sub(bytes);
}
pub const fn used_bytes(self) -> u64 {
self.used_bytes
}
pub const fn peak_bytes(self) -> u64 {
self.peak_bytes
}
fn exceeded(required: u64, limit: u64) -> HawDBError {
HawDBError::Storage(format!(
"columnar shadow existing key dictionary needs {required} metadata bytes, exceeding \
its enforced {limit} byte budget"
))
}
}
impl ShadowKeyDictionary {
pub fn load(
shadow_root: &Path,
metadata_budget_bytes: u64,
) -> Result<(Self, ShadowMetadataBudget)> {
let path = shadow_root.join(SHADOW_KEY_DICTIONARY_FILE);
let file_bytes = match fs::metadata(&path) {
Ok(metadata) => metadata.len(),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Ok((
Self::default(),
ShadowMetadataBudget::new(metadata_budget_bytes.saturating_mul(2), 0)?,
));
}
Err(error) => return Err(error.into()),
};
if file_bytes > MAX_SHADOW_KEY_DICTIONARY_BYTES {
return Err(HawDBError::Storage(
"columnar shadow key dictionary exceeds its size limit".to_string(),
));
}
let mut scratch = ShadowMetadataBudget::new(u64::MAX, 0)?;
scratch.charge(file_bytes, "key dictionary read buffer")?;
let bytes = fs::read(&path)?;
let dictionary = Self::decode(&bytes, &mut scratch)?;
scratch.release(file_bytes);
let existing_bytes = scratch.used_bytes();
let metadata_budget = ShadowMetadataBudget::new(
metadata_budget_bytes
.saturating_mul(2)
.saturating_add(existing_bytes.saturating_mul(3)),
existing_bytes,
)?;
Ok((dictionary, metadata_budget))
}
fn decode(bytes: &[u8], metadata_budget: &mut ShadowMetadataBudget) -> Result<Self> {
let corrupt = |message: &str| {
HawDBError::Storage(format!("columnar shadow key dictionary {message}"))
};
if bytes.len() as u64 > MAX_SHADOW_KEY_DICTIONARY_BYTES {
return Err(corrupt("exceeds its size limit"));
}
let magic = SHADOW_KEY_DICTIONARY_MAGIC.len();
let footer = 8 + 4 + magic;
if bytes.len() < magic + footer
|| &bytes[..magic] != SHADOW_KEY_DICTIONARY_MAGIC
|| &bytes[bytes.len() - magic..] != SHADOW_KEY_DICTIONARY_MAGIC
{
return Err(corrupt("framing is invalid"));
}
let body = &bytes[magic..bytes.len() - footer];
let footer_bytes = &bytes[bytes.len() - footer..bytes.len() - magic];
let stored_len = u64::from_le_bytes(footer_bytes[..8].try_into().expect("8 bytes"));
let stored_crc = u32::from_le_bytes(footer_bytes[8..12].try_into().expect("4 bytes"));
if stored_len != body.len() as u64 || crc32c(body).get() != stored_crc {
return Err(corrupt("checksum or length is invalid"));
}
let mut position = 0usize;
let read_u32 = |position: &mut usize| -> Result<u32> {
let end = position
.checked_add(4)
.filter(|end| *end <= body.len())
.ok_or_else(|| corrupt("is truncated"))?;
let value = u32::from_le_bytes(body[*position..end].try_into().expect("4 bytes"));
*position = end;
Ok(value)
};
if read_u32(&mut position)? != SHADOW_KEY_DICTIONARY_VERSION {
return Err(corrupt("has an unsupported version"));
}
let count = read_u32(&mut position)? as usize;
let mut dictionary = Self::default();
for _ in 0..count {
let length = read_u32(&mut position)? as usize;
let end = position
.checked_add(length)
.filter(|end| *end <= body.len())
.ok_or_else(|| corrupt("is truncated"))?;
let key = std::str::from_utf8(&body[position..end])
.map_err(|_| corrupt("holds a non-UTF-8 key"))?;
position = end;
if dictionary.ids.contains_key(key) {
return Err(corrupt("repeats a key"));
}
dictionary.intern(key, metadata_budget)?;
}
if position != body.len() {
return Err(corrupt("has trailing bytes"));
}
Ok(dictionary)
}
pub fn len(&self) -> usize {
self.keys.len()
}
pub fn is_empty(&self) -> bool {
self.keys.is_empty()
}
pub fn id(&self, key: &str) -> Option<PropertyId> {
self.ids.get(key).copied().map(PropertyId)
}
pub fn estimated_bytes(&self) -> u64 {
self.keys.iter().fold(0u64, |bytes, key| {
bytes.saturating_add(2 * (key.len() as u64 + SHADOW_KEY_ENTRY_OVERHEAD_BYTES))
})
}
pub fn intern(
&mut self,
key: &str,
metadata_budget: &mut ShadowMetadataBudget,
) -> Result<PropertyId> {
if let Some(id) = self.ids.get(key) {
return Ok(PropertyId(*id));
}
metadata_budget.charge(Self::entry_bytes(key), "key dictionary")?;
let id = u32::try_from(self.keys.len())
.ok()
.and_then(|index| index.checked_add(FIRST_DICTIONARY_COLUMN))
.ok_or_else(|| {
HawDBError::Storage(
"columnar shadow key dictionary exceeds the u32 id space".to_string(),
)
})?;
self.ids.insert(key.to_string(), id);
self.keys.push(key.to_string());
Ok(PropertyId(id))
}
fn entry_bytes(key: &str) -> u64 {
2 * (key.len() as u64 + SHADOW_KEY_ENTRY_OVERHEAD_BYTES)
}
pub fn key(&self, id: PropertyId) -> Option<&str> {
id.0.checked_sub(FIRST_DICTIONARY_COLUMN)
.and_then(|index| self.keys.get(index as usize))
.map(String::as_str)
}
fn encoded_len(&self) -> Result<u64> {
let body_len = self.keys.iter().try_fold(8u64, |bytes, key| {
bytes
.checked_add(4)
.and_then(|bytes| bytes.checked_add(key.len() as u64))
.ok_or_else(|| {
HawDBError::Storage(
"columnar shadow key dictionary length overflows u64".to_string(),
)
})
})?;
body_len
.checked_add((2 * SHADOW_KEY_DICTIONARY_MAGIC.len() + 12) as u64)
.ok_or_else(|| {
HawDBError::Storage(
"columnar shadow key dictionary framing length overflows u64".to_string(),
)
})
}
fn encode(&self, encoded_len: u64) -> Result<Vec<u8>> {
if encoded_len > MAX_SHADOW_KEY_DICTIONARY_BYTES {
return Err(HawDBError::Storage(
"columnar shadow key dictionary exceeds its size limit".to_string(),
));
}
let capacity = usize::try_from(encoded_len).map_err(|_| {
HawDBError::Storage(
"columnar shadow key dictionary exceeds the addressable memory range".to_string(),
)
})?;
let mut bytes = Vec::with_capacity(capacity);
bytes.extend(SHADOW_KEY_DICTIONARY_MAGIC);
let body_start = bytes.len();
bytes.extend(SHADOW_KEY_DICTIONARY_VERSION.to_le_bytes());
bytes.extend((self.keys.len() as u32).to_le_bytes());
for key in &self.keys {
bytes.extend((key.len() as u32).to_le_bytes());
bytes.extend(key.as_bytes());
}
let body_end = bytes.len();
let body_len = (body_end - body_start) as u64;
let body_crc = crc32c(&bytes[body_start..body_end]).get();
bytes.extend(body_len.to_le_bytes());
bytes.extend(body_crc.to_le_bytes());
bytes.extend(SHADOW_KEY_DICTIONARY_MAGIC);
debug_assert_eq!(bytes.len(), capacity);
Ok(bytes)
}
pub fn persist(
&self,
shadow_root: &Path,
metadata_budget: &mut ShadowMetadataBudget,
) -> Result<u64> {
let encoded_len = self.encoded_len()?;
metadata_budget.charge(encoded_len, "key dictionary serialization")?;
let bytes = match self.encode(encoded_len) {
Ok(bytes) => bytes,
Err(error) => {
metadata_budget.release(encoded_len);
return Err(error);
}
};
let path = shadow_root.join(SHADOW_KEY_DICTIONARY_FILE);
let tmp_path = shadow_root.join(format!(".{SHADOW_KEY_DICTIONARY_FILE}.tmp"));
let result = (|| -> Result<()> {
{
let mut file = File::create(&tmp_path)?;
file.write_all(&bytes)?;
file.sync_all()?;
}
durable_replace_file(&tmp_path, &path)?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(&tmp_path);
}
metadata_budget.release(encoded_len);
result.map(|()| bytes.len() as u64)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum InferredType {
Int,
Float,
Bool,
Str,
Binary,
Residual,
}
impl InferredType {
fn of(value: &Value) -> Self {
match value {
Value::Int(_) => Self::Int,
Value::Float(_) => Self::Float,
Value::Bool(_) => Self::Bool,
Value::String(_) => Self::Str,
Value::Binary(_) => Self::Binary,
Value::Null | Value::Uuid(_) | Value::List(_) | Value::Map(_) => Self::Residual,
}
}
fn merge(self, other: Self) -> Self {
if self == other {
self
} else {
Self::Residual
}
}
}
#[derive(Debug, Default)]
pub struct TablePropertyTypes {
inferred: BTreeMap<PropertyId, InferredType>,
}
impl TablePropertyTypes {
pub fn observe(
&mut self,
properties: &BTreeMap<String, Value>,
dictionary: &mut ShadowKeyDictionary,
metadata_budget: &mut ShadowMetadataBudget,
) -> Result<()> {
for (key, value) in properties {
let observed = InferredType::of(value);
let property_id = dictionary.intern(key, metadata_budget)?;
if !self.inferred.contains_key(&property_id) {
metadata_budget.charge(
SHADOW_PASS1_PROPERTY_OVERHEAD_BYTES,
"pass-1 property state",
)?;
}
self.inferred
.entry(property_id)
.and_modify(|current| *current = current.merge(observed))
.or_insert(observed);
}
Ok(())
}
}
pub struct ShadowTableLayout {
typed: Vec<PropertyId>,
typed_index: BTreeMap<PropertyId, usize>,
}
impl ShadowTableLayout {
pub fn typed(&self) -> &[PropertyId] {
&self.typed
}
pub fn typed_index(&self) -> &BTreeMap<PropertyId, usize> {
&self.typed_index
}
}
pub fn shadow_table_layout(
types: &TablePropertyTypes,
metadata_budget: &mut ShadowMetadataBudget,
) -> Result<ShadowTableLayout> {
metadata_budget.charge(SHADOW_LAYOUT_TABLE_OVERHEAD_BYTES, "typed table layout")?;
let mut typed = Vec::new();
let mut typed_index = BTreeMap::new();
for (property_id, inferred) in &types.inferred {
if *inferred != InferredType::Residual {
metadata_budget.charge(
SHADOW_LAYOUT_PROPERTY_OVERHEAD_BYTES,
"typed property layout",
)?;
typed_index.insert(*property_id, typed.len());
typed.push(*property_id);
}
}
Ok(ShadowTableLayout { typed, typed_index })
}
#[cfg(test)]
mod tests;