use std::collections::HashMap; use std::convert::TryFrom; use std::path::{Path, PathBuf}; use std::sync::Arc; use milli::update::IndexerConfig; use tokio::fs; use tokio::sync::RwLock; use tokio::task::spawn_blocking; use uuid::Uuid; use super::error::{IndexResolverError, Result}; use crate::index::Index; use crate::options::IndexerOpts; type AsyncMap = Arc>>; #[async_trait::async_trait] #[cfg_attr(test, mockall::automock)] pub trait IndexStore { async fn create(&self, uuid: Uuid) -> Result; async fn get(&self, uuid: Uuid) -> Result>; async fn delete(&self, uuid: Uuid) -> Result>; } pub struct MapIndexStore { index_store: AsyncMap, path: PathBuf, index_size: usize, indexer_config: Arc, } impl MapIndexStore { pub fn new( path: impl AsRef, index_size: usize, indexer_opts: &IndexerOpts, ) -> anyhow::Result { let indexer_config = Arc::new(IndexerConfig::try_from(indexer_opts)?); let path = path.as_ref().join("indexes/"); let index_store = Arc::new(RwLock::new(HashMap::new())); Ok(Self { index_store, path, index_size, indexer_config, }) } } #[async_trait::async_trait] impl IndexStore for MapIndexStore { async fn create(&self, uuid: Uuid) -> Result { // We need to keep the lock until we are sure the db file has been opened correclty, to // ensure that another db is not created at the same time. let mut lock = self.index_store.write().await; if let Some(index) = lock.get(&uuid) { return Ok(index.clone()); } let path = self.path.join(format!("{}", uuid)); if path.exists() { return Err(IndexResolverError::UuidAlreadyExists(uuid)); } let index_size = self.index_size; let update_handler = self.indexer_config.clone(); let index = spawn_blocking(move || -> Result { let index = Index::open(path, index_size, uuid, update_handler)?; Ok(index) }) .await??; lock.insert(uuid, index.clone()); Ok(index) } async fn get(&self, uuid: Uuid) -> Result> { let guard = self.index_store.read().await; match guard.get(&uuid) { Some(index) => Ok(Some(index.clone())), None => { // drop the guard here so we can perform the write after without deadlocking; drop(guard); let path = self.path.join(format!("{}", uuid)); if !path.exists() { return Ok(None); } let index_size = self.index_size; let update_handler = self.indexer_config.clone(); let index = spawn_blocking(move || Index::open(path, index_size, uuid, update_handler)) .await??; self.index_store.write().await.insert(uuid, index.clone()); Ok(Some(index)) } } } async fn delete(&self, uuid: Uuid) -> Result> { let db_path = self.path.join(format!("{}", uuid)); fs::remove_dir_all(db_path).await?; let index = self.index_store.write().await.remove(&uuid); Ok(index) } }