Merge #1176

1176: fix race condition in  document addition r=Kerollmops a=MarinPostma

As described in #1160, there was a race condition when updating settings and adding documents simultaneously. This was due to the schema being updated and document addition being processed in two different transactions. This PR moves the schema update logic for the primary key in the same transaction as the document addition, while maintaining the input checks for the validity of the primary key in the http route, in order not to break the error reporting for the document addition route.

close #1160.

Co-authored-by: mpostma <postma.marin@protonmail.com>
Co-authored-by: marin <postma.marin@protonmail.com>

meilisearch-core/src/update/documents_addition.rs

@@ -23,6 +23,8 @@ pub struct DocumentsAddition<D> {
     updates_store: store::Updates,
     updates_results_store: store::UpdatesResults,
     updates_notifier: UpdateEventsEmitter,
+    // Whether the user explicitly set the primary key in the update
+    primary_key: Option<String>,
     documents: Vec<D>,
     is_partial: bool,
 }
@@ -39,6 +41,7 @@ impl<D> DocumentsAddition<D> {
             updates_notifier,
             documents: Vec::new(),
             is_partial: false,
+            primary_key: None,
         }
     }
 
@@ -53,9 +56,14 @@ impl<D> DocumentsAddition<D> {
             updates_notifier,
             documents: Vec::new(),
             is_partial: true,
+            primary_key: None,
         }
     }
 
+    pub fn set_primary_key(&mut self, primary_key: String) {
+        self.primary_key = Some(primary_key);
+    }
+
     pub fn update_document(&mut self, document: D) {
         self.documents.push(document);
     }
@@ -71,6 +79,7 @@ impl<D> DocumentsAddition<D> {
             self.updates_results_store,
             self.documents,
             self.is_partial,
+            self.primary_key,
         )?;
         Ok(update_id)
     }
@@ -88,6 +97,7 @@ pub fn push_documents_addition<D: serde::Serialize>(
     updates_results_store: store::UpdatesResults,
     addition: Vec<D>,
     is_partial: bool,
+    primary_key: Option<String>,
 ) -> MResult<u64> {
     let mut values = Vec::with_capacity(addition.len());
     for add in addition {
@@ -99,9 +109,9 @@ pub fn push_documents_addition<D: serde::Serialize>(
     let last_update_id = next_update_id(writer, updates_store, updates_results_store)?;
 
     let update = if is_partial {
-        Update::documents_partial(values)
+        Update::documents_partial(primary_key, values)
     } else {
-        Update::documents_addition(values)
+        Update::documents_addition(primary_key, values)
     };
 
     updates_store.put_update(writer, last_update_id, &update)?;
@@ -149,7 +159,8 @@ pub fn apply_addition(
     writer: &mut heed::RwTxn<MainT>,
     index: &store::Index,
     new_documents: Vec<IndexMap<String, Value>>,
-    partial: bool
+    partial: bool,
+    primary_key: Option<String>,
 ) -> MResult<()>
 {
     let mut schema = match index.main.schema(writer)? {
@@ -162,7 +173,14 @@ pub fn apply_addition(
     let internal_docids = index.main.internal_docids(writer)?;
     let mut available_ids = DiscoverIds::new(&internal_docids);
 
-    let primary_key = schema.primary_key().ok_or(Error::MissingPrimaryKey)?;
+    let primary_key = match schema.primary_key() {
+        Some(primary_key) => primary_key.to_string(),
+        None => {
+            let name = primary_key.ok_or(Error::MissingPrimaryKey)?;
+            schema.set_primary_key(&name)?;
+            name
+        }
+    };
 
     // 1. store documents ids for future deletion
     let mut documents_additions = HashMap::new();
@@ -275,16 +293,18 @@ pub fn apply_documents_partial_addition(
     writer: &mut heed::RwTxn<MainT>,
     index: &store::Index,
     new_documents: Vec<IndexMap<String, Value>>,
+    primary_key: Option<String>,
 ) -> MResult<()> {
-    apply_addition(writer, index, new_documents, true)
+    apply_addition(writer, index, new_documents, true, primary_key)
 }
 
 pub fn apply_documents_addition(
     writer: &mut heed::RwTxn<MainT>,
     index: &store::Index,
     new_documents: Vec<IndexMap<String, Value>>,
+    primary_key: Option<String>,
 ) -> MResult<()> {
-    apply_addition(writer, index, new_documents, false)
+    apply_addition(writer, index, new_documents, false, primary_key)
 }
 
 pub fn reindex_all_documents(writer: &mut heed::RwTxn<MainT>, index: &store::Index) -> MResult<()> {

meilisearch-core/src/update/mod.rs

@@ -52,16 +52,16 @@ impl Update {
         }
     }
 
-    fn documents_addition(documents: Vec<IndexMap<String, Value>>) -> Update {
+    fn documents_addition(primary_key: Option<String>, documents: Vec<IndexMap<String, Value>>) -> Update {
         Update {
-            data: UpdateData::DocumentsAddition(documents),
+            data: UpdateData::DocumentsAddition{ documents, primary_key },
             enqueued_at: Utc::now(),
         }
     }
 
-    fn documents_partial(documents: Vec<IndexMap<String, Value>>) -> Update {
+    fn documents_partial(primary_key: Option<String>, documents: Vec<IndexMap<String, Value>>) -> Update {
         Update {
-            data: UpdateData::DocumentsPartial(documents),
+            data: UpdateData::DocumentsPartial{ documents, primary_key },
             enqueued_at: Utc::now(),
         }
     }
@@ -85,8 +85,15 @@ impl Update {
 pub enum UpdateData {
     ClearAll,
     Customs(Vec<u8>),
-    DocumentsAddition(Vec<IndexMap<String, Value>>),
-    DocumentsPartial(Vec<IndexMap<String, Value>>),
+    // (primary key, documents)
+    DocumentsAddition {
+        primary_key: Option<String>,
+        documents: Vec<IndexMap<String, Value>>
+    },
+    DocumentsPartial {
+        primary_key: Option<String>,
+        documents: Vec<IndexMap<String, Value>>,
+    },
     DocumentsDeletion(Vec<String>),
     Settings(Box<SettingsUpdate>)
 }
@@ -96,11 +103,11 @@ impl UpdateData {
         match self {
             UpdateData::ClearAll => UpdateType::ClearAll,
             UpdateData::Customs(_) => UpdateType::Customs,
-            UpdateData::DocumentsAddition(addition) => UpdateType::DocumentsAddition {
-                number: addition.len(),
+            UpdateData::DocumentsAddition{ documents, .. } => UpdateType::DocumentsAddition {
+                number: documents.len(),
             },
-            UpdateData::DocumentsPartial(addition) => UpdateType::DocumentsPartial {
-                number: addition.len(),
+            UpdateData::DocumentsPartial{ documents, .. } => UpdateType::DocumentsPartial {
+                number: documents.len(),
             },
             UpdateData::DocumentsDeletion(deletion) => UpdateType::DocumentsDeletion {
                 number: deletion.len(),
@@ -239,25 +246,25 @@ pub fn update_task(
 
             (update_type, result, start.elapsed())
         }
-        UpdateData::DocumentsAddition(documents) => {
+        UpdateData::DocumentsAddition { documents, primary_key } => {
             let start = Instant::now();
 
             let update_type = UpdateType::DocumentsAddition {
                 number: documents.len(),
             };
 
-            let result = apply_documents_addition(writer, index, documents);
+            let result = apply_documents_addition(writer, index, documents, primary_key);
 
             (update_type, result, start.elapsed())
         }
-        UpdateData::DocumentsPartial(documents) => {
+        UpdateData::DocumentsPartial{ documents, primary_key } => {
             let start = Instant::now();
 
             let update_type = UpdateType::DocumentsPartial {
                 number: documents.len(),
             };
 
-            let result = apply_documents_partial_addition(writer, index, documents);
+            let result = apply_documents_partial_addition(writer, index, documents, primary_key);
 
             (update_type, result, start.elapsed())
         }

meilisearch-http/src/dump.rs

@@ -130,13 +130,13 @@ fn import_index_v1(
         // if buffer is full, create and apply a batch, and clean buffer
         if values.len() == document_batch_size {
             let batch = std::mem::replace(&mut values, Vec::with_capacity(document_batch_size));
-            apply_documents_addition(write_txn, &index, batch)?;
+            apply_documents_addition(write_txn, &index, batch, None)?;
         }
     }
 
     // apply documents remaining in the buffer
     if !values.is_empty() {
-        apply_documents_addition(write_txn, &index, values)?;
+        apply_documents_addition(write_txn, &index, values, None)?;
     }
 
     // sync index information: stats, updated_at, last_update
@@ -289,7 +289,6 @@ fn dump_index_documents(data: &web::Data<Data>, reader: &MainReader, dir_path: &
 /// Write error with a context.
 fn fail_dump_process<E: std::error::Error>(data: &web::Data<Data>, dump_info: DumpInfo, context: &str, error: E) {
         let error_message = format!("{}; {}", context, error);
-        
         error!("Something went wrong during dump process: {}", &error_message);
         data.set_current_dump_info(dump_info.with_error(Error::dump_failed(error_message).into()))
 }

meilisearch-http/src/routes/document.rs

@@ -145,15 +145,6 @@ async fn get_all_documents(
     Ok(HttpResponse::Ok().json(documents))
 }
 
-fn find_primary_key(document: &IndexMap<String, Value>) -> Option<String> {
-    for key in document.keys() {
-        if key.to_lowercase().contains("id") {
-            return Some(key.to_string());
-        }
-    }
-    None
-}
-
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase", deny_unknown_fields)]
 struct UpdateDocumentsQuery {
@@ -168,26 +159,6 @@ async fn update_multiple_documents(
     is_partial: bool,
 ) -> Result<HttpResponse, ResponseError> {
     let update_id = data.get_or_create_index(&path.index_uid, |index| {
-        let reader = data.db.main_read_txn()?;
-
-        let mut schema = index
-            .main
-            .schema(&reader)?
-            .ok_or(meilisearch_core::Error::SchemaMissing)?;
-
-        if schema.primary_key().is_none() {
-            let id = match &params.primary_key {
-                Some(id) => id.to_string(),
-                None => body
-                    .first()
-                    .and_then(find_primary_key)
-                    .ok_or(meilisearch_core::Error::MissingPrimaryKey)?,
-            };
-
-            schema.set_primary_key(&id).map_err(Error::bad_request)?;
-
-            data.db.main_write(|w| index.main.put_schema(w, &schema))?;
-        }
 
         let mut document_addition = if is_partial {
             index.documents_partial_addition()
@@ -195,6 +166,27 @@ async fn update_multiple_documents(
             index.documents_addition()
         };
 
+        // Return an early error if primary key is already set, otherwise, try to set it up in the
+        // update later.
+        let reader = data.db.main_read_txn()?;
+        let schema = index
+            .main
+            .schema(&reader)?
+            .ok_or(meilisearch_core::Error::SchemaMissing)?;
+
+        match (params.into_inner().primary_key, schema.primary_key()) {
+            (Some(_), Some(_)) => return Err(meilisearch_schema::Error::PrimaryKeyAlreadyPresent)?,
+            (Some(key), None) => document_addition.set_primary_key(key),
+            (None, None) => {
+                let key = body
+                    .first()
+                    .and_then(find_primary_key)
+                    .ok_or(meilisearch_core::Error::MissingPrimaryKey)?;
+                document_addition.set_primary_key(key);
+            }
+            (None, Some(_)) => ()
+        }
+
         for document in body.into_inner() {
             document_addition.update_document(document);
         }
@@ -204,6 +196,15 @@ async fn update_multiple_documents(
     return Ok(HttpResponse::Accepted().json(IndexUpdateResponse::with_id(update_id)));
 }
 
+fn find_primary_key(document: &IndexMap<String, Value>) -> Option<String> {
+    for key in document.keys() {
+        if key.to_lowercase().contains("id") {
+            return Some(key.to_string());
+        }
+    }
+    None
+}
+
 #[post("/indexes/{index_uid}/documents", wrap = "Authentication::Private")]
 async fn add_documents(
     data: web::Data<Data>,

meilisearch-http/tests/index_update.rs

@@ -94,8 +94,16 @@ async fn return_update_status_of_pushed_documents() {
     ];
 
     let mut update_ids = Vec::new();
+    let mut bodies = bodies.into_iter();
 
     let url = "/indexes/test/documents?primaryKey=title";
+    let (response, status_code) = server.post_request(&url, bodies.next().unwrap()).await;
+    assert_eq!(status_code, 202);
+    let update_id = response["updateId"].as_u64().unwrap();
+    update_ids.push(update_id);
+    server.wait_update_id(update_id).await;
+
+    let url = "/indexes/test/documents";
     for body in bodies {
         let (response, status_code) = server.post_request(&url, body).await;
         assert_eq!(status_code, 202);