LSP Optimization: Use rayon for parse_ast_to_typed_tokens function (#…

…5473)

## Description

Traversing the typed tokens from `sway-lib-std` now goes from
`3.957833ms` → `1.372ms`. Collecting typed tokens on the benchmark
project goes from `139.6ms` → `127.7ms`.

I noticed that there are transient contentions when reading items from
the `TokenMap` with this change. As such, I've implemented a new
function on the `TokenMap` called `try_get_mut_with_retry`. This method
tries to access the value up to 8 times if the lock is still held using
the backoff and retry pattern.

related to #5445

sway-lsp/src/core/session.rs

@@ -24,7 +24,7 @@ use lsp_types::{
 };
 use parking_lot::RwLock;
 use pkg::{manifest::ManifestFile, BuildPlan};
-use rayon::iter::{ParallelBridge, ParallelIterator};
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use std::{
     ops::Deref,
     path::PathBuf,
@@ -561,16 +561,19 @@ fn parse_ast_to_tokens(
     ctx: &ParseContext,
     f: impl Fn(&AstNode, &ParseContext) + Sync,
 ) {
-    let root_nodes = parse_program.root.tree.root_nodes.iter();
-    let sub_nodes = parse_program
+    let nodes = parse_program
         .root
-        .submodules_recursive()
-        .flat_map(|(_, submodule)| &submodule.module.tree.root_nodes);
-
-    root_nodes
-        .chain(sub_nodes)
-        .par_bridge()
-        .for_each(|n| f(n, ctx));
+        .tree
+        .root_nodes
+        .iter()
+        .chain(
+            parse_program
+                .root
+                .submodules_recursive()
+                .flat_map(|(_, submodule)| &submodule.module.tree.root_nodes),
+        )
+        .collect::<Vec<_>>();
+    nodes.par_iter().for_each(|n| f(n, ctx));
 }
 
 /// Parse the [ty::TyProgram] AST to populate the [TokenMap] with typed AST nodes.
@@ -579,13 +582,18 @@ fn parse_ast_to_typed_tokens(
     ctx: &ParseContext,
     f: impl Fn(&ty::TyAstNode, &ParseContext) + Sync,
 ) {
-    let root_nodes = typed_program.root.all_nodes.iter();
-    let sub_nodes = typed_program
+    let nodes = typed_program
         .root
-        .submodules_recursive()
-        .flat_map(|(_, submodule)| submodule.module.all_nodes.iter());
-
-    root_nodes.chain(sub_nodes).for_each(|n| f(n, ctx));
+        .all_nodes
+        .iter()
+        .chain(
+            typed_program
+                .root
+                .submodules_recursive()
+                .flat_map(|(_, submodule)| &submodule.module.all_nodes),
+        )
+        .collect::<Vec<_>>();
+    nodes.par_iter().for_each(|n| f(n, ctx));
 }
 
 #[cfg(test)]

sway-lsp/src/core/token_map.rs

@@ -1,6 +1,7 @@
 use crate::core::token::{self, Token, TokenIdent, TypedAstToken};
-use dashmap::DashMap;
+use dashmap::{mapref::one::RefMut, try_result::TryResult, DashMap};
 use lsp_types::{Position, Url};
+use std::{thread, time::Duration};
 use sway_core::{language::ty, type_system::TypeId, Engines};
 use sway_types::{Ident, SourceEngine, Spanned};
 
@@ -14,12 +15,46 @@ pub use crate::core::token_map_ext::TokenMapExt;
 #[derive(Debug, Default)]
 pub struct TokenMap(DashMap<TokenIdent, Token>);
 
-impl TokenMap {
+impl<'a> TokenMap {
     /// Create a new token map.
     pub fn new() -> TokenMap {
         TokenMap(DashMap::with_capacity(2048))
     }
 
+    /// Attempts to get a mutable reference to a token with retries on lock.
+    /// Retries up to 10 times with increasing backoff (1ns, 10ns, 100ns, 500ns, 1µs, 10µs, 100µs, 1ms, 10ms, 50ms).
+    pub fn try_get_mut_with_retry(
+        &'a self,
+        ident: &TokenIdent,
+    ) -> Option<RefMut<TokenIdent, Token>> {
+        const MAX_RETRIES: usize = 10;
+        let backoff_times = [
+            1, 10, 100, 500, 1_000, 10_000, 100_000, 1_000_000, 10_000_000, 50_000_000,
+        ]; // Backoff times in nanoseconds
+        for (i, sleep) in backoff_times.iter().enumerate().take(MAX_RETRIES) {
+            match self.try_get_mut(ident) {
+                TryResult::Present(token) => return Some(token),
+                TryResult::Absent => return None,
+                TryResult::Locked => {
+                    tracing::warn!(
+                        "Failed to get token, retrying attmpt {}: {:#?}",
+                        i,
+                        ident.name
+                    );
+                    // Wait for the specified backoff time before retrying
+                    let backoff_time = Duration::from_nanos(*sleep);
+                    thread::sleep(backoff_time);
+                }
+            }
+        }
+        tracing::error!(
+            "Failed to get token after {} retries: {:#?}",
+            MAX_RETRIES,
+            ident
+        );
+        None // Return None if all retries are exhausted
+    }
+
     /// Create a custom iterator for the TokenMap.
     ///
     /// The iterator returns ([Ident], [Token]) pairs.

sway-lsp/src/traverse/dependency.rs

@@ -45,7 +45,7 @@ pub fn collect_typed_declaration(node: &ty::TyAstNode, ctx: &ParseContext) {
         };
 
         let token_ident = ctx.ident(&ident);
-        if let Some(mut token) = ctx.tokens.try_get_mut(&token_ident).try_unwrap() {
+        if let Some(mut token) = ctx.tokens.try_get_mut_with_retry(&token_ident) {
             token.typed = Some(typed_token);
             token.type_def = Some(TypeDefinition::Ident(ident));
         }