Skip reassignment if head deleted

Author: sanketkedia

rust/index/src/spann/types.rs

@@ -840,6 +840,10 @@ impl SpannIndexWriter {
             );
             return Ok(());
         }
+        // Don't reassign if concurrently deleted.
+        if self.is_head_deleted(prev_head_id as usize).await? {
+            return Ok(());
+        }
         // RNG query to find the nearest heads.
         let (nearest_head_ids, _, nearest_head_embeddings) = self.rng_query(doc_embedding).await?;
         // Don't reassign if empty.
@@ -934,23 +938,20 @@ impl SpannIndexWriter {
         new_head_embeddings: &[Option<&Vec<f32>>],
         old_head_embedding: &[f32],
     ) -> Result<(), SpannIndexWriterError> {
-        // Get posting list of each neighbour and check for reassignment criteria.
-        let doc_offset_ids;
-        let doc_versions;
-        let doc_embeddings;
-        {
-            // TODO(Sanket): Check if head is deleted, can happen if another concurrent thread
-            // deletes it.
-            (doc_offset_ids, doc_versions, doc_embeddings) = self
-                .posting_list_writer
-                .get_owned::<u32, &SpannPostingList<'_>>("", head_id as u32)
-                .await
-                .map_err(|e| {
-                    tracing::error!("Error getting posting list for head {}: {}", head_id, e);
-                    SpannIndexWriterError::PostingListGetError(e)
-                })?
-                .ok_or(SpannIndexWriterError::PostingListNotFound)?;
+        // Head got concurrrently deleted so abort reassignment.
+        if self.is_head_deleted(head_id).await? {
+            return Ok(());
         }
+        // Get posting list of each neighbour and check for reassignment criteria.
+        let (doc_offset_ids, doc_versions, doc_embeddings) = self
+            .posting_list_writer
+            .get_owned::<u32, &SpannPostingList<'_>>("", head_id as u32)
+            .await
+            .map_err(|e| {
+                tracing::error!("Error getting posting list for head {}: {}", head_id, e);
+                SpannIndexWriterError::PostingListGetError(e)
+            })?
+            .ok_or(SpannIndexWriterError::PostingListNotFound)?;
         for (index, doc_offset_id) in doc_offset_ids.iter().enumerate() {
             if assigned_ids.contains(doc_offset_id)
                 || self
@@ -959,6 +960,10 @@ impl SpannIndexWriter {
             {
                 continue;
             }
+            if self.is_head_deleted(head_id).await? {
+                // Head got concurrently deleted so abort reassignment.
+                return Ok(());
+            }
             let distance_function: DistanceFunction = self.params.space.clone().into();
             let distance_from_curr_center = distance_function.distance(
                 &doc_embeddings[index * self.dimensionality..(index + 1) * self.dimensionality],
@@ -1437,7 +1442,6 @@ impl SpannIndexWriter {
                     doc_versions: &[version],
                     doc_embeddings: embeddings,
                 };
-                let _write_guard = self.posting_list_partitioned_mutex.lock(&id).await;
                 self.posting_list_writer
                     .set("", next_id, &posting_list)
                     .await