From 250666d3de43439dd4026ef844616c448b6ffde7 Mon Sep 17 00:00:00 2001 From: Joshua Batty Date: Wed, 10 Jul 2024 15:35:18 +1000 Subject: [PATCH] Add LRU Session Cache to Reduce RAM Usage (#6239) ## Description This PR introduces an LRU (Least Recently Used) cache mechanism for managing sessions in large workspaces. This change aims to address the issue of high RAM usage reported in workspaces with multiple projects. closes #6222 ## Checklist - [x] I have linked to any relevant issues. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have updated the documentation where relevant (API docs, the reference, and the Sway book). - [x] If my change requires substantial documentation changes, I have [requested support from the DevRel team](https://github.com/FuelLabs/devrel-requests/issues/new/choose) - [x] I have added tests that prove my fix is effective or that my feature works. - [x] I have added (or requested a maintainer to add) the necessary `Breaking*` or `New Feature` labels where relevant. - [x] I have done my best to ensure that my PR adheres to [the Fuel Labs Code Review Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md). - [x] I have requested a review from the relevant team or maintainers. --- sway-lsp/src/server_state.rs | 111 +++++++++++++++++++++++++++++++---- 1 file changed, 98 insertions(+), 13 deletions(-) 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); + } + } +}