Fix tokio blocking deadlock in search indexing by offloading add_document to spawn_blocking

This commit is contained in:
Riz Ashraf committed 2026-09-21 04:26:11 +01:00
1 parent 9870626045
commit 2ce4fa6720
1 file changed
+41 -25
+41 -25
View File
@@ -1,12 +1,12 @@
use crate::models::{Adr, Entity, Snippet, Task}; use crate::models::{Adr, Entity, Snippet, Task};
use std::sync::Mutex; use std::sync::{Arc, Mutex};
use tantivy::schema::*; use tantivy::schema::*;
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc}; use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};
pub struct MemoryIndex { pub struct MemoryIndex {
index: Index, index: Index,
reader: IndexReader, reader: IndexReader,
writer: Mutex<IndexWriter>, writer: Arc<Mutex<IndexWriter>>,
// Schema fields // Schema fields
pub id_field: Field, pub id_field: Field,
@@ -40,7 +40,7 @@ impl MemoryIndex {
Ok(Self { Ok(Self {
index, index,
reader, reader,
writer: Mutex::new(writer), writer: Arc::new(Mutex::new(writer)),
id_field, id_field,
title_field, title_field,
body_field, body_field,
@@ -49,28 +49,36 @@ impl MemoryIndex {
}) })
} }
pub fn index_entity(&self, e: &Entity) -> tantivy::Result<()> { pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = self.writer.lock().unwrap(); let writer = Arc::clone(&self.writer);
writer.add_document(doc!( let doc = doc!(
self.id_field => e.name.clone(), self.id_field => e.name.clone(),
self.title_field => e.name.clone(), self.title_field => e.name.clone(),
self.body_field => e.observations.join(" "), self.body_field => e.observations.join(" "),
self.type_field => "entity", self.type_field => "entity",
self.namespace_field => e.namespace.clone() self.namespace_field => e.namespace.clone()
))?; );
tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap();
writer.add_document(doc)?;
Ok(()) Ok(())
})
} }
pub fn index_task(&self, t: &Task) -> tantivy::Result<()> { pub fn index_task(&self, t: &Task) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = self.writer.lock().unwrap(); let writer = Arc::clone(&self.writer);
writer.add_document(doc!( let doc = doc!(
self.id_field => t.id.clone(), self.id_field => t.id.clone(),
self.title_field => t.title.clone(), self.title_field => t.title.clone(),
self.body_field => t.description.clone(), self.body_field => t.description.clone(),
self.type_field => "task", self.type_field => "task",
self.namespace_field => "global" self.namespace_field => "global"
))?; );
tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap();
writer.add_document(doc)?;
Ok(()) Ok(())
})
} }
pub fn commit(&self) -> tantivy::Result<()> { pub fn commit(&self) -> tantivy::Result<()> {
@@ -133,28 +141,36 @@ impl MemoryIndex {
Ok(results) Ok(results)
} }
pub fn index_snippet(&self, s: &Snippet) -> tantivy::Result<()> { pub fn index_snippet(&self, s: &Snippet) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = self.writer.lock().unwrap(); let writer = Arc::clone(&self.writer);
writer.add_document(doc!( let doc = doc!(
self.id_field => s.name.clone(), self.id_field => s.name.clone(),
self.title_field => s.name.clone(), self.title_field => s.name.clone(),
self.body_field => format!("{} {}", s.language, s.description), self.body_field => format!("{} {}", s.language, s.description),
self.type_field => "snippet", self.type_field => "snippet",
self.namespace_field => "global" self.namespace_field => "global"
))?; );
tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap();
writer.add_document(doc)?;
Ok(()) Ok(())
})
} }
pub fn index_adr(&self, a: &Adr) -> tantivy::Result<()> { pub fn index_adr(&self, a: &Adr) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = self.writer.lock().unwrap(); let writer = Arc::clone(&self.writer);
writer.add_document(doc!( let doc = doc!(
self.id_field => a.id.clone(), self.id_field => a.id.clone(),
self.title_field => a.title.clone(), self.title_field => a.title.clone(),
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence), self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
self.type_field => "adr", self.type_field => "adr",
self.namespace_field => "global" self.namespace_field => "global"
))?; );
tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap();
writer.add_document(doc)?;
Ok(()) Ok(())
})
} }
} }
@@ -163,8 +179,8 @@ mod tests {
use super::*; use super::*;
use tempfile::TempDir; use tempfile::TempDir;
#[test] #[tokio::test]
fn test_search_index_and_retrieve() { async fn test_search_index_and_retrieve() {
let temp_dir = TempDir::new().unwrap(); let temp_dir = TempDir::new().unwrap();
let index = MemoryIndex::new(temp_dir.path()).unwrap(); let index = MemoryIndex::new(temp_dir.path()).unwrap();
@@ -175,7 +191,7 @@ mod tests {
namespace: "global".to_string(), namespace: "global".to_string(),
git_branch: None, git_branch: None,
}; };
index.index_entity(&entity).unwrap(); index.index_entity(&entity).await.unwrap();
let task = Task { let task = Task {
id: "task-1".to_string(), id: "task-1".to_string(),
@@ -189,7 +205,7 @@ mod tests {
dependencies: vec![], dependencies: vec![],
parent_id: None, parent_id: None,
}; };
index.index_task(&task).unwrap(); index.index_task(&task).await.unwrap();
let snippet = Snippet { let snippet = Snippet {
name: "test_snippet".to_string(), name: "test_snippet".to_string(),
@@ -198,7 +214,7 @@ mod tests {
description: "A test snippet".to_string(), description: "A test snippet".to_string(),
updated_at: 0, updated_at: 0,
}; };
index.index_snippet(&snippet).unwrap(); index.index_snippet(&snippet).await.unwrap();
let adr = Adr { let adr = Adr {
id: "adr-1".to_string(), id: "adr-1".to_string(),
@@ -208,7 +224,7 @@ mod tests {
consequence: "Test consequence".to_string(), consequence: "Test consequence".to_string(),
timestamp: 0, timestamp: 0,
}; };
index.index_adr(&adr).unwrap(); index.index_adr(&adr).await.unwrap();
index.commit().unwrap(); index.commit().unwrap();
index.reader.reload().unwrap(); index.reader.reload().unwrap();