diff options
Diffstat (limited to 'autere_db/src')
| -rw-r--r-- | autere_db/src/common.rs | 574 | ||||
| -rw-r--r-- | autere_db/src/config.rs | 140 | ||||
| -rw-r--r-- | autere_db/src/engine.rs | 852 | ||||
| -rw-r--r-- | autere_db/src/lib.rs | 549 | ||||
| -rw-r--r-- | autere_db/src/lock.rs | 117 | ||||
| -rw-r--r-- | autere_db/src/log_reader_forward.rs | 134 | ||||
| -rw-r--r-- | autere_db/src/memtable_primary.rs | 40 | ||||
| -rw-r--r-- | autere_db/src/memtable_secondary.rs | 66 | ||||
| -rw-r--r-- | autere_db/src/record.rs | 38 | ||||
| -rw-r--r-- | autere_db/src/row.rs | 70 | ||||
| -rw-r--r-- | autere_db/src/schema.rs | 7 |
11 files changed, 2587 insertions, 0 deletions
diff --git a/autere_db/src/common.rs b/autere_db/src/common.rs new file mode 100644 index 0000000..0f6a058 --- /dev/null +++ b/autere_db/src/common.rs @@ -0,0 +1,574 @@ +use super::*; + +use std::collections::btree_map::Values; +// For Unix-like systems +#[cfg(unix)] +use std::os::unix::fs::MetadataExt; + +// For Windows +#[cfg(windows)] +use std::os::windows::fs::MetadataExt; + +pub const ACTIVE_SYMLINK_FILENAME: &str = "active"; +pub const LOCK_FILENAME: &str = "lock"; +pub const EXCL_LOCK_REQ_FILENAME: &str = "excl_lock_req"; +pub const INITIALIZED_FILENAME: &str = "initialized"; + +pub const METADATA_FILE_HEADER_SIZE: usize = 24; +pub const METADATA_ROW_LENGTH: usize = 16; +pub const LOCK_WAIT_MAX_MS: u64 = 1000; + +// Serialized value tags +pub const B_NULL: u8 = 0x0; +pub const B_INT: u8 = 0x1; +pub const B_DECIMAL: u8 = 0x2; +pub const B_STRING: u8 = 0x3; +pub const B_BYTES: u8 = 0x4; +// Tombstone marker tags +pub const B_LIVE: u8 = 0x0; +pub const B_TOMBSTONE: u8 = 0xFF; + +pub fn metadata_filename(num: u16) -> String { + format!("metadata.{}", num) +} + +pub type DBResult<A> = Result<A, DBError>; + +#[derive(Debug, Error)] +pub enum DBError { + #[error("lock request failed: {0}")] + LockRequestError(String), + #[error("validation failed: {0}")] + ValidationError(String), + #[error("consistency check failed: {0}")] + ConsistencyError(String), + #[error("invalid transaction: {0}")] + TransactionError(String), + #[error("unexpected IO error: {0}")] + IOError(#[from] io::Error), +} + +#[derive(Debug, Error)] +pub enum LogKeyMapError { + #[error("log key not found in map")] + NotFoundError, + #[error("attempted to remove last element of non-empty map")] + RemovingLastElementError, +} + +/// LogKey is a packed struct that contains: +/// - a log segment number (16 bits) +/// - a log index within the segment (48 bits) +#[derive(Debug, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)] +pub struct LogKey(u64); + +impl LogKey { + pub fn new(segment_num: u16, index: u64) -> Self { + assert!(index < (1 << 48), "Index must fit in 48 bits"); + LogKey((segment_num as u64) << 48 | index) + } + + pub fn segment_num(&self) -> u16 { + (self.0 >> 48) as u16 + } + + pub fn index(&self) -> u64 { + self.0 & 0x0000_FFFF_FFFF_FFFF + } +} + +/// LogKeyMap is a non-empty map of PK => LogKey mappings. +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct LogKeyMap { + map: BTreeMap<IndexableValue, LogKey>, +} + +impl LogKeyMap { + /// Create a new LogKeyMap with an initial mapping. + /// The initial mapping is required since LogKeyMap must be non-empty. + pub fn new_with_initial(pk: IndexableValue, log_key: LogKey) -> Self { + let mut map = BTreeMap::new(); + map.insert(pk, log_key); + LogKeyMap { map } + } + + pub fn contains_pk(&self, key: &IndexableValue) -> bool { + self.map.contains_key(key) + } + + /// The number of LogKeys in the map. + pub fn len(&self) -> usize { + self.map.len() + } + + /// Insert a PK -> LogKey mapping into the map. + pub fn insert(&mut self, key: IndexableValue, log_key: LogKey) { + self.map.insert(key, log_key); + } + + /// Remove a mapping from the map. Return Ok(()) if the key was found and removed. + /// Return `LogKeyMapError::RemovingLastElementError` if trying to remove the last element. + /// Return `LogKeyMapError::NotFoundError` if the key was not found. + pub fn remove_pk(&mut self, key: &IndexableValue) -> Result<(), LogKeyMapError> { + if self.map.len() == 1 { + return Err(LogKeyMapError::RemovingLastElementError); + } + let removed = self.map.remove(key); + + if removed.is_none() { + return Err(LogKeyMapError::NotFoundError); + } + + assert!( + self.map.len() > 0, + "LogKeyMap should not be empty after removal" + ); + + Ok(()) + } + + /// Get a reference to the set of LogKeys. + pub fn log_keys(&self) -> Values<IndexableValue, LogKey> { + self.map.values() + } +} + +pub static APPEND_MODE: Lazy<fs::OpenOptions> = Lazy::new(|| { + let mut options = fs::OpenOptions::new(); + options.read(true).append(true); + options +}); +pub static READ_MODE: Lazy<fs::OpenOptions> = Lazy::new(|| { + let mut options = fs::OpenOptions::new(); + options.read(true); + options +}); +pub static WRITE_MODE: Lazy<fs::OpenOptions> = Lazy::new(|| { + let mut options = fs::OpenOptions::new(); + options.read(true).write(true); + options +}); + +pub struct MetadataHeader { + pub version: u8, + pub uuid: Uuid, +} + +const METADATA_HEADER_PADDING: &[u8] = &[0; 7]; +impl MetadataHeader { + pub fn serialize(&self) -> [u8; METADATA_FILE_HEADER_SIZE] { + let mut header = [0u8; METADATA_FILE_HEADER_SIZE]; + header[0] = self.version; + header[1..8].copy_from_slice(METADATA_HEADER_PADDING); + header[8..].copy_from_slice(self.uuid.as_bytes()); + + header + } + + pub fn deserialize(bytes: &[u8]) -> Self { + assert_eq!(bytes.len(), METADATA_FILE_HEADER_SIZE); + + let version = bytes[0]; + let uuid = Uuid::from_slice(&bytes[8..24]).expect("Failed to deserialize Uuid"); + + MetadataHeader { version, uuid } + } +} + +#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)] +pub enum IndexableValue { + Null, + Int(i64), + Decimal(Decimal), + String(String), +} + +#[derive(Debug, Clone)] +pub enum Value { + Null, + Int(i64), + Decimal(Decimal), + String(String), + Bytes(Vec<u8>), +} + +impl PartialEq for Value { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Value::Int(a), Value::Int(b)) => a == b, + (Value::Decimal(a), Value::Decimal(b)) => a == b, + (Value::String(a), Value::String(b)) => a == b, + (Value::Bytes(a), Value::Bytes(b)) => a == b, + (Value::Null, Value::Null) => true, + _ => false, + } + } +} +impl Eq for Value {} + +impl Value { + pub fn serialize(&self) -> Vec<u8> { + match self { + Value::Null => vec![B_NULL], + Value::Int(i) => { + let mut bytes = Vec::with_capacity(1 + 16); + bytes.push(B_INT); + bytes.extend_from_slice(&i.to_be_bytes()); + bytes + } + Value::Decimal(d) => { + let mut bytes = Vec::with_capacity(1 + 16); + bytes.push(B_DECIMAL); + bytes.extend_from_slice(&d.serialize()); + bytes + } + Value::String(s) => { + let len = s.len(); + let mut bytes = Vec::with_capacity(1 + 8 + len); + bytes.push(B_STRING); + bytes.extend_from_slice(&(len as u64).to_be_bytes()); + bytes.extend_from_slice(s.as_bytes()); + bytes + } + Value::Bytes(b) => { + let len = b.len(); + let mut bytes = Vec::with_capacity(1 + 8 + len); + bytes.push(B_BYTES); + bytes.extend_from_slice(&(len as u64).to_be_bytes()); + bytes.extend_from_slice(b); + bytes + } + } + } + + /// Deserialize a Value from a byte slice. + /// Returns the deserialized Value and the number of bytes consumed. + pub fn deserialize(bytes: &[u8]) -> (Value, usize) { + match bytes[0] { + B_NULL => (Value::Null, 1), + B_INT => { + let mut int_bytes = [0; 8]; + int_bytes.copy_from_slice(&bytes[1..1 + 8]); + (Value::Int(i64::from_be_bytes(int_bytes)), 1 + 8) + } + B_DECIMAL => { + let mut decimal_bytes = [0; 16]; + decimal_bytes.copy_from_slice(&bytes[1..1 + 16]); + (Value::Decimal(Decimal::deserialize(decimal_bytes)), 1 + 16) + } + B_STRING => { + let length_bytes = &bytes[1..1 + 8]; + let length = u64::from_be_bytes(length_bytes.try_into().unwrap()) as usize; + ( + Value::String( + String::from_utf8(bytes[1 + 8..1 + 8 + length].to_vec()).unwrap(), + ), + 1 + 8 + length, + ) + } + B_BYTES => { + let length_bytes = &bytes[1..1 + 8]; + let length = u64::from_be_bytes(length_bytes.try_into().unwrap()) as usize; + ( + Value::Bytes(bytes[1 + 8..1 + 8 + length].to_vec()), + 1 + 8 + length, + ) + } + _ => panic!("Invalid tag: {}", bytes[0]), + } + } + + pub fn as_indexable(&self) -> Option<IndexableValue> { + match self { + Value::Null => Some(IndexableValue::Null), + Value::Int(i) => Some(IndexableValue::Int(*i)), + Value::Decimal(d) => Some(IndexableValue::Decimal(d.clone())), + Value::String(s) => Some(IndexableValue::String(s.clone())), + _ => None, + } + } +} + +pub fn get_secondary_memtable_index_by_field(sks: &Vec<String>, field: &str) -> Option<usize> { + sks.iter().position(|schema_field| schema_field == field) +} + +pub fn is_file_same_as_path(file: &File, path: &PathBuf) -> DBResult<bool> { + // Get the metadata for the open file handle + let file_metadata = file.metadata()?; + + // Get the metadata for the file at the specified path + let path_metadata = metadata(path)?; + + // Platform-specific comparison + #[cfg(unix)] + { + Ok( + file_metadata.dev() == path_metadata.dev() + && file_metadata.ino() == path_metadata.ino(), + ) + } + + #[cfg(windows)] + { + Ok(file_metadata.file_index() == path_metadata.file_index() + && file_metadata.volume_serial_number() == path_metadata.volume_serial_number()) + } +} + +pub fn symlink(original: &Path, link: &Path) -> io::Result<()> { + #[cfg(unix)] + { + std::os::unix::fs::symlink(original, link) + } + + #[cfg(windows)] + { + std::os::windows::fs::symlink_file(original, link) + } +} + +/// Set the active segment to the segment with the given ordinal number. +pub fn set_active_segment(data_dir_path: &Path, segment_num: u16) -> DBResult<()> { + let tmp_uuid = Uuid::new_v4(); + let tmp_filename = format!("active_{}", tmp_uuid.to_string()); + let tmp_path = data_dir_path.join(tmp_filename); + + let metadata_filename = format!("metadata.{}", segment_num); + let metadata_path = Path::new(&metadata_filename); + let active_symlink = data_dir_path.join(ACTIVE_SYMLINK_FILENAME); + + symlink(&metadata_path, &tmp_path)?; + fs::rename(&tmp_path, &active_symlink)?; + + Ok(()) +} + +/// Create a new segment metadata file and return its number and path. +/// A metadata file contains the segment metadata, including the UUID of the data file. +/// See `ARCHITECTURE.md` for the file format. +pub fn create_segment_metadata_file( + data_dir_path: &Path, + data_file_uuid: &Uuid, +) -> DBResult<(u16, PathBuf)> { + let current_greatest_num = greatest_segment_number(data_dir_path)?; + let new_num = current_greatest_num + 1; + + let metadata_filename = format!("metadata.{}", new_num); + let metadata_path = data_dir_path.join(metadata_filename); + + let mut metadata_file = fs::OpenOptions::new() + .create(true) + .append(true) + .open(&metadata_path)?; + + let metadata_header = MetadataHeader { + version: 1, + uuid: *data_file_uuid, + }; + + metadata_file.write_all(&metadata_header.serialize())?; + metadata_file.flush()?; + + let len = metadata_file.seek(io::SeekFrom::End(0))?; + assert!(len >= METADATA_FILE_HEADER_SIZE as u64); + assert_eq!((len - METADATA_FILE_HEADER_SIZE as u64) % 16, 0); + + Ok((new_num, metadata_path)) +} + +/// Parse the segment number from a metadata file path +pub fn parse_segment_number(metadata_path: &Path) -> DBResult<u16> { + let filename = metadata_path + .file_name() + .expect("No filename in symlink") + .to_str() + .expect("Filename was not valid UTF-8"); + + // parse number from format "metadata.1" + let segment_number = filename + .split('.') + .last() + .expect("Filename did not have a number") + .parse::<u16>(); + + segment_number.map_err(|_| { + DBError::ValidationError("Failed to parse segment number from filename".to_owned()) + }) +} + +/// Get the number of the segment with the greatest ordinal. +/// This is the newest segment, i.e. the one that is pointed to by the `active` symlink. +/// If there are no segments yet, returns 0. +pub fn greatest_segment_number(data_dir_path: &Path) -> DBResult<u16> { + let active_symlink = data_dir_path.join(ACTIVE_SYMLINK_FILENAME); + + if !fs::exists(&active_symlink)? { + return Ok(0); + } + + let segment_metadata_path = fs::read_link(&active_symlink)?; + parse_segment_number(&segment_metadata_path) +} + +/// Create a new segment data file and return its UUID. +/// A data file contains the segment data, tightly packed without separators. +/// An accompanying metadata file is required to interpret the data. +pub fn create_segment_data_file(data_dir_path: &Path) -> DBResult<(Uuid, PathBuf)> { + let uuid = Uuid::new_v4(); + let new_segment_path = data_dir_path.join(uuid.to_string()); + fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_segment_path)?; + + Ok((uuid, new_segment_path)) +} + +/// Reads the metadata header from the metadata file. +/// Leaves the file seek head at the beginning of the records, after the header. +pub fn read_metadata_header(metadata_file: &mut fs::File) -> DBResult<MetadataHeader> { + metadata_file.seek(SeekFrom::Start(0))?; + let mut buf = [0u8; METADATA_FILE_HEADER_SIZE]; + metadata_file.read_exact(&mut buf)?; + + let header = MetadataHeader::deserialize(&buf); + Ok(header) +} + +pub fn validate_metadata_header(header: &MetadataHeader) -> DBResult<()> { + if header.version != 1 { + return Err(DBError::ValidationError( + "Unsupported metadata file version".to_owned(), + )); + } + + Ok(()) +} + +pub enum IsMetadatafileValidResult { + Ok, + ReplaceFile, + TruncateToSize(u64), +} + +pub fn is_metadata_file_valid(metadata_file: &mut fs::File) -> DBResult<IsMetadatafileValidResult> { + let size = metadata_file.seek(SeekFrom::End(0))? as usize; + + if size < METADATA_FILE_HEADER_SIZE { + return Ok(IsMetadatafileValidResult::ReplaceFile); + } + + // The data section must be a multiple of 16 bytes. + // Otherwise, the non-aligned part of the file is dropped. + let data_section_len = size - METADATA_FILE_HEADER_SIZE; + let remainder = data_section_len % 16; + if remainder != 0 { + return Ok(IsMetadatafileValidResult::TruncateToSize( + (size - remainder) as u64, + )); + } + + Ok(IsMetadatafileValidResult::Ok) +} + +/// Check that the active metadata file is well-formed and repair it if necessary. +/// The metadata file is considered well-formed if its size is, in pseudocode, `header_size + n * record_size`. +/// If the file is not well-formed, it is truncated to the last well-formed record using +/// a temporary file and an atomic move operation. +/// +/// `self.active_metadata_file` must be a locked file handle opened with read permissions. +/// The function leaves the seek head in an unspecified position. +/// +/// Returns `false` if the file was repaired and rotated, `true` if no action was taken. +pub fn ensure_active_metadata_is_valid( + data_dir: &Path, + metadata_file: &mut fs::File, +) -> DBResult<bool> { + let current_len = metadata_file.seek(SeekFrom::End(0))? as usize; + + match is_metadata_file_valid(metadata_file)? { + IsMetadatafileValidResult::Ok => return Ok(true), + IsMetadatafileValidResult::ReplaceFile => { + let active_target = fs::read_link(data_dir.join(ACTIVE_SYMLINK_FILENAME))?; + let active_path = data_dir.join(&active_target); + warn!( + "Metadata file \"{}\" is malformed ({} bytes), replacing it with an empty file", + active_target.display(), + current_len, + ); + let mut tmp_file = tempfile::NamedTempFile::new()?; + + let header = MetadataHeader { + version: 1, + uuid: Uuid::new_v4(), + }; + + tmp_file.write_all(&header.serialize())?; + tmp_file.flush()?; + + fs::rename(tmp_file.path(), active_path)?; + + debug!("Replaced metadata file"); + return Ok(false); + } + IsMetadatafileValidResult::TruncateToSize(new_size) => { + let active_target = fs::read_link(data_dir.join(ACTIVE_SYMLINK_FILENAME))?; + let active_path = data_dir.join(&active_target); + warn!( + "Metadata file \"{}\" is malformed ({} bytes), truncating it to {} bytes", + active_target.display(), + current_len, + new_size + ); + + let mut tmp_file = tempfile::NamedTempFile::new()?; + + let mut buf = vec![0; new_size as usize]; + metadata_file.seek(SeekFrom::Start(0))?; + metadata_file.read_exact(&mut buf)?; + + tmp_file.write_all(&buf)?; + tmp_file.flush()?; + + fs::rename(tmp_file.path(), active_path)?; + + debug!("Truncated metadata file"); + return Ok(false); + } + } +} + +pub struct OwnedBounds<T> { + start: Bound<T>, + end: Bound<T>, +} + +impl<T> OwnedBounds<T> { + pub fn new(start: Bound<T>, end: Bound<T>) -> Self { + OwnedBounds { start, end } + } +} + +impl<T> RangeBounds<T> for OwnedBounds<T> { + fn start_bound(&self) -> Bound<&T> { + self.start.as_ref() + } + + fn end_bound(&self) -> Bound<&T> { + self.end.as_ref() + } +} + +#[derive(Debug, Clone)] +pub struct QueryParams { + pub offset: usize, + pub limit: usize, + pub sort_asc: bool, +} + +pub static DEFAULT_QUERY_PARAMS: QueryParams = QueryParams { + offset: 0, + limit: usize::MAX, + sort_asc: true, +}; diff --git a/autere_db/src/config.rs b/autere_db/src/config.rs new file mode 100644 index 0000000..e8c6fa1 --- /dev/null +++ b/autere_db/src/config.rs @@ -0,0 +1,140 @@ +use super::*; + +pub struct ConfigBuilder { + data_dir: Option<String>, + segment_size: Option<usize>, + write_durability: Option<WriteDurability>, + read_consistency: Option<ReadConsistency>, + + fields: Option<Vec<String>>, + primary_key: Option<String>, + secondary_keys: Option<Vec<String>>, +} + +impl ConfigBuilder { + pub fn new() -> ConfigBuilder { + ConfigBuilder { + data_dir: None, + segment_size: None, + write_durability: None, + read_consistency: None, + + fields: None, + primary_key: None, + secondary_keys: None, + } + } + + /// The directory where the database will store its data. + pub fn data_dir(mut self, data_dir: impl Into<String>) -> Self { + self.data_dir = Some(data_dir.into()); + self + } + + /// The maximum size of a segment file in bytes. + /// Once a segment file reaches this size, it can be closed, rotated and compacted. + /// Note that this is not a hard limit: if `db.do_maintenance_tasks()` is not called, + /// the segment file may continue to grow. + pub fn segment_size(mut self, segment_size: usize) -> Self { + self.segment_size = Some(segment_size); + self + } + + /// The write durability policy for the database. + /// This determines how writes are persisted to disk. + /// The default is WriteDurability::Flush. + pub fn write_durability(mut self, write_durability: WriteDurability) -> Self { + self.write_durability = Some(write_durability); + self + } + + /// The read consistency policy for the database. + /// This determines how recent writes are visible when reading. + /// See individual `ReadConsistency` enum values for more information. + /// The default is ReadConsistency::Strong. + pub fn read_consistency(mut self, read_consistency: ReadConsistency) -> Self { + self.read_consistency = Some(read_consistency); + self + } + + pub fn fields(mut self, schema: Vec<impl Into<String>>) -> Self { + self.fields = Some(schema.into_iter().map(|s| s.into()).collect()); + self + } + + pub fn primary_key(mut self, primary_key: impl Into<String>) -> Self { + self.primary_key = Some(primary_key.into()); + self + } + + pub fn secondary_keys(mut self, secondary_keys: Vec<impl Into<String>>) -> Self { + self.secondary_keys = Some(secondary_keys.into_iter().map(|s| s.into()).collect()); + self + } + + pub fn initialize(self) -> DBResult<DB> { + let schema = self + .fields + .ok_or_else(|| DBError::ValidationError("Schema not set".to_string()))?; + let primary_key = self + .primary_key + .ok_or_else(|| DBError::ValidationError("Primary key not set".to_string()))?; + + let config = Config { + schema, + primary_key, + secondary_keys: self.secondary_keys.unwrap_or_default(), + + data_dir: self.data_dir.clone().unwrap_or("db_data".to_string()), + segment_size: self.segment_size.unwrap_or(4 * 1024 * 1024), // 4MB + write_durability: self + .write_durability + .clone() + .unwrap_or(WriteDurability::Flush), + read_consistency: self + .read_consistency + .clone() + .unwrap_or(ReadConsistency::Strong), + }; + + DB::initialize(config) + } +} + +#[derive(Clone)] +pub struct Config { + pub schema: Vec<String>, + pub primary_key: String, + pub secondary_keys: Vec<String>, + pub data_dir: String, + pub segment_size: usize, + pub write_durability: WriteDurability, + pub read_consistency: ReadConsistency, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum ReadConsistency { + /// Reads by client A are guaranteed to see writes by themselves and any writes by other clients B + /// that were done before last index refresh. You must call `refresh_indexes()` manually to refresh indexes. + Eventual, + /// Reads by client A are guaranteed to see all writes. This is slower: all reads must first + /// refresh indexes. + Strong, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum WriteDurability { + /// Changes are written to the OS write buffer but not immediately synced to disk. + /// This is generally recommended. Most OSes will sync the write buffer to disk within a few seconds. + Flush, + /// Changes are written to the OS write buffer and synced to disk immediately. + /// Offers the best durability guarantees but is a lot slower. + FlushSync, +} + +impl Display for WriteDurability { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { + write!(f, "{:?}", self)?; + Ok(()) + } +} diff --git a/autere_db/src/engine.rs b/autere_db/src/engine.rs new file mode 100644 index 0000000..6fb9edf --- /dev/null +++ b/autere_db/src/engine.rs @@ -0,0 +1,852 @@ +use super::*; + +pub struct Engine { + pub config: Config, + pub lock_manager: LockManager, + + data_dir_path: PathBuf, + primary_key_index: usize, + refresh_next_logkey: LogKey, + + pub tx_active: bool, + pub tx_log: Vec<TxEntry>, + + active_metadata_file: fs::File, + active_data_file: fs::File, + + // TODO: these could be made private. Currently they are public for testing in lib.rs. + pub primary_memtable: PrimaryMemtable, + pub secondary_memtables: Vec<SecondaryMemtable>, +} + +impl Engine { + pub fn initialize(config: Config) -> DBResult<Engine> { + info!("Initializing DB..."); + // If data_dir does not exist or is empty, create it and any necessary files. + // After creation, the directory should always be in a complete state without missing files. + + // Ensure the data directory exists + let data_dir_path = Path::new(&config.data_dir).to_path_buf(); + match fs::create_dir(&data_dir_path) { + Ok(_) => {} + Err(e) => { + if e.kind() != io::ErrorKind::AlreadyExists { + return Err(DBError::IOError(e)); + } + } + } + + // Create the lock file first to prevent multiple concurrent initializations + let mut lock_manager = LockManager::new(data_dir_path.clone())?; + lock_manager.lock_exclusive()?; + + // We have acquired the lock, check if the data directory is in a complete state + // If not, initialize it, otherwise skip. + if !fs::exists(data_dir_path.join(INITIALIZED_FILENAME))? { + // Delete all files except the lock files to ensure a clean state + for entry in fs::read_dir(&data_dir_path)? { + let entry = entry?; + let path = entry.path(); + if path.is_file() + && path.file_name().unwrap() != LOCK_FILENAME + && path.file_name().unwrap() != EXCL_LOCK_REQ_FILENAME + { + fs::remove_file(&path)?; + } + } + + // Create the initial segment files + let (segment_uuid, _) = create_segment_data_file(&data_dir_path)?; + let (segment_num, _) = create_segment_metadata_file(&data_dir_path, &segment_uuid)?; + set_active_segment(&data_dir_path, segment_num)?; + + // Create the initialized file to indicate that the directory is in a complete state + fs::File::create(data_dir_path.join(INITIALIZED_FILENAME))?; + } + + // Calculate the index of the primary value in a record + let primary_key_index = config + .schema + .iter() + .position(|field| field == &config.primary_key) + .ok_or(DBError::ValidationError( + "Primary key not found in schema after initialize".to_owned(), + ))?; + + // Join primary key and secondary keys vec into a single vec + let mut all_keys = vec![&config.primary_key]; + all_keys.extend(&config.secondary_keys); + + // If any of the keys is not in the schema or + // is not an IndexableValue, return an error + for &key in &all_keys { + let _ = config.schema.iter().find(|&field| field == key).ok_or( + DBError::ValidationError("Key must be present in the field schema".to_owned()), + )?; + } + let primary_memtable = PrimaryMemtable::new(); + let secondary_memtables = config + .secondary_keys + .iter() + .map(|_| SecondaryMemtable::new()) + .collect(); + + let active_symlink = Path::new(&config.data_dir).join(ACTIVE_SYMLINK_FILENAME); + + let active_target = fs::read_link(&active_symlink)?; + let active_metadata_path = Path::new(&config.data_dir).join(active_target); + let mut active_metadata_file = APPEND_MODE.open(&active_metadata_path)?; + + let active_metadata_header = read_metadata_header(&mut active_metadata_file)?; + validate_metadata_header(&active_metadata_header)?; + + let active_data_path = + Path::new(&config.data_dir).join(active_metadata_header.uuid.to_string()); + let active_data_file = APPEND_MODE.open(&active_data_path)?; + + let mut engine = Engine { + config, + lock_manager, + data_dir_path, + primary_key_index, + primary_memtable, + secondary_memtables, + active_metadata_file, + active_data_file, + refresh_next_logkey: LogKey::new(1, 0), + tx_active: false, + tx_log: vec![], + }; + + info!("Rebuilding memtable indexes..."); + engine.refresh_indexes()?; + + info!("Database ready."); + + engine.lock_manager.unlock()?; + Ok(engine) + } + + pub fn refresh_indexes(&mut self) -> DBResult<()> { + let active_symlink_path = self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME); + let active_target = fs::read_link(active_symlink_path)?; + let active_metadata_path = self.data_dir_path.join(active_target); + + let to_segnum = parse_segment_number(&active_metadata_path)?; + let from_segnum = self.refresh_next_logkey.segment_num(); + let mut from_index = self.refresh_next_logkey.index(); + + for segnum in from_segnum..=to_segnum { + let metadata_path = self.data_dir_path.join(metadata_filename(segnum)); + let mut metadata_file = READ_MODE.open(&metadata_path)?; + + let metadata_len = metadata_file.seek(SeekFrom::End(0))?; + if (metadata_len - METADATA_FILE_HEADER_SIZE as u64) % METADATA_ROW_LENGTH as u64 != 0 { + return Err(DBError::ConsistencyError(format!( + "Metadata file {} has invalid size: {}", + metadata_path.display(), + metadata_len + ))); + } + + let metadata_header = read_metadata_header(&mut metadata_file)?; + validate_metadata_header(&metadata_header)?; + + let data_path = self.data_dir_path.join(metadata_header.uuid.to_string()); + let data_file = READ_MODE.open(data_path)?; + + for ForwardLogReaderItem { row, index } in + ForwardLogReader::new_with_index(metadata_file, data_file, from_index) + { + let log_key = LogKey::new(segnum, index); + + if row.tombstone { + self.remove_row_from_memtables(&row.values); + } else { + self.insert_row_to_memtables(log_key, row.values); + } + + // Update from_index in case this is the last iteration: we need to know the next + // index that should be read on later invocations of refresh_indexes. + from_index = index + 1 + } + + // If there are still segments to read, set from_index to zero to read them + // from beginning. Otherwise we leave from_index as the index of the next record to read. + if segnum != to_segnum { + from_index = 0 + } + } + + self.refresh_next_logkey = LogKey::new(to_segnum, from_index); + + Ok(()) + } + + fn insert_row_to_memtables(&mut self, log_key: LogKey, row_values: Vec<Value>) { + let pk = row_values[self.primary_key_index].as_indexable().unwrap(); + + for (sk_index, sk_field) in self.config.secondary_keys.iter().enumerate() { + let secondary_memtable = &mut self.secondary_memtables[sk_index]; + let sk_field_index = self + .config + .schema + .iter() + .position(|f| sk_field == f) + .unwrap(); + let sk = row_values[sk_field_index].as_indexable().unwrap(); + + secondary_memtable.set(pk.clone(), sk, log_key.clone()); + } + + // Doing this last because this moves log_key + self.primary_memtable.set(pk, log_key); + } + + fn remove_row_from_memtables(&mut self, row_values: &Vec<Value>) { + let pk = row_values[self.primary_key_index].as_indexable().unwrap(); + + if let Some(_) = self.primary_memtable.remove(&pk) { + for (sk_index, sk_field) in self.config.secondary_keys.iter_mut().enumerate() { + let secondary_memtable = &mut self.secondary_memtables[sk_index]; + let sk_field_index = self + .config + .schema + .iter() + .position(|f| sk_field == f) + .unwrap(); + let sk = row_values[sk_field_index].as_indexable().unwrap(); + + secondary_memtable.remove(&pk, &sk); + } + } + } + + pub fn upsert_record(&mut self, record: Row) -> DBResult<()> { + debug!("Opening file in append mode..."); + + if !self.ensure_metadata_file_is_active()? + || !ensure_active_metadata_is_valid( + &self.data_dir_path, + &mut self.active_metadata_file, + )? + { + // The log file has been rotated, so we must try again + return self.upsert_record(record); + } + + self.tx_log.push(TxEntry::Upsert { row: record }); + + if !self.tx_active { + self.commit_transaction()?; + self.tx_log.clear(); + } + + Ok(()) + } + + pub fn batch_find_by_records<'a>( + &mut self, + field: &str, + values: impl Iterator<Item = &'a Value>, + params: &QueryParams, + ) -> DBResult<Vec<(usize, Row)>> { + let indexables = values + .map(|value| { + value.as_indexable().ok_or(DBError::ValidationError( + "Queried value must be indexable".to_owned(), + )) + }) + .collect::<DBResult<Vec<IndexableValue>>>()?; + + // Otherwise, continue with querying secondary indexes. + debug!("Finding all records with matching fields"); + + if self.config.read_consistency == ReadConsistency::Strong { + self.refresh_indexes()?; + } + + let log_key_batches = indexables + .into_iter() + .map(|query_key| { + if field == &self.config.primary_key { + let opt = self.primary_memtable.get(&query_key); + let log_keys = match opt { + Some(log_key) => vec![log_key], + None => vec![], + }; + Ok(log_keys) + } else { + let smemtable_index = match get_secondary_memtable_index_by_field( + &self.config.secondary_keys, + field, + ) { + Some(index) => index, + None => { + return Err(DBError::ValidationError( + "Cannot find_by by non-indexed key".to_owned(), + )) + } + }; + + let log_keys = self.secondary_memtables[smemtable_index] + .find_by(&query_key) + .into_iter() + .collect(); + Ok(log_keys) + } + }) + .collect::<DBResult<Vec<Vec<&LogKey>>>>()?; + + debug!("Found log keys in memtable: {:?}", log_key_batches); + + let mut tagged = vec![]; + for (tag, batch) in log_key_batches.into_iter().enumerate() { + let mapped = batch.into_iter().map(|log_key| (tag, log_key)); + tagged.extend(mapped); + } + + if !params.sort_asc { + tagged.reverse(); + } + let bound_low = params.offset; + let bound_high = (params.offset + params.limit).min(tagged.len()); + let sliced = &tagged[bound_low..bound_high]; + let mut tagged_records = self.read_tagged_log_keys(sliced.into_iter())?; + + debug!("Read {} records", tagged_records.len()); + + if !params.sort_asc { + tagged_records.reverse(); + } + Ok(tagged_records) + } + + /// Read records from segment files based on log keys. + /// The log keys are accompanied by an integer tag that can be used to identify and group them later. + fn read_tagged_log_keys<'a>( + &self, + log_keys: impl Iterator<Item = &'a (usize, &'a LogKey)>, + ) -> DBResult<Vec<(usize, Row)>> { + let mut records = vec![]; + let mut log_keys_map = BTreeMap::new(); + + for (tag, log_key) in log_keys { + if !log_keys_map.contains_key(&log_key.segment_num()) { + log_keys_map.insert(log_key.segment_num(), vec![(tag, log_key.index())]); + } else { + log_keys_map + .get_mut(&log_key.segment_num()) + .unwrap() + .push((tag, log_key.index())); + } + } + + for (segment_num, mut segment_indexes) in log_keys_map { + segment_indexes.sort_unstable(); + + let metadata_path = &self.data_dir_path.join(metadata_filename(segment_num)); + let mut metadata_file = READ_MODE.open(&metadata_path)?; + + let metadata_header = read_metadata_header(&mut metadata_file)?; + + let data_path = &self.data_dir_path.join(metadata_header.uuid.to_string()); + let mut data_file = READ_MODE.open(&data_path)?; + + let header_size = METADATA_FILE_HEADER_SIZE as i64; + let row_length = METADATA_ROW_LENGTH as i64; + let mut current_metadata_offset = header_size; + for (tag, segment_index) in segment_indexes { + let new_metadata_offset = header_size + segment_index as i64 * row_length; + metadata_file.seek_relative(new_metadata_offset - current_metadata_offset)?; + + let mut metadata_buf = [0; METADATA_ROW_LENGTH]; + metadata_file.read_exact(&mut metadata_buf)?; + + let data_offset = u64::from_be_bytes(metadata_buf[0..8].try_into().unwrap()); + let data_length = u64::from_be_bytes(metadata_buf[8..16].try_into().unwrap()); + assert!(data_length > 0); + + data_file.seek(SeekFrom::Start(data_offset))?; + + let mut data_buf = vec![0; data_length as usize]; + data_file.read_exact(&mut data_buf)?; + + let record = Row::deserialize(&data_buf); + records.push((*tag, record)); + + current_metadata_offset = new_metadata_offset + row_length; + } + } + + Ok(records) + } + + pub fn range_by_records<B: RangeBounds<Value>>( + &mut self, + field: &str, + range: B, + params: &QueryParams, + ) -> DBResult<Vec<Row>> { + fn range_bound_to_indexable(bound: Bound<&Value>) -> DBResult<Bound<IndexableValue>> { + match bound { + Bound::Included(value) => value + .as_indexable() + .ok_or(DBError::ValidationError( + "Queried value must be indexable".to_owned(), + )) + .map(Bound::Included), + Bound::Excluded(value) => value + .as_indexable() + .ok_or(DBError::ValidationError( + "Queried value must be indexable".to_owned(), + )) + .map(Bound::Excluded), + Bound::Unbounded => Ok(Bound::Unbounded), + } + } + + let start_indexable = range_bound_to_indexable(range.start_bound())?; + let end_indexable = range_bound_to_indexable(range.end_bound())?; + + let indexable_bounds = OwnedBounds::new(start_indexable, end_indexable); + + if self.config.read_consistency == ReadConsistency::Strong { + self.refresh_indexes()?; + } + + let log_keys = if field == &self.config.primary_key { + self.primary_memtable.range(indexable_bounds) + } else { + let index = get_secondary_memtable_index_by_field(&self.config.secondary_keys, field) + .ok_or_else(|| { + DBError::ValidationError("Cannot range_by by non-indexed key".to_owned()) + })?; + + self.secondary_memtables[index].range(indexable_bounds) + }; + + let mut log_key_batches: Vec<(usize, &LogKey)> = + log_keys.into_iter().map(|log_key| (0, log_key)).collect(); + + if !params.sort_asc { + log_key_batches.reverse(); + } + + let bound_low = params.offset; + let bound_high = (params.offset + params.limit).min(log_key_batches.len()); + let sliced = &log_key_batches[bound_low..bound_high]; + let tagged_records = self.read_tagged_log_keys(sliced.into_iter()); + + let mut result_records: Vec<Row> = + tagged_records?.into_iter().map(|(_, rec)| rec).collect(); + + if !params.sort_asc { + result_records.reverse(); + } + + Ok(result_records) + } + + /// Ensures that the `self.metadata_file` and `self.data_file` handles are still pointing to the correct files. + /// If the segment has been rotated, the handle will be closed and reopened. + /// Returns `false` if the file has been rotated and the handle has been reopened, `true` otherwise. + fn ensure_metadata_file_is_active(&mut self) -> DBResult<bool> { + let active_target = fs::read_link(&self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME))?; + let active_metadata_path = &self.data_dir_path.join(active_target); + + let correct = is_file_same_as_path(&self.active_metadata_file, &active_metadata_path)?; + if !correct { + debug!("Metadata file has been rotated. Reopening..."); + let metadata_file = APPEND_MODE.open(&active_metadata_path)?; + + let metadata_header = read_metadata_header(&mut self.active_metadata_file)?; + + validate_metadata_header(&metadata_header)?; + + let data_file_path = &self.data_dir_path.join(metadata_header.uuid.to_string()); + + self.active_metadata_file = metadata_file; + self.active_data_file = APPEND_MODE.open(&data_file_path)?; + + return Ok(false); + } else { + return Ok(true); + } + } + + pub fn delete_by_field(&mut self, field: &str, value: &Value) -> DBResult<Vec<Row>> { + let recs: Vec<Row> = self + .batch_find_by_records(field, std::iter::once(value), &DEFAULT_QUERY_PARAMS)? + .into_iter() + .map(|(_, mut rec)| { + rec.tombstone = true; + rec + }) + .collect(); + + // TODO: refactor the clone out of here + for record in &recs { + self.tx_log.push(TxEntry::Delete { + row: record.clone(), + }); + } + + if !self.tx_active { + self.commit_transaction()?; + self.tx_log.clear(); + } + + debug!("Records deleted"); + + Ok(recs) + } + + pub fn commit_transaction(&mut self) -> DBResult<()> { + let active_symlink_path = self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME); + let active_target = fs::read_link(active_symlink_path)?; + let segment_num = parse_segment_number(&active_target)?; + + let initial_data_offset = self.active_data_file.seek(SeekFrom::End(0))?; + let initial_metadata_offset = self.active_metadata_file.seek(SeekFrom::End(0))?; + let mut serialized_data: Vec<u8> = vec![]; + let mut serialized_metadata: Vec<u8> = Vec::with_capacity(self.tx_log.len() * 16); + let mut pending_memtable_ops: Vec<(LogKey, TxEntry)> = vec![]; + + let mut metadata_buf = [0u8; 16]; + + debug!("Serializing tx_log to byte arrays"); + for tx_entry in &self.tx_log { + let record = match tx_entry { + TxEntry::Upsert { row: record } => record, + TxEntry::Delete { row: record } => record, + }; + + let serialized = record.serialize(); + let record_offset = initial_data_offset + serialized_data.len() as u64; + let record_length = serialized.len() as u64; + assert!(record_length > 0); + + serialized_data.extend(serialized); + + let metadata_pos = initial_metadata_offset + serialized_metadata.len() as u64; + let metadata_index = + (metadata_pos - METADATA_FILE_HEADER_SIZE as u64) / METADATA_ROW_LENGTH as u64; + + // Write the record metadata to the fixed-size metadata buffer + metadata_buf[..8].copy_from_slice(&record_offset.to_be_bytes()); + metadata_buf[8..].copy_from_slice(&record_length.to_be_bytes()); + + serialized_metadata.extend_from_slice(&metadata_buf); + + let log_key = LogKey::new(segment_num, metadata_index); + pending_memtable_ops.push((log_key, tx_entry.clone())); + } + + debug!("Writing serialized bytearrays to log files"); + self.active_data_file.write_all(&serialized_data)?; + self.active_metadata_file.write_all(&serialized_metadata)?; + + // Flush and sync data and metadata to disk + if self.config.write_durability == WriteDurability::Flush { + self.active_data_file.flush()?; + self.active_metadata_file.flush()?; + } else if self.config.write_durability == WriteDurability::FlushSync { + self.active_data_file.flush()?; + self.active_data_file.sync_all()?; + self.active_metadata_file.flush()?; + self.active_metadata_file.sync_all()?; + } + + debug!("Updating memtables"); + for (log_key, tx_entry) in pending_memtable_ops { + match tx_entry { + TxEntry::Upsert { row } => self.insert_row_to_memtables(log_key, row.values), + TxEntry::Delete { row } => self.remove_row_from_memtables(&row.values), + } + } + debug!("Commit done"); + + Ok(()) + } + + pub fn do_maintenance_tasks(&mut self) -> DBResult<()> { + ensure_active_metadata_is_valid(&self.data_dir_path, &mut self.active_metadata_file)?; + + let metadata_size = self.active_metadata_file.seek(SeekFrom::End(0))?; + if metadata_size >= self.config.segment_size as u64 { + self.rotate_and_compact()?; + } + + Ok(()) + } + + fn rotate_and_compact(&mut self) -> DBResult<()> { + debug!("Active log size exceeds threshold, starting rotation and compaction..."); + + let original_data_len = self.active_data_file.seek(SeekFrom::End(0))?; + + let active_target = fs::read_link(&self.data_dir_path.join(ACTIVE_SYMLINK_FILENAME))?; + let active_num = parse_segment_number(&active_target)?; + + debug!("Reading segment data into a BTreeMap"); + let mut pk_to_item_map: BTreeMap<&IndexableValue, &Row> = BTreeMap::new(); + let forward_read_items: Vec<(IndexableValue, Row)> = ForwardLogReader::new( + self.active_metadata_file.try_clone()?, + self.active_data_file.try_clone()?, + ) + .map(|item| { + ( + item.row.values[self.primary_key_index] + .as_indexable() + .expect("Primary key was not indexable"), + item.row, + ) + }) + .collect(); + + for (pk, record) in forward_read_items.iter() { + pk_to_item_map.insert(pk, record); + } + + debug!( + "Read {} records, out of which {} were unique", + forward_read_items.len(), + pk_to_item_map.len() + ); + + // Create a new log data file and write it + debug!("Opening new data file and writing compacted data"); + let (new_data_uuid, new_data_path) = create_segment_data_file(&self.data_dir_path)?; + let mut new_data_file = APPEND_MODE.open(&new_data_path)?; + + let mut pk_to_data_map = BTreeMap::new(); + let mut offset = 0u64; + for (pk, record) in pk_to_item_map.into_iter() { + let serialized = record.serialize(); + let len = serialized.len() as u64; + new_data_file.write_all(&serialized)?; + + pk_to_data_map.insert(pk, (offset, len)); + offset += len; + } + + // Sync the data file to disk. + // This is fine to do without consulting WriteDurability because this is a one-off + // operation that is not part of the normal write path. + new_data_file.flush()?; + new_data_file.sync_all()?; + + let final_data_len = new_data_file.seek(io::SeekFrom::End(0))?; + debug!( + "Wrote compacted data, reduced data size: {} -> {}", + original_data_len, final_data_len + ); + + // Create a new log metadata file and write it + debug!("Opening temp metadata file and writing pointers to compacted data file"); + let temp_metadata_file = tempfile::NamedTempFile::new()?; + let temp_metadata_path = temp_metadata_file.as_ref(); + let mut temp_metadata_file = WRITE_MODE.open(temp_metadata_path)?; + + let metadata_header = MetadataHeader { + version: 1, + uuid: new_data_uuid, + }; + + temp_metadata_file.write_all(&metadata_header.serialize())?; + + let mut metadata_buf = [0u8; 16]; + for (pk, _) in forward_read_items.iter() { + let (offset, len) = pk_to_data_map.get(&pk).unwrap(); + + metadata_buf[..8].copy_from_slice(&offset.to_be_bytes()); + metadata_buf[8..].copy_from_slice(&len.to_be_bytes()); + + temp_metadata_file.write_all(&metadata_buf)?; + } + + // Sync the metadata file to disk, see comment above about sync. + temp_metadata_file.flush()?; + temp_metadata_file.sync_all()?; + + debug!("Moving temporary files to their final locations"); + let new_data_path = &self.data_dir_path.join(new_data_uuid.to_string()); + let active_metadata_path = &self.data_dir_path.join(metadata_filename(active_num)); // overwrite active + + fs::rename(&temp_metadata_path, &active_metadata_path)?; + + debug!("Compaction complete, creating new segment"); + + let new_segment_num = active_num + 1; + let new_metadata_path = self.data_dir_path.join(metadata_filename(new_segment_num)); + let mut new_metadata_file = APPEND_MODE.clone().create(true).open(&new_metadata_path)?; + + let new_metadata_header = MetadataHeader { + version: 1, + uuid: new_data_uuid, + }; + + new_metadata_file.write_all(&new_metadata_header.serialize())?; + + set_active_segment(&self.data_dir_path, new_segment_num)?; + + self.active_metadata_file = APPEND_MODE.open(&new_metadata_path)?; + self.active_data_file = APPEND_MODE.open(&new_data_path)?; + + debug!( + "Active log file {} rotated and compacted, new segment: {}", + active_num, new_segment_num + ); + + Ok(()) + } + + #[inline] + pub fn with_exclusive_lock<A>( + &mut self, + f: impl FnOnce(&mut Self) -> DBResult<A>, + ) -> DBResult<A> { + // No need to acquire a lock if a transaction is already active + // because the lock is already held. + if !self.tx_active { + self.lock_manager.lock_exclusive()?; + } + let result = f(self); + if !self.tx_active { + self.lock_manager.unlock()?; + } + result + } + + #[inline] + pub fn with_shared_lock<A>(&mut self, f: impl FnOnce(&mut Self) -> DBResult<A>) -> DBResult<A> { + // No need to acquire a lock if a transaction is already active + // because the lock is already held. + if !self.tx_active { + self.lock_manager.lock_shared()?; + } + let result = f(self); + if !self.tx_active { + self.lock_manager.unlock()?; + } + + result + } +} + +#[cfg(test)] +mod tests { + use ctor::ctor; + use env_logger; + + use super::*; + + #[ctor] + fn init_logger() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + #[derive(Eq, PartialEq, Clone, Debug)] + enum Field { + Id, + Name, + } + + impl Into<String> for Field { + fn into(self) -> String { + match self { + Field::Id => "id".to_owned(), + Field::Name => "name".to_owned(), + } + } + } + + #[derive(PartialEq, Eq, Debug, Clone)] + struct TestInst2 { + id: i64, + name: String, + } + + impl From<TestInst2> for Vec<Value> { + fn from(inst: TestInst2) -> Self { + vec![Value::Int(inst.id), Value::String(inst.name)] + } + } + + impl From<Vec<Value>> for TestInst2 { + fn from(record: Vec<Value>) -> Self { + let mut it = record.into_iter(); + TestInst2 { + id: match it.next().unwrap() { + Value::Int(i) => i, + _ => panic!("Expected int"), + }, + name: match it.next().unwrap() { + Value::String(s) => s, + _ => panic!("Expected string"), + }, + } + } + } + + #[test] + fn test_memtable_insert_and_delete() { + let temp_dir = tempfile::tempdir().unwrap(); + let data_dir = temp_dir.path(); + + let capacity = 5; + let segment_size = capacity * 2 * 8 + METADATA_FILE_HEADER_SIZE; + + let mut db = DB::configure() + .data_dir(data_dir.to_str().unwrap()) + .fields(vec![Field::Id, Field::Name]) + .primary_key(Field::Id) + .secondary_keys(vec![Field::Name]) + .segment_size(segment_size) + .initialize() + .expect("Failed to create DB"); + + let engine = &mut db.engine; + + let inst = TestInst2 { + id: 0, + name: "foo".to_owned(), + }; + let id = IndexableValue::Int(0); + + assert_eq!(engine.primary_memtable.get(&id), None); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 0 + ); + engine.insert_row_to_memtables(LogKey::new(1, 0), inst.clone().into()); + assert_eq!(engine.primary_memtable.get(&id), Some(&LogKey::new(1, 0))); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 1 + ); + + engine.insert_row_to_memtables(LogKey::new(1, 1), inst.clone().into()); + assert_eq!(engine.primary_memtable.get(&id), Some(&LogKey::new(1, 1))); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 1 + ); + + engine.remove_row_from_memtables(&inst.into()); + assert_eq!(engine.primary_memtable.get(&id), None); + assert_eq!( + engine.secondary_memtables[0] + .find_by(&IndexableValue::String("foo".to_owned())) + .len(), + 0 + ); + } +} diff --git a/autere_db/src/lib.rs b/autere_db/src/lib.rs new file mode 100644 index 0000000..006056c --- /dev/null +++ b/autere_db/src/lib.rs @@ -0,0 +1,549 @@ +#[macro_use] +extern crate log; + +use once_cell::sync::Lazy; +use rust_decimal::Decimal; +use std::collections::BTreeMap; +use std::fmt::Debug; +use std::fmt::Display; +use std::fs::{self, metadata, File}; +use std::io::{self, Read, Seek, SeekFrom, Write}; +use std::ops::*; +use std::path::{Path, PathBuf}; +use std::thread; +use thiserror::Error; +use uuid::Uuid; + +#[macro_use] +mod common; +mod config; +mod engine; +mod lock; +mod log_reader_forward; +mod memtable_primary; +mod memtable_secondary; +mod record; +mod row; +mod schema; + +pub use common::{DBError, DBResult, OwnedBounds, QueryParams, Value, DEFAULT_QUERY_PARAMS}; +pub use config::{ReadConsistency, WriteDurability}; +pub use record::Record; +pub use schema::Schema; + +use common::*; +use config::*; +use engine::*; +use lock::*; +use log_reader_forward::*; +use memtable_primary::PrimaryMemtable; +use memtable_secondary::SecondaryMemtable; +use row::*; + +pub struct DB { + engine: Engine, +} + +impl DB { + /// Create a new database configuration builder. + pub fn configure() -> ConfigBuilder { + ConfigBuilder::new() + } + + fn initialize(config: Config) -> DBResult<DB> { + let engine = Engine::initialize(config)?; + Ok(DB { engine }) + } + + /// Insert a record into the database. If the primary key value already exists, + /// the existing record will be replaced by the supplied one. + pub fn upsert(&mut self, record: impl Into<Record>) -> DBResult<()> { + let row = Row { + values: record.into().into(), + tombstone: false, + }; + debug!("Upserting record: {:?}", row); + + self.engine + .with_exclusive_lock(move |engine| engine.upsert_record(row))?; + + Ok(()) + } + + /// Get a record by its primary index value. + /// E.g. `db.get(Value::Int(10))`. + pub fn get(&mut self, value: &Value) -> DBResult<Option<Record>> { + let tagged_rows = self.engine.with_shared_lock(|engine| { + engine.batch_find_by_records( + // TODO: This clone is only here to appease the borrow checker + &engine.config.primary_key.clone(), + std::iter::once(value), + &DEFAULT_QUERY_PARAMS, + ) + })?; + + assert!(tagged_rows.len() <= 1); + + Ok(tagged_rows + .into_iter() + .next() + .map(|(_, row)| Record::from(row))) + } + + /// Get a collection of records based on an indexed field value. + pub fn find_by(&mut self, field: impl AsRef<str>, value: &Value) -> DBResult<Vec<Record>> { + let tagged_rows = self.engine.with_shared_lock(|engine| { + engine.batch_find_by_records( + field.as_ref(), + std::iter::once(value), + &DEFAULT_QUERY_PARAMS, + ) + })?; + + Ok(tagged_rows + .into_iter() + .map(|(_, row)| Record::from(row)) + .collect()) + } + + /// Get a collection of records based on an indexed field value, with additional parameters. + pub fn find_by_with_params( + &mut self, + field: impl AsRef<str>, + value: &Value, + params: &QueryParams, + ) -> DBResult<Vec<Record>> { + let recs = self.engine.with_shared_lock(|engine| { + engine.batch_find_by_records(field.as_ref(), std::iter::once(value), params) + })?; + + Ok(recs.into_iter().map(|(_, row)| Record::from(row)).collect()) + } + + /// Get a collection of records based on a sequence of indexed field values. + /// Returns a vector of pairs where the first value is an index into the given sequence of values, + /// and the second value is the record. + pub fn batch_find_by( + &mut self, + field: impl Into<String>, + values: &[Value], + ) -> DBResult<Vec<(usize, Record)>> { + let recs = self.engine.with_shared_lock(|engine| { + engine.batch_find_by_records(&field.into(), values.iter(), &DEFAULT_QUERY_PARAMS) + })?; + + Ok(recs + .into_iter() + .map(|(tag, row)| (tag, Record::from(row))) + .collect()) + } + + /// Get a collection of records based on a sequence of indexed field values, with additional parameters. + /// Returns a vector of pairs where the first value is an index into the given sequence of values, + /// and the second value is the record. + pub fn batch_find_by_with_params( + &mut self, + field: impl AsRef<str>, + values: &[Value], + params: &QueryParams, + ) -> DBResult<Vec<(usize, Record)>> { + let recs = self.engine.with_shared_lock(|engine| { + engine.batch_find_by_records(field.as_ref(), values.iter(), params) + })?; + + Ok(recs + .into_iter() + .map(|(tag, row)| (tag, Record::from(row))) + .collect()) + } + + /// Get a collection of records based on a range of indexed field values. + /// This method can be used to run comparison-like queries, e.g. `field >= 10` + /// could be expressed as `db.range_by(Field::Id, 10..)`. + pub fn range_by<B: RangeBounds<Value>>( + &mut self, + field: impl AsRef<str>, + range: B, + ) -> DBResult<Vec<Record>> { + let recs = self.engine.with_shared_lock(|engine| { + engine.range_by_records(field.as_ref(), range, &DEFAULT_QUERY_PARAMS) + })?; + + Ok(recs.into_iter().map(|row| Record::from(row)).collect()) + } + + /// Get a collection of records based on a range of indexed field values, with additional parameters. + /// This method can be used to run comparison-like queries, e.g. `field >= 10` + /// could be expressed as `db.range_by(Field::Id, 10..)`. + pub fn range_by_with_params<B: RangeBounds<Value>>( + &mut self, + field: impl AsRef<str>, + range: B, + params: &QueryParams, + ) -> DBResult<Vec<Record>> { + let recs = self + .engine + .with_shared_lock(|engine| engine.range_by_records(field.as_ref(), range, params))?; + + Ok(recs.into_iter().map(|row| Record::from(row)).collect()) + } + + /// Delete records by a field value. + /// E.g. `db.delete_by(Field::Name, "John")`, assuming `Field` is the DB field type and `Field::Name` is secondary indexed. + /// Returns a vector of deleted records. If no records were deleted, the vector will be empty. + /// + /// Deletion is done by marking the record as a tombstone. The record will still be present in the log file, + /// but will be ignored by reads. Upon compaction, tombstoned records will be removed. + pub fn delete_by(&mut self, field: impl AsRef<str>, value: &Value) -> DBResult<Vec<Record>> { + let recs = self + .engine + .with_exclusive_lock(|engine| engine.delete_by_field(field.as_ref(), value))?; + + Ok(recs + .into_iter() + .map(|row| Record::from(row.values)) + .collect()) + } + + /// Delete record by primary key. + pub fn delete(&mut self, pk: &Value) -> DBResult<Option<Record>> { + let recs = self.engine.with_exclusive_lock(|engine| { + engine + // TODO: This clone is only here to appease the borrow checker + .delete_by_field(&engine.config.primary_key.clone(), pk) + })?; + + assert!(recs.len() <= 1); + + Ok(recs.into_iter().next().map(|row| Record::from(row.values))) + } + + /// Check if there are any pending tasks and do them. Tasks include: + /// - Rotating the active log file if it has reached capacity and compacting it. + /// + /// This function should be called periodically to ensure that the database remains in an optimal state. + /// Note that this function is synchronous and may block for a relatively long time. + /// You may call this function in a separate thread or process to avoid blocking the main thread. + /// However, the database will be exclusively locked, so all writes and reads will be blocked during the tasks. + pub fn do_maintenance_tasks(&mut self) -> DBResult<()> { + self.engine + .with_exclusive_lock(|engine| engine.do_maintenance_tasks()) + } + + /// Refresh the in-memory indexes from the log files. + /// This needs to only be called if the read consistency is set to `ReadConsistency::Eventual`. + pub fn refresh_indexes(&mut self) -> DBResult<()> { + self.engine + .with_exclusive_lock(|engine| engine.refresh_indexes()) + } + + /// Begin a transaction. This will acquire an exclusive lock on the database, + /// preventing other clients from using the database until the transaction is committed or rolled back. + pub fn tx_begin(&mut self) -> DBResult<()> { + if self.engine.tx_active { + return Err(DBError::TransactionError( + "Transaction already active".to_string(), + )); + } + + self.engine.lock_manager.lock_exclusive()?; + self.engine.tx_active = true; + Ok(()) + } + + /// Commit the active transaction. A transaction must be active, otherwise + /// a `DBError::TransactionError` will be returned. + pub fn tx_commit(&mut self) -> DBResult<()> { + if !self.engine.tx_active { + return Err(DBError::TransactionError( + "No active transaction to commit".to_string(), + )); + } + + self.engine.commit_transaction()?; + self.engine.tx_log.clear(); + self.engine.tx_active = false; + self.engine.lock_manager.unlock()?; + Ok(()) + } + + /// Rollback the active transaction. A transaction must be active, otherwise + /// a `DBError::TransactionError` will be returned. + pub fn tx_rollback(&mut self) -> DBResult<()> { + if !self.engine.tx_active { + return Err(DBError::TransactionError( + "No active transaction to roll back".to_string(), + )); + } + + self.engine.tx_log.clear(); + self.engine.tx_active = false; + self.engine.lock_manager.unlock()?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use ctor::ctor; + use env_logger; + + use super::*; + + #[ctor] + fn init_logger() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + #[derive(Eq, PartialEq, Clone, Debug)] + enum Field { + Id, + Name, + } + + impl Into<String> for Field { + fn into(self) -> String { + match self { + Field::Id => "id".to_string(), + Field::Name => "name".to_string(), + } + } + } + + struct TestInst1 { + id: i64, + } + + impl From<TestInst1> for Record { + fn from(inst: TestInst1) -> Self { + vec![Value::Int(inst.id)].into() + } + } + + impl From<Record> for TestInst1 { + fn from(record: Record) -> Self { + let mut it = record.into_iter(); + TestInst1 { + id: match it.next().unwrap() { + Value::Int(i) => i, + _ => panic!("Expected int"), + }, + } + } + } + + struct TestInst2 { + id: i64, + name: String, + } + + impl From<TestInst2> for Record { + fn from(inst: TestInst2) -> Self { + vec![Value::Int(inst.id), Value::String(inst.name)].into() + } + } + + impl From<Record> for TestInst2 { + fn from(record: Record) -> Self { + let mut it = record.into_iter(); + TestInst2 { + id: match it.next().unwrap() { + Value::Int(i) => i, + _ => panic!("Expected int"), + }, + name: match it.next().unwrap() { + Value::String(s) => s, + _ => panic!("Expected string"), + }, + } + } + } + + #[test] + fn test_compaction() { + let temp_dir = tempfile::tempdir().unwrap(); + let data_dir = temp_dir.path(); + + let capacity = 5; + let segment_size = capacity * 2 * 8 + METADATA_FILE_HEADER_SIZE; + + let mut db = DB::configure() + .data_dir(data_dir.to_str().unwrap()) + .fields(vec![Field::Id]) + .primary_key(Field::Id) + .segment_size(segment_size) + .initialize() + .expect("Failed to create DB"); + + // Insert records with same value until we reach the capacity + for _ in 0..capacity { + db.upsert(TestInst1 { id: 0 }) + .expect("Failed to insert record"); + } + + let mut segment1_file = READ_MODE.open(data_dir.join(metadata_filename(1))).unwrap(); + let segment1_metadata_size_original = segment1_file.seek(io::SeekFrom::End(0)).unwrap(); + + let segment1_header = read_metadata_header(&mut segment1_file).unwrap(); + let mut segment1_data_file = READ_MODE + .open(data_dir.join(segment1_header.uuid.to_string())) + .unwrap(); + let segment1_data_size_original = segment1_data_file.seek(io::SeekFrom::End(0)).unwrap(); + + // Rotate and compact + db.do_maintenance_tasks() + .expect("Failed to do maintenance tasks"); + + // Insert one extra with different value, this goes into another segment + db.upsert(TestInst1 { id: 1 }) + .expect("Failed to insert record"); + + // Check that rotation resulted in 2 segments + assert!(fs::exists(data_dir.join(metadata_filename(1))).unwrap()); + assert!(fs::exists(data_dir.join(metadata_filename(2))).unwrap()); + // Note negation here + assert!(!fs::exists(data_dir.join(metadata_filename(3))).unwrap()); + + // Check that the compacted metadata file has the same size + let mut segment1_metadata_file_compacted = + READ_MODE.open(data_dir.join(metadata_filename(1))).unwrap(); + let segment1_metadata_size_compacted = segment1_metadata_file_compacted + .seek(io::SeekFrom::End(0)) + .unwrap(); + assert_eq!( + segment1_metadata_size_compacted, + segment1_metadata_size_original + ); + + // Check that the compacted data file is smaller + let segment1_header_compacted = + read_metadata_header(&mut segment1_metadata_file_compacted).unwrap(); + let mut segment1_data_file_compacted = READ_MODE + .open(data_dir.join(segment1_header_compacted.uuid.to_string())) + .unwrap(); + let segment1_data_size_compacted = segment1_data_file_compacted + .seek(io::SeekFrom::End(0)) + .unwrap(); + assert!( + segment1_data_size_compacted < segment1_data_size_original, + "Original: {}, Compacted: {}", + segment1_data_size_original, + segment1_data_size_compacted + ); + + // Check that the records can be read + let inst0: TestInst1 = db + .get(&Value::Int(0 as i64)) + .expect("Failed to get record") + .expect("Record not found") + .into(); + + assert!(inst0.id == 0); + + let inst1: TestInst1 = db + .get(&Value::Int(1 as i64)) + .expect("Failed to get record") + .expect("Record not found") + .into(); + + assert!(inst1.id == 1); + } + + #[test] + fn test_repair() { + let temp_dir = tempfile::tempdir().unwrap(); + let data_dir = temp_dir.path(); + + let mut db = DB::configure() + .data_dir(data_dir.to_str().unwrap()) + .fields(vec![Field::Id]) + .primary_key(Field::Id) + .initialize() + .expect("Failed to create DB"); + + // Insert records + let n_recs: u64 = 100; + for i in 0..n_recs { + db.upsert(TestInst1 { id: i as i64 }) + .expect("Failed to insert record"); + } + + // Open the segment file and write garbage to it to simulate corruption + let segment_metadata_path = data_dir.join(metadata_filename(1)); + let mut file = APPEND_MODE + .open(&segment_metadata_path) + .expect("Failed to open file"); + + file.write_all(&[1, 0, 0, 0]) // A partially written integer value ([1] + some bytes) + .expect("Failed to write garbage"); + file.flush().unwrap(); + + let len = file.seek(SeekFrom::End(0)).expect("Failed to seek"); + assert_ne!(len, METADATA_FILE_HEADER_SIZE as u64 + n_recs * 16); + + // Try to refresh indexes, reading the file from beginning to end: should lead to error + db.refresh_indexes() + .expect_err("refresh_indexes should fail because of partial write"); + + // Trigger autorepair + db.do_maintenance_tasks() + .expect("Failed to run maintenance tasks"); + + // Try to refresh indexes, reading the file from beginning to end: should work now + db.refresh_indexes() + .expect("refresh_indexes should succeed"); + + // Reopen file and check that it has the correct size + let mut file = READ_MODE + .open(&segment_metadata_path) + .expect("Failed to open file"); + let len = file.seek(SeekFrom::End(0)).expect("Failed to seek"); + assert_eq!(len, METADATA_FILE_HEADER_SIZE as u64 + n_recs * 16); + } + + #[test] + fn test_memtables_updated_on_write() { + let temp_dir = tempfile::tempdir().unwrap(); + let data_dir = temp_dir.path(); + + let mut db = DB::configure() + .data_dir(data_dir.to_str().unwrap()) + .fields(vec![Field::Id, Field::Name]) + .primary_key(Field::Id) + .secondary_keys(vec![Field::Name]) + .initialize() + .expect("Failed to create DB"); + + // Check that the key is not indexed before write + assert_eq!( + db.engine.primary_memtable.get(&IndexableValue::Int(0)), + None + ); + assert_eq!( + db.engine.secondary_memtables[0] + .find_by(&IndexableValue::String("John".to_string())) + .len(), + 0 + ); + + // Insert record + db.upsert(TestInst2 { + id: 0, + name: "John".to_owned(), + }) + .expect("Failed to insert record"); + + // Check that the key is now indexed + let expected_log_key = LogKey::new(1, 0); + let expected_pk = IndexableValue::Int(0); + assert_eq!( + db.engine.primary_memtable.get(&expected_pk), + Some(&expected_log_key) + ); + let expected_vals = vec![&expected_log_key]; + let actual_vals = db.engine.secondary_memtables[0] + .find_by(&IndexableValue::String("John".to_string())) + .collect::<Vec<&LogKey>>(); + assert_eq!(actual_vals, expected_vals); + } +} diff --git a/autere_db/src/lock.rs b/autere_db/src/lock.rs new file mode 100644 index 0000000..f00429f --- /dev/null +++ b/autere_db/src/lock.rs @@ -0,0 +1,117 @@ +use super::*; + +pub struct LockManager { + lock_file: fs::File, + excl_lock_file: fs::File, + + state: LockState, +} + +#[derive(Debug, PartialEq, Eq)] +enum LockState { + NotLocked, + Shared, + Exclusive, +} + +impl LockManager { + pub fn new(data_dir_path: PathBuf) -> DBResult<LockManager> { + let lock_file = fs::File::create(data_dir_path.join(LOCK_FILENAME))?; + let excl_lock_file = fs::File::create(data_dir_path.join(EXCL_LOCK_REQ_FILENAME))?; + + Ok(LockManager { + lock_file, + excl_lock_file, + state: LockState::NotLocked, + }) + } + + fn is_exclusive_lock_requested(&self) -> DBResult<bool> { + // Attempt to acquire a shared lock on the lock request file + // If the file is already locked, return false + match fs2::FileExt::try_lock_shared(&self.excl_lock_file) { + Err(e) => { + if e.kind() == fs2::lock_contended_error().kind() { + return Ok(true); + } + return Err(DBError::IOError(e)); + } + + Ok(_) => { + fs2::FileExt::unlock(&self.excl_lock_file)?; + return Ok(false); + } + } + } + + pub fn lock_shared(&mut self) -> DBResult<()> { + if self.state == LockState::Shared { + return Err(DBError::LockRequestError( + "Already holding a shared lock".to_owned(), + )); + } else if self.state == LockState::Exclusive { + return Err(DBError::LockRequestError( + "Cannot acquire shared lock while holding an exclusive lock".to_owned(), + )); + } + + let mut timeout = 5; + loop { + if self.is_exclusive_lock_requested()? { + debug!( + "Exclusive lock requested, waiting for {}ms before requesting a shared lock again", + timeout + ); + thread::sleep(std::time::Duration::from_millis(timeout)); + timeout *= 2; + + if timeout > LOCK_WAIT_MAX_MS { + return Err(DBError::LockRequestError( + "Acquisition of shared lock timed out after {LOCK_WAIT_MAX_MS}".to_owned(), + )); + } + } else { + fs2::FileExt::lock_shared(&self.lock_file)?; + self.state = LockState::Shared; + return Ok(()); + } + } + } + + pub fn lock_exclusive(&mut self) -> DBResult<()> { + if self.state == LockState::Exclusive { + return Err(DBError::LockRequestError( + "Already holding an exclusive lock".to_owned(), + )); + } else if self.state == LockState::Shared { + return Err(DBError::LockRequestError( + "Cannot acquire exclusive lock while holding a shared lock".to_owned(), + )); + } + + // Create a lock on the exclusive lock request file to signal to readers that they should wait + // This will block until the lock is acquired + fs2::FileExt::lock_exclusive(&self.excl_lock_file)?; + + // Acquire an exclusive lock on the actual lock files + fs2::FileExt::lock_exclusive(&self.lock_file)?; + self.state = LockState::Exclusive; + + // Unlock the request file + fs2::FileExt::unlock(&self.excl_lock_file)?; + + Ok(()) + } + + pub fn unlock(&mut self) -> DBResult<()> { + if self.state == LockState::NotLocked { + return Err(DBError::LockRequestError( + "Not holding any locks".to_owned(), + )); + } + + fs2::FileExt::unlock(&self.lock_file)?; + self.state = LockState::NotLocked; + Ok(()) + } +} diff --git a/autere_db/src/log_reader_forward.rs b/autere_db/src/log_reader_forward.rs new file mode 100644 index 0000000..f3fc16c --- /dev/null +++ b/autere_db/src/log_reader_forward.rs @@ -0,0 +1,134 @@ +use super::*; + +pub struct ForwardLogReader { + metadata_reader: io::BufReader<fs::File>, + data_reader: io::BufReader<fs::File>, +} + +pub struct ForwardLogReaderItem { + pub row: Row, + pub index: u64, +} + +impl ForwardLogReader { + pub fn new(metadata_file: fs::File, data_file: fs::File) -> ForwardLogReader { + let mut ret = ForwardLogReader { + metadata_reader: io::BufReader::new(metadata_file), + data_reader: io::BufReader::new(data_file), + }; + + ret.metadata_reader + .seek(io::SeekFrom::Start(METADATA_FILE_HEADER_SIZE as u64)) + .expect("Seek failed"); + + ret + } + + pub fn new_with_index( + metadata_file: fs::File, + data_file: fs::File, + index: u64, + ) -> ForwardLogReader { + let mut ret = ForwardLogReader { + metadata_reader: io::BufReader::new(metadata_file), + data_reader: io::BufReader::new(data_file), + }; + + ret.metadata_reader + .seek(io::SeekFrom::Start( + METADATA_FILE_HEADER_SIZE as u64 + METADATA_ROW_LENGTH as u64 * index, + )) + .expect("Seek failed"); + + ret + } + + fn read_record(&mut self) -> Result<Option<ForwardLogReaderItem>, io::Error> { + loop { + let pos = self.metadata_reader.stream_position()?; + let index = (pos - METADATA_FILE_HEADER_SIZE as u64) / METADATA_ROW_LENGTH as u64; + + let mut metadata_entry_buf = vec![0; 16]; // 2x u64 + if let Err(e) = self.metadata_reader.read_exact(&mut metadata_entry_buf) { + if e.kind() == io::ErrorKind::UnexpectedEof { + return Ok(None); + } else { + return Err(e); + } + } + + // First u64 is the offset of the record in the data file, second is the length of the record + let entry_offset = u64::from_be_bytes(metadata_entry_buf[0..8].try_into().unwrap()); + let entry_length = u64::from_be_bytes(metadata_entry_buf[8..16].try_into().unwrap()); + + if entry_offset == 0 && entry_length == 0 { + // This is an unused entry in the metadata file, skip + continue; + } + + // Use .seek_relative instead of .seek to avoid dropping the BufReader internal buffer when + // the seek distance is small + let seek_distance = entry_offset as i64 - self.data_reader.stream_position()? as i64; + self.data_reader.seek_relative(seek_distance)?; + + let mut result_buf = vec![0; entry_length as usize]; + self.data_reader.read_exact(&mut result_buf)?; + + let row = Row::deserialize(&result_buf); + return Ok(Some(ForwardLogReaderItem { row, index })); + } + } +} + +impl Iterator for ForwardLogReader { + type Item = ForwardLogReaderItem; + + fn next(&mut self) -> Option<Self::Item> { + self.read_record().unwrap_or_else(|err| { + panic!("Error reading record: {:?}", err); + }) + } +} + +#[cfg(test)] +mod tests { + use ctor::ctor; + use env_logger; + + use super::*; + + #[ctor] + fn init_logger() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + const TEST_RESOURCES_DIR: &str = "tests/resources"; + + #[test] + fn test_forward_log_reader_fixture_db1() { + let metadata_path = Path::new(TEST_RESOURCES_DIR).join("test_metadata_1"); + let data_path = Path::new(TEST_RESOURCES_DIR).join("test_data_1"); + let metadata_file = fs::OpenOptions::new() + .read(true) + .open(&metadata_path) + .expect("Failed to open metadata file"); + let data_file = fs::OpenOptions::new() + .read(true) + .open(&data_path) + .expect("Failed to open data file"); + + let mut forward_log_reader = ForwardLogReader::new(metadata_file, data_file); + + // There are two records in the log with "schema" with one field: Bytes + + let ForwardLogReaderItem { row, index: _ } = forward_log_reader + .next() + .expect("Failed to read the first record"); + assert!(match &row.values[..] { + [Value::Bytes(bytes)] => bytes.len() == 256, + _ => false, + }); + + assert!(forward_log_reader.next().is_none()); + } +} diff --git a/autere_db/src/memtable_primary.rs b/autere_db/src/memtable_primary.rs new file mode 100644 index 0000000..573592e --- /dev/null +++ b/autere_db/src/memtable_primary.rs @@ -0,0 +1,40 @@ +use super::*; +use std::collections::BTreeMap; + +pub struct PrimaryMemtable { + /// Map of records indexed by key. Used as a shared heap of records + /// for all secondary memtables also. Secondary memtables store an + /// IndexableValue as their record value, which is used to get + /// the actual record from the primary memtable `records` map. + /// + /// Note: it must be invariant that all memtables (primary and secondary) + /// contain the same keys. + records: BTreeMap<IndexableValue, LogKey>, +} + +impl PrimaryMemtable { + pub fn new() -> PrimaryMemtable { + PrimaryMemtable { + records: BTreeMap::new(), + } + } + + pub fn set(&mut self, key: IndexableValue, value: LogKey) { + self.records.insert(key, value); + } + + pub fn get(&self, key: &IndexableValue) -> Option<&LogKey> { + self.records.get(key) + } + + pub fn remove(&mut self, key: &IndexableValue) -> Option<LogKey> { + self.records.remove(key) + } + + pub fn range<B: RangeBounds<IndexableValue>>(&self, range: B) -> Vec<&LogKey> { + self.records + .range(range) + .map(|(_, log_key)| log_key) + .collect() + } +} diff --git a/autere_db/src/memtable_secondary.rs b/autere_db/src/memtable_secondary.rs new file mode 100644 index 0000000..e1ca835 --- /dev/null +++ b/autere_db/src/memtable_secondary.rs @@ -0,0 +1,66 @@ +use once_cell::sync::Lazy; + +use super::*; +use std::collections::{btree_map::Values, BTreeMap}; + +pub struct SecondaryMemtable { + /// A 2-layer map of records indexed by SK => PK => LogKey. + /// The PK information is required to tell two records apart. + records: BTreeMap<IndexableValue, LogKeyMap>, +} + +static EMPTY_MAP: Lazy<BTreeMap<IndexableValue, LogKey>> = Lazy::new(|| BTreeMap::new()); + +impl SecondaryMemtable { + pub fn new() -> SecondaryMemtable { + SecondaryMemtable { + records: BTreeMap::new(), + } + } + + pub fn set(&mut self, pk: IndexableValue, sk: IndexableValue, value: LogKey) { + match self.records.get_mut(&sk) { + Some(map) => { + map.insert(pk, value); + } + None => { + self.records + .insert(sk, LogKeyMap::new_with_initial(pk, value)); + } + }; + } + + pub fn find_by(&self, key: &IndexableValue) -> Values<IndexableValue, LogKey> { + match self.records.get(key) { + Some(set) => set.log_keys(), + None => EMPTY_MAP.values(), + } + } + + // Remove a single mapping associated with the given PK and SK. Returns `true` + // if the log key existed and was removed, `false` otherwise. + pub fn remove(&mut self, pk: &IndexableValue, sk: &IndexableValue) -> bool { + let map = match self.records.get_mut(sk) { + Some(set) => set, + None => return false, + }; + if map.len() == 1 && map.contains_pk(pk) { + self.records.remove(sk); + true + } else { + return match map.remove_pk(pk) { + Ok(_) => true, + Err(LogKeyMapError::NotFoundError) => false, + Err(e) => panic!("{:?}", e), + }; + } + } + + pub fn range<B: RangeBounds<IndexableValue>>(&self, range: B) -> Vec<&LogKey> { + let mut keys = Vec::new(); + for (_, map) in self.records.range(range) { + keys.extend(map.log_keys()); + } + keys + } +} diff --git a/autere_db/src/record.rs b/autere_db/src/record.rs new file mode 100644 index 0000000..b349853 --- /dev/null +++ b/autere_db/src/record.rs @@ -0,0 +1,38 @@ +use super::*; + +pub struct Record { + values: Vec<Value>, +} + +impl Record { + pub fn values(&self) -> &[Value] { + &self.values + } +} + +impl IntoIterator for Record { + type Item = Value; + type IntoIter = std::vec::IntoIter<Self::Item>; + + fn into_iter(self) -> Self::IntoIter { + self.values.into_iter() + } +} + +impl From<Vec<Value>> for Record { + fn from(values: Vec<Value>) -> Self { + Record { values } + } +} + +impl From<Record> for Vec<Value> { + fn from(record: Record) -> Self { + record.values + } +} + +impl From<Row> for Record { + fn from(row: Row) -> Self { + Record { values: row.values } + } +} diff --git a/autere_db/src/row.rs b/autere_db/src/row.rs new file mode 100644 index 0000000..d8140d2 --- /dev/null +++ b/autere_db/src/row.rs @@ -0,0 +1,70 @@ +use super::*; + +#[derive(Debug, Clone)] +pub struct Row { + pub values: Vec<Value>, + pub tombstone: bool, +} + +impl Row { + pub fn serialize(&self) -> Vec<u8> { + let mut bytes = Vec::new(); + + if self.tombstone { + bytes.extend(&[B_TOMBSTONE]); + } else { + bytes.extend(&[B_LIVE]); + } + + for value in &self.values { + bytes.extend(value.serialize()); + } + bytes + } + + pub fn deserialize(bytes: &[u8]) -> Row { + assert!(bytes.len() > 0); + + let mut values = Vec::new(); + + let tombstone = bytes[0] == B_TOMBSTONE; + + let mut start = 1; + while start < bytes.len() { + let (rv, consumed) = Value::deserialize(&bytes[start..]); + values.push(rv); + start += consumed; + } + Row { values, tombstone } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_record_serialize_deserialize() { + let record = Row { + values: vec![ + Value::Int(1), + Value::String("hello".to_string()), + Value::Bytes(vec![0, 1, 2, 3]), + ], + tombstone: true, + }; + + let serialized = record.serialize(); + let deserialized = Row::deserialize(&serialized); + let reserialized = deserialized.serialize(); + + assert_eq!(serialized.len(), reserialized.len()); + assert_eq!(record.values, deserialized.values); + } +} + +#[derive(Clone, Debug)] +pub enum TxEntry { + Upsert { row: Row }, + Delete { row: Row }, +} diff --git a/autere_db/src/schema.rs b/autere_db/src/schema.rs new file mode 100644 index 0000000..98d0df7 --- /dev/null +++ b/autere_db/src/schema.rs @@ -0,0 +1,7 @@ +use super::*; + +pub struct Schema { + pub fields: Vec<String>, + pub primary_key: String, + pub secondary_keys: Vec<String>, +} |