minisql/storage_engine/src/index.rs
2024-02-05 01:34:16 +01:00

290 lines
9 KiB
Rust

use std::path::PathBuf;
use tokio::fs::{File, OpenOptions};
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, BufWriter};
use std::collections::{BTreeMap, HashSet};
use std::hash::Hash;
use std::io::SeekFrom;
use crate::binary_coding::{decode, encode};
use bincode;
use bincode::{Decode, Encode};
use crate::error::{DecodeErrorKind, Error};
type Result<T> = std::result::Result<T, Error>;
#[derive(Debug)]
pub struct Index<K, V> {
file: File,
data: BTreeMap<K, HashSet<V>>,
}
#[derive(Debug)]
pub struct IndexHeader {}
impl<K, V> Index<K, V>
where
K: Encode + Decode + Ord,
V: Encode + Decode + Clone + Eq + Hash,
{
pub async fn new(file_name: PathBuf) -> Result<Index<K, V>> {
let file: File = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(file_name)
.await?;
let data = BTreeMap::new();
Ok(Index { file, data })
}
pub async fn connect(file_name: PathBuf) -> Result<Index<K, V>> {
let file: File = OpenOptions::new()
.read(true)
.write(true)
.open(file_name)
.await?;
let mut index = Index {
file,
data: BTreeMap::new(),
};
index.load_from_file().await?;
Ok(index)
}
pub async fn insert(&mut self, k: K, v: V) -> Result<()> {
self.append_to_file(&k, &v).await?;
self.data.entry(k).or_insert_with(HashSet::new).insert(v);
Ok(())
}
pub fn insert_desynced(&mut self, k: K, v: V) -> () {
self.data.entry(k).or_insert_with(HashSet::new).insert(v);
}
pub async fn lookup(&self, k: &K) -> Result<Option<HashSet<V>>> {
let hashset = self.data.get(k).cloned();
Ok(hashset)
}
pub async fn delete(&mut self, k: K, v: V) -> Result<()> {
self.data.entry(k).and_modify(|values| {
values.remove(&v);
});
self.sync_to_disk().await
}
pub async fn sync_to_disk(&mut self) -> Result<()> {
let mut writer = BufWriter::new(&mut self.file);
writer.seek(SeekFrom::Start(0)).await?;
let mut written: u64 = 0;
let mut encoded = Vec::new();
for (key, value) in &self.data {
for v in value {
encoded.clear();
encoded.extend(encode(key)?);
encoded.extend(encode(v)?);
writer.write(&encoded).await?;
written += encoded.len() as u64;
}
}
writer.flush().await?;
self.file.set_len(written).await?;
Ok(())
}
pub async fn reset(&mut self, data: BTreeMap<K, HashSet<V>>) -> Result<()> {
self.data = data;
self.sync_to_disk().await
}
async fn append_to_file(&mut self, key: &K, value: &V) -> Result<()> {
let mut encoded = Vec::new();
encoded.extend(encode(key)?);
encoded.extend(encode(value)?);
self.file.seek(SeekFrom::End(0)).await?;
self.file.write(&encoded).await?;
Ok(())
}
async fn load_from_file(&mut self) -> Result<()> {
let mut bytes = vec![];
self.file.seek(SeekFrom::Start(0)).await?;
self.file.read_to_end(&mut bytes).await?;
let mut cursor = 0;
while cursor < bytes.len() {
let (key, len) = decode(&bytes[cursor..])
.map_err(|e| Error::DecodeError(DecodeErrorKind::CorruptedData, e))?;
cursor += len;
let (value, len) = decode(&bytes[cursor..])
.map_err(|e| Error::DecodeError(DecodeErrorKind::CorruptedData, e))?;
cursor += len;
self.insert_desynced(key, value);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::fs::remove_file;
#[tokio::test]
async fn connect_to_new() {
let file_name = PathBuf::from("connect_to_new");
if file_name.exists() {
remove_file(&file_name).await.unwrap();
}
{
let index = Index::<u64, u64>::new(file_name.clone()).await.unwrap();
assert_eq!(index.data.len(), 0);
}
{
let index = Index::<u64, u64>::connect(file_name.clone()).await.unwrap();
assert_eq!(index.data.len(), 0);
}
remove_file(&file_name).await.unwrap();
}
#[tokio::test]
async fn inserting() {
let file_name = PathBuf::from("inserting");
if file_name.exists() {
remove_file(&file_name).await.unwrap();
}
{
let mut index = Index::<u64, u64>::new(file_name.clone()).await.unwrap();
index.insert(1, 2).await.unwrap();
index.insert(1, 3).await.unwrap();
index.insert(1, 4).await.unwrap();
index.insert(2, 3).await.unwrap();
index.insert(2, 4).await.unwrap();
index.insert(2, 5).await.unwrap();
assert_eq!(index.data.len(), 2);
assert_eq!(index.data.get(&1).unwrap().len(), 3);
assert_eq!(index.data.get(&2).unwrap().len(), 3);
}
{
let index = Index::<u64, u64>::connect(file_name.clone()).await.unwrap();
assert_eq!(index.data.len(), 2);
assert_eq!(index.data.get(&1).unwrap().len(), 3);
assert_eq!(index.data.get(&2).unwrap().len(), 3);
}
remove_file(&file_name).await.unwrap();
}
#[tokio::test]
async fn lookuping() {
let file_name = PathBuf::from("lookuping");
if file_name.exists() {
remove_file(&file_name).await.unwrap();
}
{
let mut index = Index::<u64, u64>::new(file_name.clone()).await.unwrap();
index.insert(1, 2).await.unwrap();
index.insert(1, 3).await.unwrap();
index.insert(1, 4).await.unwrap();
index.insert(2, 3).await.unwrap();
index.insert(2, 4).await.unwrap();
index.insert(2, 5).await.unwrap();
assert_eq!(index.lookup(&1).await.unwrap().unwrap().len(), 3);
assert_eq!(index.lookup(&2).await.unwrap().unwrap().len(), 3);
assert_eq!(index.lookup(&3).await.unwrap(), None);
let first = index.lookup(&1).await.unwrap().unwrap();
assert!(first.contains(&2));
assert!(first.contains(&3));
assert!(first.contains(&4));
let second = index.lookup(&2).await.unwrap().unwrap();
assert!(second.contains(&3));
assert!(second.contains(&4));
assert!(second.contains(&5));
}
{
let index = Index::<u64, u64>::connect(file_name.clone()).await.unwrap();
assert_eq!(index.lookup(&1).await.unwrap().unwrap().len(), 3);
assert_eq!(index.lookup(&2).await.unwrap().unwrap().len(), 3);
assert_eq!(index.lookup(&3).await.unwrap(), None);
let first = index.lookup(&1).await.unwrap().unwrap();
assert!(first.contains(&2));
assert!(first.contains(&3));
assert!(first.contains(&4));
let second = index.lookup(&2).await.unwrap().unwrap();
assert!(second.contains(&3));
assert!(second.contains(&4));
assert!(second.contains(&5));
}
remove_file(&file_name).await.unwrap();
}
#[tokio::test]
async fn deleting() {
let file_name = PathBuf::from("deleting");
if file_name.exists() {
remove_file(&file_name).await.unwrap();
}
{
let mut index = Index::<u64, u64>::new(file_name.clone()).await.unwrap();
index.insert(1, 2).await.unwrap();
index.insert(1, 3).await.unwrap();
index.insert(1, 4).await.unwrap();
index.insert(2, 3).await.unwrap();
index.insert(2, 4).await.unwrap();
index.insert(2, 5).await.unwrap();
assert!(index.lookup(&1).await.unwrap().unwrap().contains(&2));
index.delete(1, 2).await.unwrap();
assert!(!index.lookup(&1).await.unwrap().unwrap().contains(&2));
assert!(index.lookup(&2).await.unwrap().unwrap().contains(&3));
index.delete(2, 3).await.unwrap();
assert!(!index.lookup(&2).await.unwrap().unwrap().contains(&3));
}
{
let mut index = Index::<u64, u64>::connect(file_name.clone()).await.unwrap();
assert!(!index.lookup(&1).await.unwrap().unwrap().contains(&2));
assert!(!index.lookup(&2).await.unwrap().unwrap().contains(&3));
assert!(index.lookup(&1).await.unwrap().unwrap().contains(&3));
index.delete(1, 3).await.unwrap();
assert!(!index.lookup(&1).await.unwrap().unwrap().contains(&3));
assert!(index.lookup(&1).await.unwrap().unwrap().contains(&4));
assert!(index.lookup(&2).await.unwrap().unwrap().contains(&4));
assert!(index.lookup(&2).await.unwrap().unwrap().contains(&5));
}
remove_file(&file_name).await.unwrap();
}
}