Introduce the Database struct to manage indexes

This commit is contained in:
Clément Renault 2019-10-04 16:49:17 +02:00
parent effbbc7370
commit 851cc38216
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
3 changed files with 93 additions and 81 deletions

View File

@ -0,0 +1,87 @@
use std::collections::HashMap;
use std::io;
use std::path::Path;
use std::sync::{Arc, RwLock};
use crate::{store, Index, MResult};
pub struct Database {
pub rkv: Arc<RwLock<rkv::Rkv>>,
main_store: rkv::SingleStore,
indexes: RwLock<HashMap<String, Index>>,
}
impl Database {
pub fn open_or_create(path: impl AsRef<Path>) -> MResult<Database> {
let manager = rkv::Manager::singleton();
let mut rkv_write = manager.write().unwrap();
let rkv = rkv_write
.get_or_create(path.as_ref(), |path| {
let mut builder = rkv::Rkv::environment_builder();
builder.set_max_dbs(3000).set_map_size(10 * 1024 * 1024 * 1024); // 10GB
rkv::Rkv::from_env(path, builder)
})?;
drop(rkv_write);
let rkv_read = rkv.read().unwrap();
let create_options = rkv::store::Options::create();
let main_store = rkv_read.open_single("indexes", create_options)?;
// list all indexes that needs to be opened
let mut must_open = Vec::new();
let reader = rkv_read.read()?;
for result in main_store.iter_start(&reader)? {
let (key, _) = result?;
if let Ok(index_name) = std::str::from_utf8(key) {
must_open.push(index_name.to_owned());
}
}
drop(reader);
// open the previously aggregated indexes
let mut indexes = HashMap::new();
for index_name in must_open {
let index = store::open(&rkv_read, &index_name)?;
indexes.insert(index_name, index);
}
drop(rkv_read);
Ok(Database { rkv, main_store, indexes: RwLock::new(indexes) })
}
pub fn open_index(&self, name: impl Into<String>) -> MResult<Index> {
let indexes_lock = self.indexes.read().unwrap();
let name = name.into();
match indexes_lock.get(&name) {
Some(index) => Ok(*index),
None => {
drop(indexes_lock);
let rkv_lock = self.rkv.read().unwrap();
let index = store::create(&rkv_lock, &name)?;
let mut writer = rkv_lock.write()?;
let value = rkv::Value::Blob(&[]);
self.main_store.put(&mut writer, &name, &value)?;
{
let mut indexes_write = self.indexes.write().unwrap();
indexes_write.entry(name).or_insert(index);
}
writer.commit()?;
Ok(index)
},
}
}
pub fn indexes_names(&self) -> MResult<Vec<String>> {
let indexes = self.indexes.read().unwrap();
Ok(indexes.keys().cloned().collect())
}
}

View File

@ -1,4 +1,5 @@
mod automaton;
mod database;
mod error;
mod number;
mod query_builder;
@ -11,92 +12,14 @@ pub mod raw_indexer;
pub mod serde;
pub mod store;
pub use self::query_builder::QueryBuilder;
pub use self::raw_document::RawDocument;
pub use self::database::Database;
pub use self::error::{Error, MResult};
pub use self::number::{Number, ParseNumberError};
pub use self::query_builder::QueryBuilder;
pub use self::ranked_map::RankedMap;
pub use self::raw_document::RawDocument;
pub use self::store::Index;
use std::collections::HashMap;
use std::io;
use std::path::Path;
use std::sync::{Arc, RwLock};
pub struct Database {
rkv: Arc<RwLock<rkv::Rkv>>,
main_store: rkv::SingleStore,
indexes: RwLock<HashMap<String, Index>>,
}
impl Database {
pub fn open_or_create(path: impl AsRef<Path>) -> io::Result<Database> {
let manager = rkv::Manager::singleton();
let mut rkv_write = manager.write().unwrap();
let rkv = rkv_write
.get_or_create(path.as_ref(), |path| {
let mut builder = rkv::Rkv::environment_builder();
builder.set_max_dbs(3000).set_map_size(10 * 1024 * 1024 * 1024); // 10GB
rkv::Rkv::from_env(path, builder)
})
.unwrap();
drop(rkv_write);
let mut indexes = HashMap::new();
let main_store;
{
let rkv_read = rkv.read().unwrap();
main_store = rkv_read
.open_single("indexes", rkv::store::Options::create())
.unwrap();
let mut must_open = Vec::new();
let reader = rkv_read.read().unwrap();
for result in main_store.iter_start(&reader).unwrap() {
let (key, _) = result.unwrap();
if let Ok(index_name) = std::str::from_utf8(key) {
println!("{:?}", index_name);
must_open.push(index_name.to_owned());
}
}
drop(reader);
for index_name in must_open {
let index = store::open(&rkv_read, &index_name).unwrap();
indexes.insert(index_name, index);
}
}
Ok(Database { rkv, main_store, indexes: RwLock::new(indexes) })
}
pub fn open_index(&self, name: impl Into<String>) -> MResult<Index> {
let read = self.indexes.read().unwrap();
let name = name.into();
match read.get(&name) {
Some(index) => Ok(*index),
None => {
drop(read);
let rkv = self.rkv.read().unwrap();
let mut write = self.indexes.write().unwrap();
let index = store::create(&rkv, &name).unwrap();
let mut writer = rkv.write().unwrap();
let value = rkv::Value::Blob(&[]);
self.main_store.put(&mut writer, &name, &value).unwrap();
writer.commit().unwrap();
Ok(*write.entry(name.clone()).or_insert(index))
},
}
}
}
use zerocopy::{AsBytes, FromBytes};
use ::serde::{Serialize, Deserialize};

View File

@ -7,6 +7,8 @@ fn main() {
fs::create_dir_all(path).unwrap();
let database = Database::open_or_create(path).unwrap();
println!("{:?}", database.indexes_names());
let hello = database.open_index("hello").unwrap();
let hello1 = database.open_index("hello1").unwrap();
let hello2 = database.open_index("hello2").unwrap();