mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-12-27 07:00:05 +01:00
implement document push
This commit is contained in:
mpostma
parent
6bcc302950
commit
658166c05e
@ -1,15 +1,20 @@
|
||||
use uuid::Uuid;
|
||||
use std::fs::{File, create_dir_all};
|
||||
use std::path::{PathBuf, Path};
|
||||
use chrono::Utc;
|
||||
use tokio::sync::{mpsc, oneshot, RwLock};
|
||||
use thiserror::Error;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use milli::Index;
|
||||
use std::collections::hash_map::Entry;
|
||||
use std::fs::create_dir_all;
|
||||
|
||||
use chrono::Utc;
|
||||
use heed::EnvOpenOptions;
|
||||
use milli::Index;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::hash_map::Entry;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::{mpsc, oneshot, RwLock};
|
||||
use uuid::Uuid;
|
||||
use log::info;
|
||||
|
||||
use super::update_handler::UpdateHandler;
|
||||
use crate::index_controller::{IndexMetadata, UpdateMeta, updates::{Processed, Failed, Processing}, UpdateResult as UResult};
|
||||
use crate::option::IndexerOpts;
|
||||
|
||||
pub type Result<T> = std::result::Result<T, IndexError>;
|
||||
type AsyncMap<K, V> = Arc<RwLock<HashMap<K, V>>>;
|
||||
@ -22,6 +27,7 @@ enum IndexMsg {
|
||||
|
||||
struct IndexActor<S> {
|
||||
inbox: mpsc::Receiver<IndexMsg>,
|
||||
update_handler: Arc<UpdateHandler>,
|
||||
store: S,
|
||||
}
|
||||
|
||||
@ -36,18 +42,22 @@ pub enum IndexError {
|
||||
#[async_trait::async_trait]
|
||||
trait IndexStore {
|
||||
async fn create_index(&self, uuid: Uuid, primary_key: Option<String>) -> Result<IndexMetadata>;
|
||||
async fn get_or_create(&self, uuid: Uuid) -> Result<Arc<Index>>;
|
||||
}
|
||||
|
||||
impl<S: IndexStore> IndexActor<S> {
|
||||
impl<S: IndexStore + Sync + Send> IndexActor<S> {
|
||||
fn new(inbox: mpsc::Receiver<IndexMsg>, store: S) -> Self {
|
||||
Self { inbox, store }
|
||||
let options = IndexerOpts::default();
|
||||
let update_handler = UpdateHandler::new(&options).unwrap();
|
||||
let update_handler = Arc::new(update_handler);
|
||||
Self { inbox, store, update_handler }
|
||||
}
|
||||
|
||||
async fn run(mut self) {
|
||||
loop {
|
||||
match self.inbox.recv().await {
|
||||
Some(IndexMsg::CreateIndex { uuid, primary_key, ret }) => self.handle_create_index(uuid, primary_key, ret).await,
|
||||
Some(IndexMsg::Update { ret, meta, data }) => self.handle_update().await,
|
||||
Some(IndexMsg::Update { ret, meta, data }) => self.handle_update(meta, data, ret).await,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
@ -58,8 +68,14 @@ impl<S: IndexStore> IndexActor<S> {
|
||||
let _ = ret.send(result);
|
||||
}
|
||||
|
||||
async fn handle_update(&self) {
|
||||
println!("processing update!!!");
|
||||
async fn handle_update(&self, meta: Processing<UpdateMeta>, data: File, ret: oneshot::Sender<UpdateResult>) {
|
||||
info!("processing update");
|
||||
let uuid = meta.index_uuid().clone();
|
||||
let index = self.store.get_or_create(uuid).await.unwrap();
|
||||
let update_handler = self.update_handler.clone();
|
||||
let result = tokio::task::spawn_blocking(move || update_handler.handle_update(meta, data, index.as_ref())).await;
|
||||
let result = result.unwrap();
|
||||
let _ = ret.send(result);
|
||||
}
|
||||
}
|
||||
|
||||
@ -96,7 +112,7 @@ impl IndexActorHandle {
|
||||
struct MapIndexStore {
|
||||
root: PathBuf,
|
||||
meta_store: AsyncMap<Uuid, IndexMetadata>,
|
||||
index_store: AsyncMap<Uuid, Index>,
|
||||
index_store: AsyncMap<Uuid, Arc<Index>>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
@ -126,10 +142,26 @@ impl IndexStore for MapIndexStore {
|
||||
Ok(index)
|
||||
}).await.expect("thread died");
|
||||
|
||||
self.index_store.write().await.insert(meta.uuid.clone(), index?);
|
||||
self.index_store.write().await.insert(meta.uuid.clone(), Arc::new(index?));
|
||||
|
||||
Ok(meta)
|
||||
}
|
||||
|
||||
async fn get_or_create(&self, uuid: Uuid) -> Result<Arc<Index>> {
|
||||
match self.index_store.write().await.entry(uuid.clone()) {
|
||||
Entry::Vacant(entry) => {
|
||||
match self.meta_store.write().await.entry(uuid.clone()) {
|
||||
Entry::Vacant(_) => {
|
||||
todo!()
|
||||
}
|
||||
Entry::Occupied(entry) => {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
}
|
||||
Entry::Occupied(entry) => Ok(entry.get().clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl MapIndexStore {
|
||||
|
||||
@ -2,6 +2,7 @@ mod index_actor;
|
||||
mod update_actor;
|
||||
mod uuid_resolver;
|
||||
mod update_store;
|
||||
mod update_handler;
|
||||
|
||||
use tokio::sync::oneshot;
|
||||
use super::IndexController;
|
||||
@ -10,7 +11,6 @@ use super::IndexMetadata;
|
||||
use tokio::fs::File;
|
||||
use super::UpdateMeta;
|
||||
|
||||
|
||||
pub struct ActorIndexController {
|
||||
uuid_resolver: uuid_resolver::UuidResolverHandle,
|
||||
index_handle: index_actor::IndexActorHandle,
|
||||
|
||||
@ -53,11 +53,11 @@ impl UpdateActor {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_update(&self, _uuid: Uuid, meta: UpdateMeta, payload: Option<File>, ret: oneshot::Sender<Result<UpdateStatus>>) {
|
||||
async fn handle_update(&self, uuid: Uuid, meta: UpdateMeta, payload: Option<File>, ret: oneshot::Sender<Result<UpdateStatus>>) {
|
||||
let mut buf = Vec::new();
|
||||
let mut payload = payload.unwrap();
|
||||
payload.read_to_end(&mut buf).await.unwrap();
|
||||
let result = self.store.register_update(meta, &buf).unwrap();
|
||||
let result = self.store.register_update(meta, &buf, uuid).unwrap();
|
||||
let _ = ret.send(Ok(UpdateStatus::Pending(result)));
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,255 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
use flate2::read::GzDecoder;
|
||||
use grenad::CompressionType;
|
||||
use log::info;
|
||||
use milli::update::{IndexDocumentsMethod, UpdateBuilder, UpdateFormat};
|
||||
use milli::Index;
|
||||
use rayon::ThreadPool;
|
||||
|
||||
use super::update_store::HandleUpdate;
|
||||
use crate::index_controller::updates::{Failed, Processed, Processing};
|
||||
use crate::index_controller::{Facets, Settings, UpdateMeta, UpdateResult};
|
||||
use crate::option::IndexerOpts;
|
||||
|
||||
pub struct UpdateHandler {
|
||||
index: Arc<Index>,
|
||||
max_nb_chunks: Option<usize>,
|
||||
chunk_compression_level: Option<u32>,
|
||||
thread_pool: Arc<ThreadPool>,
|
||||
log_frequency: usize,
|
||||
max_memory: usize,
|
||||
linked_hash_map_size: usize,
|
||||
chunk_compression_type: CompressionType,
|
||||
chunk_fusing_shrink_size: u64,
|
||||
}
|
||||
|
||||
impl UpdateHandler {
|
||||
pub fn new(
|
||||
opt: &IndexerOpts,
|
||||
index: Arc<Index>,
|
||||
thread_pool: Arc<ThreadPool>,
|
||||
) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
index,
|
||||
max_nb_chunks: opt.max_nb_chunks,
|
||||
chunk_compression_level: opt.chunk_compression_level,
|
||||
thread_pool,
|
||||
log_frequency: opt.log_every_n,
|
||||
max_memory: opt.max_memory.get_bytes() as usize,
|
||||
linked_hash_map_size: opt.linked_hash_map_size,
|
||||
chunk_compression_type: opt.chunk_compression_type,
|
||||
chunk_fusing_shrink_size: opt.chunk_fusing_shrink_size.get_bytes(),
|
||||
})
|
||||
}
|
||||
|
||||
fn update_buidler(&self, update_id: u64) -> UpdateBuilder {
|
||||
// We prepare the update by using the update builder.
|
||||
let mut update_builder = UpdateBuilder::new(update_id);
|
||||
if let Some(max_nb_chunks) = self.max_nb_chunks {
|
||||
update_builder.max_nb_chunks(max_nb_chunks);
|
||||
}
|
||||
if let Some(chunk_compression_level) = self.chunk_compression_level {
|
||||
update_builder.chunk_compression_level(chunk_compression_level);
|
||||
}
|
||||
update_builder.thread_pool(&self.thread_pool);
|
||||
update_builder.log_every_n(self.log_frequency);
|
||||
update_builder.max_memory(self.max_memory);
|
||||
update_builder.linked_hash_map_size(self.linked_hash_map_size);
|
||||
update_builder.chunk_compression_type(self.chunk_compression_type);
|
||||
update_builder.chunk_fusing_shrink_size(self.chunk_fusing_shrink_size);
|
||||
update_builder
|
||||
}
|
||||
|
||||
fn update_documents(
|
||||
&self,
|
||||
format: UpdateFormat,
|
||||
method: IndexDocumentsMethod,
|
||||
content: &[u8],
|
||||
update_builder: UpdateBuilder,
|
||||
primary_key: Option<&str>,
|
||||
) -> anyhow::Result<UpdateResult> {
|
||||
// We must use the write transaction of the update here.
|
||||
let mut wtxn = self.index.write_txn()?;
|
||||
|
||||
// Set the primary key if not set already, ignore if already set.
|
||||
match (self.index.primary_key(&wtxn)?, primary_key) {
|
||||
(None, Some(ref primary_key)) => {
|
||||
self.index.put_primary_key(&mut wtxn, primary_key)?;
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
|
||||
let mut builder = update_builder.index_documents(&mut wtxn, &self.index);
|
||||
builder.update_format(format);
|
||||
builder.index_documents_method(method);
|
||||
|
||||
let gzipped = true;
|
||||
let reader = if gzipped && !content.is_empty() {
|
||||
Box::new(GzDecoder::new(content))
|
||||
} else {
|
||||
Box::new(content) as Box<dyn io::Read>
|
||||
};
|
||||
|
||||
let result = builder.execute(reader, |indexing_step, update_id| {
|
||||
info!("update {}: {:?}", update_id, indexing_step)
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(addition_result) => wtxn
|
||||
.commit()
|
||||
.and(Ok(UpdateResult::DocumentsAddition(addition_result)))
|
||||
.map_err(Into::into),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_documents(&self, update_builder: UpdateBuilder) -> anyhow::Result<UpdateResult> {
|
||||
// We must use the write transaction of the update here.
|
||||
let mut wtxn = self.index.write_txn()?;
|
||||
let builder = update_builder.clear_documents(&mut wtxn, &self.index);
|
||||
|
||||
match builder.execute() {
|
||||
Ok(_count) => wtxn
|
||||
.commit()
|
||||
.and(Ok(UpdateResult::Other))
|
||||
.map_err(Into::into),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn update_settings(
|
||||
&self,
|
||||
settings: &Settings,
|
||||
update_builder: UpdateBuilder,
|
||||
) -> anyhow::Result<UpdateResult> {
|
||||
// We must use the write transaction of the update here.
|
||||
let mut wtxn = self.index.write_txn()?;
|
||||
let mut builder = update_builder.settings(&mut wtxn, &self.index);
|
||||
|
||||
// We transpose the settings JSON struct into a real setting update.
|
||||
if let Some(ref names) = settings.searchable_attributes {
|
||||
match names {
|
||||
Some(names) => builder.set_searchable_fields(names.clone()),
|
||||
None => builder.reset_searchable_fields(),
|
||||
}
|
||||
}
|
||||
|
||||
// We transpose the settings JSON struct into a real setting update.
|
||||
if let Some(ref names) = settings.displayed_attributes {
|
||||
match names {
|
||||
Some(names) => builder.set_displayed_fields(names.clone()),
|
||||
None => builder.reset_displayed_fields(),
|
||||
}
|
||||
}
|
||||
|
||||
// We transpose the settings JSON struct into a real setting update.
|
||||
if let Some(ref facet_types) = settings.faceted_attributes {
|
||||
let facet_types = facet_types.clone().unwrap_or_else(|| HashMap::new());
|
||||
builder.set_faceted_fields(facet_types);
|
||||
}
|
||||
|
||||
// We transpose the settings JSON struct into a real setting update.
|
||||
if let Some(ref criteria) = settings.criteria {
|
||||
match criteria {
|
||||
Some(criteria) => builder.set_criteria(criteria.clone()),
|
||||
None => builder.reset_criteria(),
|
||||
}
|
||||
}
|
||||
|
||||
let result = builder
|
||||
.execute(|indexing_step, update_id| info!("update {}: {:?}", update_id, indexing_step));
|
||||
|
||||
match result {
|
||||
Ok(()) => wtxn
|
||||
.commit()
|
||||
.and(Ok(UpdateResult::Other))
|
||||
.map_err(Into::into),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn update_facets(
|
||||
&self,
|
||||
levels: &Facets,
|
||||
update_builder: UpdateBuilder,
|
||||
) -> anyhow::Result<UpdateResult> {
|
||||
// We must use the write transaction of the update here.
|
||||
let mut wtxn = self.index.write_txn()?;
|
||||
let mut builder = update_builder.facets(&mut wtxn, &self.index);
|
||||
if let Some(value) = levels.level_group_size {
|
||||
builder.level_group_size(value);
|
||||
}
|
||||
if let Some(value) = levels.min_level_size {
|
||||
builder.min_level_size(value);
|
||||
}
|
||||
match builder.execute() {
|
||||
Ok(()) => wtxn
|
||||
.commit()
|
||||
.and(Ok(UpdateResult::Other))
|
||||
.map_err(Into::into),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn delete_documents(
|
||||
&self,
|
||||
document_ids: &[u8],
|
||||
update_builder: UpdateBuilder,
|
||||
) -> anyhow::Result<UpdateResult> {
|
||||
let ids: Vec<String> = serde_json::from_slice(document_ids)?;
|
||||
let mut txn = self.index.write_txn()?;
|
||||
let mut builder = update_builder.delete_documents(&mut txn, &self.index)?;
|
||||
|
||||
// We ignore unexisting document ids
|
||||
ids.iter().for_each(|id| { builder.delete_external_id(id); });
|
||||
|
||||
match builder.execute() {
|
||||
Ok(deleted) => txn
|
||||
.commit()
|
||||
.and(Ok(UpdateResult::DocumentDeletion { deleted }))
|
||||
.map_err(Into::into),
|
||||
Err(e) => Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HandleUpdate<UpdateMeta, UpdateResult, String> for UpdateHandler {
|
||||
fn handle_update(
|
||||
&mut self,
|
||||
meta: Processing<UpdateMeta>,
|
||||
content: &[u8],
|
||||
) -> Result<Processed<UpdateMeta, UpdateResult>, Failed<UpdateMeta, String>> {
|
||||
use UpdateMeta::*;
|
||||
|
||||
let update_id = meta.id();
|
||||
|
||||
let update_builder = self.update_buidler(update_id);
|
||||
|
||||
let result = match meta.meta() {
|
||||
DocumentsAddition {
|
||||
method,
|
||||
format,
|
||||
primary_key,
|
||||
} => self.update_documents(
|
||||
*format,
|
||||
*method,
|
||||
content,
|
||||
update_builder,
|
||||
primary_key.as_deref(),
|
||||
),
|
||||
ClearDocuments => self.clear_documents(update_builder),
|
||||
DeleteDocuments => self.delete_documents(content, update_builder),
|
||||
Settings(settings) => self.update_settings(settings, update_builder),
|
||||
Facets(levels) => self.update_facets(levels, update_builder),
|
||||
};
|
||||
|
||||
match result {
|
||||
Ok(result) => Ok(meta.process(result)),
|
||||
Err(e) => Err(meta.fail(e.to_string())),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,5 +1,6 @@
|
||||
use chrono::{Utc, DateTime};
|
||||
use serde::{Serialize, Deserialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Clone)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@ -7,14 +8,16 @@ pub struct Pending<M> {
|
||||
pub update_id: u64,
|
||||
pub meta: M,
|
||||
pub enqueued_at: DateTime<Utc>,
|
||||
pub index_uuid: Uuid,
|
||||
}
|
||||
|
||||
impl<M> Pending<M> {
|
||||
pub fn new(meta: M, update_id: u64) -> Self {
|
||||
pub fn new(meta: M, update_id: u64, index_uuid: Uuid) -> Self {
|
||||
Self {
|
||||
enqueued_at: Utc::now(),
|
||||
meta,
|
||||
update_id,
|
||||
index_uuid,
|
||||
}
|
||||
}
|
||||
|
||||
@ -73,6 +76,10 @@ impl<M> Processing<M> {
|
||||
self.from.meta()
|
||||
}
|
||||
|
||||
pub fn index_uuid(&self) -> &Uuid {
|
||||
&self.from.index_uuid
|
||||
}
|
||||
|
||||
pub fn process<N>(self, meta: N) -> Processed<M, N> {
|
||||
Processed {
|
||||
success: meta,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user