diff --git a/sway-lsp/src/server_state.rs b/sway-lsp/src/server_state.rs index 4415f0b2d91..f744adc8847 100644 --- a/sway-lsp/src/server_state.rs +++ b/sway-lsp/src/server_state.rs @@ -10,12 +10,15 @@ use crate::{ utils::{debug, keyword_docs::KeywordDocs}, }; use crossbeam_channel::{Receiver, Sender}; -use dashmap::DashMap; +use dashmap::{mapref::multiple::RefMulti, DashMap}; use forc_pkg::manifest::GenericManifestFile; use forc_pkg::PackageManifestFile; use lsp_types::{Diagnostic, Url}; -use parking_lot::RwLock; -use std::{collections::BTreeMap, process::Command}; +use parking_lot::{Mutex, RwLock}; +use std::{ + collections::{BTreeMap, VecDeque}, + process::Command, +}; use std::{ mem, path::PathBuf, @@ -28,13 +31,17 @@ use sway_core::LspConfig; use tokio::sync::Notify; use tower_lsp::{jsonrpc, Client}; +const DEFAULT_SESSION_CACHE_CAPACITY: usize = 4; + /// `ServerState` is the primary mutable state of the language server pub struct ServerState { pub(crate) client: Option, pub config: Arc>, pub(crate) keyword_docs: Arc, - /// A collection of [Session]s, each of which represents a project that has been opened in the users workspace - pub(crate) sessions: Arc>>, + /// A Least Recently Used (LRU) cache of [Session]s, each representing a project opened in the user's workspace. + /// This cache limits memory usage by maintaining a fixed number of active sessions, automatically + /// evicting the least recently used sessions when the capacity is reached. + pub(crate) sessions: LruSessionCache, pub documents: Documents, // Compilation thread related fields pub(crate) retrigger_compilation: Arc, @@ -52,7 +59,7 @@ impl Default for ServerState { client: None, config: Arc::new(RwLock::new(Config::default())), keyword_docs: Arc::new(KeywordDocs::new()), - sessions: Arc::new(DashMap::new()), + sessions: LruSessionCache::new(DEFAULT_SESSION_CACHE_CAPACITY), documents: Documents::new(), retrigger_compilation: Arc::new(AtomicBool::new(false)), is_compiling: Arc::new(AtomicBool::new(false)), @@ -357,17 +364,95 @@ impl ServerState { .ok_or(DirectoryError::ManifestDirNotFound)? .to_path_buf(); - let session = if let Some(item) = self.sessions.try_get(&manifest_dir).try_unwrap() { - item.value().clone() - } else { + let session = self.sessions.get(&manifest_dir).unwrap_or({ // If no session can be found, then we need to call init and insert a new session into the map self.init_session(uri).await?; self.sessions - .try_get(&manifest_dir) - .try_unwrap() - .map(|item| item.value().clone()) + .get(&manifest_dir) .expect("no session found even though it was just inserted into the map") - }; + }); Ok(session) } } + +/// A Least Recently Used (LRU) cache for storing and managing `Session` objects. +/// This cache helps limit memory usage by maintaining a fixed number of active sessions. +pub(crate) struct LruSessionCache { + /// Stores the actual `Session` objects, keyed by their file paths. + sessions: Arc>>, + /// Keeps track of the order in which sessions were accessed, with most recent at the front. + usage_order: Arc>>, + /// The maximum number of sessions that can be stored in the cache. + capacity: usize, +} + +impl LruSessionCache { + /// Creates a new `LruSessionCache` with the specified capacity. + pub(crate) fn new(capacity: usize) -> Self { + LruSessionCache { + sessions: Arc::new(DashMap::new()), + usage_order: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))), + capacity, + } + } + + pub(crate) fn iter(&self) -> impl Iterator>> { + self.sessions.iter() + } + + /// Retrieves a session from the cache and updates its position to the front of the usage order. + pub(crate) fn get(&self, path: &PathBuf) -> Option> { + if let Some(session) = self.sessions.try_get(path).try_unwrap() { + if self.sessions.len() >= self.capacity { + self.move_to_front(path); + } + Some(session.clone()) + } else { + None + } + } + + /// Inserts or updates a session in the cache. + /// If at capacity and inserting a new session, evicts the least recently used one. + /// For existing sessions, updates their position in the usage order if at capacity. + pub(crate) fn insert(&self, path: PathBuf, session: Arc) { + if self.sessions.get(&path).is_some() { + tracing::trace!("Updating existing session for path: {:?}", path); + // Session already exists, just update its position in the usage order if at capacity + if self.sessions.len() >= self.capacity { + self.move_to_front(&path); + } + } else { + // New session + tracing::trace!("Inserting new session for path: {:?}", path); + if self.sessions.len() >= self.capacity { + self.evict_least_used(); + } + self.sessions.insert(path.clone(), session); + let mut order = self.usage_order.lock(); + order.push_front(path); + } + } + + /// Moves the specified path to the front of the usage order, marking it as most recently used. + fn move_to_front(&self, path: &PathBuf) { + tracing::trace!("Moving path to front of usage order: {:?}", path); + let mut order = self.usage_order.lock(); + if let Some(index) = order.iter().position(|p| p == path) { + order.remove(index); + } + order.push_front(path.clone()); + } + + /// Removes the least recently used session from the cache when the capacity is reached. + fn evict_least_used(&self) { + let mut order = self.usage_order.lock(); + if let Some(old_path) = order.pop_back() { + tracing::trace!( + "Cache at capacity. Evicting least used session: {:?}", + old_path + ); + self.sessions.remove(&old_path); + } + } +}