Fix leak of LMDB connection in semantic index (#17992)

Max Brunsfeld , Richard Feldman , and Jason created 1 year ago

Apparently, to close LMDB's file descriptors when using the `heed`
library, you need to explicitly call `prepare_for_closing`.

Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <[e-mail]>
Co-authored-by: Jason <[e-mail]>

Change summary

crates/evals/src/eval.rs                    |  9 +++++++++
crates/semantic_index/src/semantic_index.rs | 12 +++++++++---
2 files changed, 18 insertions(+), 3 deletions(-)

Detailed changes

crates/evals/src/eval.rs

@@ -446,6 +446,15 @@ async fn run_evaluation(
 
             println!("{}", serde_json::to_string(&query_results).unwrap());
         }
+
+        user_store
+            .update(cx, |_, _| {
+                drop(semantic_index);
+                drop(project);
+                drop(worktree);
+                drop(project_index);
+            })
+            .unwrap();
     }
 
     eprint!(

crates/semantic_index/src/semantic_index.rs

@@ -25,7 +25,7 @@ pub use summary_index::FileSummary;
 
 pub struct SemanticDb {
     embedding_provider: Arc<dyn EmbeddingProvider>,
-    db_connection: heed::Env,
+    db_connection: Option<heed::Env>,
     project_indices: HashMap<WeakModel<Project>, Model<ProjectIndex>>,
 }
 
@@ -70,7 +70,7 @@ impl SemanticDb {
         .ok();
 
         Ok(SemanticDb {
-            db_connection,
+            db_connection: Some(db_connection),
             embedding_provider,
             project_indices: HashMap::default(),
         })
@@ -148,7 +148,7 @@ impl SemanticDb {
         let project_index = cx.new_model(|cx| {
             ProjectIndex::new(
                 project.clone(),
-                self.db_connection.clone(),
+                self.db_connection.clone().unwrap(),
                 self.embedding_provider.clone(),
                 cx,
             )
@@ -171,6 +171,12 @@ impl SemanticDb {
     }
 }
 
+impl Drop for SemanticDb {
+    fn drop(&mut self) {
+        self.db_connection.take().unwrap().prepare_for_closing();
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;