fix(internal): use handrolled async lock hashmap

Cargo.toml

@@ -77,6 +77,7 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 tree-sitter-javascript = "0.23.1"
 tree-sitter-python = "0.23.4"
 tree-sitter.workspace = true
+ouroboros = "0.18.4"
 
 [target.'cfg(all(target_os = "linux", any(not(target_env = "gnu"), not(target_pointer_width = "64"))))'.dependencies]
 self_update = { version = "0.41.0", default-features = false, features = ["archive-tar", "archive-zip", "compression-flate2", "compression-zip-deflate", "rustls"] }

src/backend.rs

@@ -30,7 +30,7 @@ use odoo_lsp::{format_loc, some, utils::*};
 pub struct Backend {
 	pub client: Client,
 	/// fs path -> rope, diagnostics etc.
-	pub document_map: DashMap<String, Document>,
+	pub document_map: RwMap<String, Document>,
 	pub record_ranges: DashMap<String, Box<[ByteRange]>>,
 	pub ast_map: DashMap<String, Tree>,
 	pub index: Index,
@@ -110,14 +110,14 @@ impl Backend {
 	#[instrument(skip_all)]
 	pub async fn on_change(&self, params: TextDocumentItem) -> miette::Result<()> {
 		let split_uri = params.uri.path().rsplit_once('.');
-		self.root_setup.wait().await;
 		let mut document = self
 			.document_map
-			.try_get_mut(params.uri.path())
-			.expect(format_loc!("deadlock"))
+			.get_mut(params.uri.path())
+			.await
 			.expect(format_loc!("(on_change) did not build document"));
 		match &params.text {
 			Text::Full(full) => {
+				let mut document = document.as_mut();
 				document.rope = ropey::Rope::from_str(full);
 				document.damage_zone = None;
 			}
@@ -140,9 +140,10 @@ impl Backend {
 						params.uri.path(),
 						&rope,
 						params.text.damage_zone(&rope, None),
-						&mut document.diagnostics_cache,
+						&mut document.as_mut().diagnostics_cache,
 					);
 				} else {
+					let mut document = document.as_mut();
 					document.damage_zone = params.text.damage_zone(&rope, document.damage_zone.take());
 				}
 			}
@@ -171,22 +172,21 @@ impl Backend {
 		let root = self.find_root_of(&path).ok_or_else(|| diagnostic!("out of root"))?;
 		let root = interner().get_or_intern_path(&root);
 
-		let mut document = self
-			.document_map
-			.try_get_mut(uri.path())
-			.expect(format_loc!("deadlock"))
-			.ok_or_else(|| diagnostic!("(did_save) did not build document"))?;
-
 		match extension {
-			"py" => self.did_save_python(uri, root, &mut document).await,
-			"xml" => self.did_save_xml(uri, root, &mut document).await?,
+			"py" => self.did_save_python(uri, root).await?,
+			"xml" => self.did_save_xml(uri, root).await?,
 			_ => {}
 		}
 		Ok(())
 	}
-	pub async fn did_save_python(&self, uri: Url, root: Spur, document: &mut Document) {
+	pub async fn did_save_python(&self, uri: Url, root: Spur) -> miette::Result<()> {
 		let path = uri.to_file_path().unwrap();
-		let zone = document.damage_zone.take();
+		let mut document = self
+			.document_map
+			.get_mut(uri.path())
+			.await
+			.ok_or_else(|| diagnostic!("(did_save) did not build document"))?;
+		let zone = document.as_mut().damage_zone.take();
 		let rope = &document.rope;
 		let text = Cow::from(rope);
 		_ = self
@@ -199,14 +199,21 @@ impl Backend {
 				uri.path(),
 				&document.rope.clone(),
 				zone,
-				&mut document.diagnostics_cache,
+				&mut document.as_mut().diagnostics_cache,
 			);
 			self.client
 				.publish_diagnostics(uri, document.diagnostics_cache.clone(), None)
 				.await;
 		}
+
+		Ok(())
 	}
-	pub async fn did_save_xml(&self, uri: Url, root: Spur, document: &mut Document) -> miette::Result<()> {
+	pub async fn did_save_xml(&self, uri: Url, root: Spur) -> miette::Result<()> {
+		let document = self
+			.document_map
+			.get(uri.path())
+			.await
+			.ok_or_else(|| diagnostic!("(did_save) did not build document"))?;
 		self.update_xml(root, &Text::Delta(vec![]), &uri, &document.rope, true)
 			.await?;
 		Ok(())
@@ -876,7 +883,8 @@ impl Backend {
 		let symbols_usage = interner.current_memory_usage();
 		Ok(serde_json::json! {{
 			"debug": cfg!(debug_assertions),
-			"documents": document_map.usage(),
+			// TODO
+			// "documents": document_map.usage(),
 			"records": record_ranges.usage(),
 			"ast": ast_map.usage(),
 			"index": index.statistics(),

src/main.rs

@@ -240,7 +240,7 @@ impl LanguageServer for Backend {
 			references_limit: _,
 			completions_limit: _,
 		} = self;
-		document_map.remove(path);
+		document_map.remove(path).await;
 		record_ranges.remove(path);
 		ast_map.remove(path);
 
@@ -335,7 +335,8 @@ impl LanguageServer for Backend {
 
 		let rope = Rope::from_str(&params.text_document.text);
 		self.document_map
-			.insert(params.text_document.uri.path().to_string(), Document::new(rope.clone()));
+			.insert(params.text_document.uri.path().to_string(), Document::new(rope.clone()))
+			.await;
 
 		self.root_setup.wait().await;
 		let path = params.text_document.uri.to_file_path().unwrap();
@@ -400,8 +401,8 @@ impl LanguageServer for Backend {
 		{
 			let mut document = self
 				.document_map
-				.try_get_mut(params.text_document.uri.path())
-				.expect(format_loc!("deadlock"))
+				.get_mut(params.text_document.uri.path())
+				.await
 				.expect("Did not build a document");
 			old_rope = document.rope.clone();
 			// Update the rope
@@ -410,12 +411,12 @@ impl LanguageServer for Backend {
 			// deltas are applied SEQUENTIALLY so we don't have to do any extra processing.
 			for change in &params.content_changes {
 				if change.range.is_none() && change.range_length.is_none() {
-					document.value_mut().rope = ropey::Rope::from_str(&change.text);
+					document.as_mut().rope = ropey::Rope::from_str(&change.text);
 				} else {
 					let range = change.range.expect("LSP change event must have a range");
 					let range =
 						lsp_range_to_char_range(range, &document.rope).expect("did_change applying delta: no range");
-					let rope = &mut document.value_mut().rope;
+					let rope = &mut document.as_mut().rope;
 					rope.remove(range.erase());
 					if !change.text.is_empty() {
 						rope.insert(range.start.0, &change.text);
@@ -453,7 +454,7 @@ impl LanguageServer for Backend {
 			debug!("(goto_definition) unsupported: {}", uri.path());
 			return Ok(None);
 		};
-		let Some(document) = self.document_map.get(uri.path()) else {
+		let Some(document) = self.document_map.get(uri.path()).await else {
 			panic!("Bug: did not build a document for {}", uri.path());
 		};
 		let location = match ext {
@@ -483,7 +484,7 @@ impl LanguageServer for Backend {
 		let Some((_, ext)) = uri.path().rsplit_once('.') else {
 			return Ok(None); // hit a directory, super unlikely
 		};
-		let Some(document) = self.document_map.get(uri.path()) else {
+		let Some(document) = self.document_map.get(uri.path()).await else {
 			debug!("Bug: did not build a document for {}", uri.path());
 			return Ok(None);
 		};
@@ -508,7 +509,7 @@ impl LanguageServer for Backend {
 			return Ok(None); // hit a directory, super unlikely
 		};
 
-		let Some(document) = self.document_map.get(uri.path()) else {
+		let Some(document) = self.document_map.get(uri.path()).await else {
 			debug!("Bug: did not build a document for {}", uri.path());
 			return Ok(None);
 		};
@@ -623,7 +624,7 @@ impl LanguageServer for Backend {
 			return Ok(None);
 		}
 		let uri = &params.text_document_position_params.text_document.uri;
-		let document = some!(self.document_map.get(uri.path()));
+		let document = some!(self.document_map.get(uri.path()).await);
 		let (_, ext) = some!(uri.path().rsplit_once('.'));
 		let hover = match ext {
 			"py" => self.python_hover(params, document.rope.clone()),
@@ -759,7 +760,8 @@ impl LanguageServer for Backend {
 		debug!("{path}");
 		let mut diagnostics = vec![];
 		if let Some((_, "py")) = path.rsplit_once('.') {
-			if let Some(mut document) = self.document_map.try_get_mut(path).expect(format_loc!("deadlock")) {
+			if let Some(mut document) = self.document_map.get_mut(path).await {
+				let mut document = document.as_mut();
 				let damage_zone = document.damage_zone.take();
 				let rope = &document.rope.clone();
 				self.diagnose_python(
@@ -790,7 +792,7 @@ impl LanguageServer for Backend {
 			return Ok(None);
 		};
 
-		let document = some!(self.document_map.get(params.text_document.uri.path()));
+		let document = some!(self.document_map.get(params.text_document.uri.path()).await);
 
 		Ok(self
 			.xml_code_actions(params, document.rope.clone())
@@ -887,7 +889,7 @@ async fn main() {
 	let (service, socket) = LspService::build(|client| Backend {
 		client,
 		index: Default::default(),
-		document_map: DashMap::with_shard_amount(4),
+		document_map: RwMap::new(),
 		record_ranges: DashMap::with_shard_amount(4),
 		roots: DashSet::default(),
 		capabilities: Default::default(),

src/utils.rs

@@ -21,6 +21,9 @@ use crate::index::PathSymbol;
 mod usage;
 pub use usage::{Usage, UsageInfo};
 
+mod map;
+pub use map::RwMap;
+
 /// Unwraps the option in the context of a function that returns [`Result<Option<_>>`].
 #[macro_export]
 macro_rules! some {