Centralize Mutex Locking (#51)

server/src/diagnostic_ext.rs

@@ -1,15 +1,13 @@
 // Copyright (c) ZeroC, Inc.
 
 use crate::configuration_set::ConfigurationSet;
-use crate::session::Session;
 use crate::utils::convert_slice_path_to_uri;
 use crate::{notifications, show_popup};
 
 use slicec::diagnostics::{Diagnostic, DiagnosticLevel, Note};
 use std::collections::{HashMap, HashSet};
-use tokio::sync::Mutex;
 use tower_lsp::lsp_types::{
-    DiagnosticRelatedInformation, Location, MessageType, NumberOrString, Position, Range, Url,
+    DiagnosticRelatedInformation, Location, NumberOrString, Position, Range, Url,
 };
 use tower_lsp::Client;
 
@@ -47,26 +45,6 @@ pub async fn publish_diagnostics_for_set(
     }
 }
 
-/// Triggers and compilation and publishes any diagnostics that are reported.
-/// It does this for all configuration sets.
-pub async fn compile_and_publish_diagnostics(client: &Client, session: &Session) {
-    let mut configuration_sets = session.configuration_sets.lock().await;
-    let server_config = session.server_config.read().await;
-
-    client
-        .log_message(
-            MessageType::INFO,
-            "Publishing diagnostics for all configuration sets.",
-        )
-        .await;
-    for configuration_set in configuration_sets.iter_mut() {
-        // Trigger a compilation and get any diagnostics that were reported during it.
-        let diagnostics = configuration_set.trigger_compilation(&server_config);
-        // Publish those diagnostics.
-        publish_diagnostics_for_set(client, diagnostics, configuration_set).await;
-    }
-}
-
 /// Processes a list of diagnostics and updates the publish map with LSP-compatible diagnostics.
 ///
 /// This function filters out any diagnostics that do not have a span or cannot be converted
@@ -106,8 +84,7 @@ pub fn process_diagnostics(
 ///
 /// This function iterates over all configuration sets, collects all tracked file URIs,
 /// and then publishes empty diagnostics to clear existing ones for each URI.
-pub async fn clear_diagnostics(client: &Client, configuration_sets: &Mutex<Vec<ConfigurationSet>>) {
-    let configuration_sets = configuration_sets.lock().await;
+pub async fn clear_diagnostics(client: &Client, configuration_sets: &[ConfigurationSet]) {
     let mut all_tracked_files = HashSet::new();
     for configuration_set in configuration_sets.iter() {
         configuration_set

server/src/main.rs

@@ -1,12 +1,14 @@
 // Copyright (c) ZeroC, Inc.
 
-use crate::diagnostic_ext::{clear_diagnostics, compile_and_publish_diagnostics, process_diagnostics};
+use crate::diagnostic_ext::{clear_diagnostics, process_diagnostics, publish_diagnostics_for_set};
 use crate::hover::try_into_hover_result;
 use crate::jump_definition::get_definition_span;
 use crate::notifications::{ShowNotification, ShowNotificationParams};
 use crate::session::Session;
 use crate::slice_config::compute_slice_options;
+use std::ops::DerefMut;
 use std::{collections::HashMap, path::Path};
+use tokio::sync::Mutex;
 use tower_lsp::{jsonrpc::Error, lsp_types::*, Client, LanguageServer, LspService, Server};
 use utils::{convert_slice_path_to_uri, url_to_sanitized_file_path};
 
@@ -30,12 +32,12 @@ async fn main() {
 
 struct Backend {
     client: Client,
-    session: Session,
+    session: Mutex<Session>,
 }
 
 impl Backend {
     pub fn new(client: tower_lsp::Client) -> Self {
-        let session = Session::default();
+        let session = Mutex::new(Session::default());
         Self { client, session }
     }
 
@@ -76,15 +78,15 @@ impl Backend {
             .log_message(MessageType::INFO, format!("File '{}' changed", file_path.display()))
             .await;
 
-        let mut configuration_sets = self.session.configuration_sets.lock().await;
-        let server_config = self.session.server_config.read().await;
+        let mut session_guard = self.session.lock().await;
+        let Session { configuration_sets, server_config } = session_guard.deref_mut();
 
         let mut publish_map = HashMap::new();
         let mut diagnostics = Vec::new();
 
         // Process each configuration set that contains the changed file
         for set in configuration_sets.iter_mut().filter(|set| {
-            compute_slice_options(&server_config, &set.slice_config)
+            compute_slice_options(server_config, &set.slice_config)
                 .references
                 .into_iter()
                 .any(|f| {
@@ -93,7 +95,7 @@ impl Backend {
                 })
         }) {
             // `trigger_compilation` compiles the configuration set's files and returns any diagnostics.
-            diagnostics.extend(set.trigger_compilation(&server_config));
+            diagnostics.extend(set.trigger_compilation(server_config));
 
             // Update publish_map with files to be updated
             publish_map.extend(
@@ -134,6 +136,26 @@ impl Backend {
                 .await;
         }
     }
+
+    /// Triggers and compilation and publishes any diagnostics that are reported.
+    /// It does this for all configuration sets.
+    pub async fn compile_and_publish_diagnostics(&self) {
+        let mut session_guard = self.session.lock().await;
+        let Session { configuration_sets, server_config } = session_guard.deref_mut();
+
+        self.client
+            .log_message(
+                MessageType::INFO,
+                "Publishing diagnostics for all configuration sets.",
+            )
+            .await;
+        for configuration_set in configuration_sets.iter_mut() {
+            // Trigger a compilation and get any diagnostics that were reported during it.
+            let diagnostics = configuration_set.trigger_compilation(server_config);
+            // Publish those diagnostics.
+            publish_diagnostics_for_set(&self.client, diagnostics, configuration_set).await;
+        }
+    }
 }
 
 #[tower_lsp::async_trait]
@@ -142,9 +164,10 @@ impl LanguageServer for Backend {
         &self,
         params: InitializeParams,
     ) -> tower_lsp::jsonrpc::Result<InitializeResult> {
-        self.session.update_from_initialize_params(params).await;
-        let capabilities = Backend::capabilities();
+        let mut session_guard = self.session.lock().await;
+        session_guard.update_from_initialize_params(params);
 
+        let capabilities = Backend::capabilities();
         Ok(InitializeResult {
             capabilities,
             ..InitializeResult::default()
@@ -153,7 +176,7 @@ impl LanguageServer for Backend {
 
     async fn initialized(&self, _: InitializedParams) {
         // Now that the server and client are fully initialized, it's safe to compile and publish any diagnostics.
-        compile_and_publish_diagnostics(&self.client, &self.session).await;
+        self.compile_and_publish_diagnostics().await;
     }
 
     async fn shutdown(&self) -> tower_lsp::jsonrpc::Result<()> {
@@ -165,15 +188,20 @@ impl LanguageServer for Backend {
             .log_message(MessageType::INFO, "Extension settings changed")
             .await;
 
-        // When the configuration changes, any of the files in the workspace could be impacted. Therefore, we need to
-        // clear the diagnostics for all files and then re-publish them.
-        clear_diagnostics(&self.client, &self.session.configuration_sets).await;
+        // Explicit scope to ensure the session lock guard is dropped before we start compilation.
+        {
+            let mut session_guard = self.session.lock().await;
+
+            // When the configuration changes, any of the files in the workspace could be impacted. Therefore, we need to
+            // clear the diagnostics for all files and then re-publish them.
+            clear_diagnostics(&self.client, &session_guard.configuration_sets).await;
 
-        // Update the stored configuration sets from the data provided in the client notification
-        self.session.update_configurations_from_params(params).await;
+            // Update the stored configuration sets from the data provided in the client notification
+            session_guard.update_configurations_from_params(params);
+        }
 
         // Trigger a compilation and publish the diagnostics for all files
-        compile_and_publish_diagnostics(&self.client, &self.session).await;
+        self.compile_and_publish_diagnostics().await;
     }
 
     async fn goto_definition(
@@ -187,7 +215,8 @@ impl LanguageServer for Backend {
         let file_path = url_to_sanitized_file_path(&uri).ok_or_else(Error::internal_error)?;
 
         // Find the configuration set that contains the file
-        let configuration_sets = self.session.configuration_sets.lock().await;
+        let session_guard = self.session.lock().await;
+        let configuration_sets = &session_guard.configuration_sets;
 
         // Get the definition span and convert it to a GotoDefinitionResponse
         Ok(configuration_sets.iter().find_map(|set| {
@@ -220,7 +249,8 @@ impl LanguageServer for Backend {
         let file_path = url_to_sanitized_file_path(&uri).ok_or_else(Error::internal_error)?;
 
         // Find the configuration set that contains the file and get the hover info
-        let configuration_sets = self.session.configuration_sets.lock().await;
+        let session_guard = self.session.lock().await;
+        let configuration_sets = &session_guard.configuration_sets;
 
         Ok(configuration_sets.iter().find_map(|set| {
             let files = &set.compilation_data.files;

server/src/session.rs

@@ -3,25 +3,21 @@
 use crate::configuration_set::ConfigurationSet;
 use crate::slice_config::ServerConfig;
 use crate::utils::{sanitize_path, url_to_sanitized_file_path};
-use tokio::sync::{Mutex, RwLock};
-use tower_lsp::lsp_types::DidChangeConfigurationParams;
+use tower_lsp::lsp_types::{DidChangeConfigurationParams, InitializeParams};
 
 #[derive(Debug, Default)]
 pub struct Session {
     /// This vector contains all of the configuration sets for the language server. Each element is a tuple containing
     /// `SliceConfig` and `CompilationState`. The `SliceConfig` is used to determine which configuration set to use when
     /// publishing diagnostics. The `CompilationState` is used to retrieve the diagnostics for a given file.
-    pub configuration_sets: Mutex<Vec<ConfigurationSet>>,
+    pub configuration_sets: Vec<ConfigurationSet>,
     /// Configuration that affects the entire server.
-    pub server_config: RwLock<ServerConfig>,
+    pub server_config: ServerConfig,
 }
 
 impl Session {
     // Update the properties of the session from `InitializeParams`
-    pub async fn update_from_initialize_params(
-        &self,
-        params: tower_lsp::lsp_types::InitializeParams,
-    ) {
+    pub fn update_from_initialize_params(&mut self, params: InitializeParams) {
         let initialization_options = params.initialization_options;
 
         // Use the root_uri if it exists temporarily as we cannot access configuration until
@@ -41,7 +37,7 @@ impl Session {
             .map(sanitize_path)
             .expect("builtInSlicePath not found in initialization options");
 
-        *self.server_config.write().await = ServerConfig { workspace_root_path, built_in_slice_path };
+        self.server_config = ServerConfig { workspace_root_path, built_in_slice_path };
 
         // Load any user configuration from the 'slice.configurations' option.
         let configuration_sets = initialization_options
@@ -51,11 +47,11 @@ impl Session {
             .map(|arr| ConfigurationSet::parse_configuration_sets(arr))
             .unwrap_or_default();
 
-        self.update_configurations(configuration_sets).await;
+        self.update_configurations(configuration_sets);
     }
 
     // Update the configuration sets from the `DidChangeConfigurationParams` notification.
-    pub async fn update_configurations_from_params(&self, params: DidChangeConfigurationParams) {
+    pub fn update_configurations_from_params(&mut self, params: DidChangeConfigurationParams) {
         // Parse the configurations from the notification
         let configurations = params
             .settings
@@ -66,17 +62,17 @@ impl Session {
             .unwrap_or_default();
 
         // Update the configuration sets
-        self.update_configurations(configurations).await;
+        self.update_configurations(configurations);
     }
 
     // Update the configuration sets by replacing it with the new configurations. If there are no configuration sets
     // after updating, insert the default configuration set.
-    async fn update_configurations(&self, mut configurations: Vec<ConfigurationSet>) {
+    fn update_configurations(&mut self, mut configurations: Vec<ConfigurationSet>) {
         // Insert the default configuration set if needed
         if configurations.is_empty() {
             configurations.push(ConfigurationSet::default());
         }
 
-        *self.configuration_sets.lock().await = configurations;
+        self.configuration_sets = configurations;
     }
 }