Refactor Language Server I/O to Use Tokio Async Primitives

sway-lsp/Cargo.toml

@@ -36,6 +36,7 @@ syn = { version = "1.0.73", features = ["full"] }
 tempfile = "3"
 thiserror = "1.0.30"
 tokio = { version = "1.3", features = [
+    "fs",
     "io-std",
     "io-util",
     "macros",

sway-lsp/benches/lsp_benchmarks/compile.rs

@@ -2,14 +2,18 @@ use criterion::{black_box, criterion_group, Criterion};
 use lsp_types::Url;
 use sway_core::Engines;
 use sway_lsp::core::session::{self, Session};
+use tokio::runtime::Runtime;
 
 const NUM_DID_CHANGE_ITERATIONS: usize = 4;
 
 fn benchmarks(c: &mut Criterion) {
     // Load the test project
     let uri = Url::from_file_path(super::benchmark_dir().join("src/main.sw")).unwrap();
-    let session = Session::new();
-    session.handle_open_file(&uri);
+    let session = Runtime::new().unwrap().block_on(async {
+        let session = Session::new();
+        session.handle_open_file(&uri).await;
+        session
+    });
 
     c.bench_function("compile", |b| {
         b.iter(|| {

sway-lsp/benches/lsp_benchmarks/mod.rs

@@ -6,11 +6,11 @@ use lsp_types::Url;
 use std::{path::PathBuf, sync::Arc};
 use sway_lsp::core::session::{self, Session};
 
-pub fn compile_test_project() -> (Url, Arc<Session>) {
+pub async fn compile_test_project() -> (Url, Arc<Session>) {
     let session = Session::new();
     // Load the test project
     let uri = Url::from_file_path(benchmark_dir().join("src/main.sw")).unwrap();
-    session.handle_open_file(&uri);
+    session.handle_open_file(&uri).await;
     // Compile the project and write the parse result to the session
     let parse_result = session::parse_project(&uri, &session.engines.read()).unwrap();
     session.write_parse_result(parse_result);

sway-lsp/benches/lsp_benchmarks/requests.rs

@@ -4,9 +4,12 @@ use lsp_types::{
     TextDocumentIdentifier,
 };
 use sway_lsp::{capabilities, lsp_ext::OnEnterParams, utils::keyword_docs::KeywordDocs};
+use tokio::runtime::Runtime;
 
 fn benchmarks(c: &mut Criterion) {
-    let (uri, session) = black_box(super::compile_test_project());
+    let (uri, session) = Runtime::new()
+        .unwrap()
+        .block_on(async { black_box(super::compile_test_project().await) });
     let config = sway_lsp::config::Config::default();
     let keyword_docs = KeywordDocs::new();
     let position = Position::new(1717, 24);

sway-lsp/benches/lsp_benchmarks/token_map.rs

@@ -1,8 +1,11 @@
 use criterion::{black_box, criterion_group, Criterion};
 use lsp_types::Position;
+use tokio::runtime::Runtime;
 
 fn benchmarks(c: &mut Criterion) {
-    let (uri, session) = black_box(super::compile_test_project());
+    let (uri, session) = Runtime::new()
+        .unwrap()
+        .block_on(async { black_box(super::compile_test_project().await) });
     let engines = session.engines.read();
     let position = Position::new(1716, 24);
 

sway-lsp/src/capabilities/on_enter.rs

@@ -119,12 +119,13 @@ mod tests {
         }
     }
 
-    #[test]
-    fn get_comment_workspace_edit_double_slash_indented() {
+    #[tokio::test]
+    async fn get_comment_workspace_edit_double_slash_indented() {
         let path = get_absolute_path("sway-lsp/tests/fixtures/diagnostics/dead_code/src/main.sw");
         let uri = Url::from_file_path(path.clone()).unwrap();
-        let text_document =
-            TextDocument::build_from_path(path.as_str()).expect("failed to build document");
+        let text_document = TextDocument::build_from_path(path.as_str())
+            .await
+            .expect("failed to build document");
         let params = OnEnterParams {
             text_document: TextDocumentIdentifier { uri },
             content_changes: vec![TextDocumentContentChangeEvent {
@@ -156,12 +157,13 @@ mod tests {
         assert_text_edit(&edits[0].edits[0], "// ".to_string(), 48, 4);
     }
 
-    #[test]
-    fn get_comment_workspace_edit_triple_slash_paste() {
+    #[tokio::test]
+    async fn get_comment_workspace_edit_triple_slash_paste() {
         let path = get_absolute_path("sway-lsp/tests/fixtures/diagnostics/dead_code/src/main.sw");
         let uri = Url::from_file_path(path.clone()).unwrap();
-        let text_document =
-            TextDocument::build_from_path(path.as_str()).expect("failed to build document");
+        let text_document = TextDocument::build_from_path(path.as_str())
+            .await
+            .expect("failed to build document");
         let params = OnEnterParams {
             text_document: TextDocumentIdentifier { uri },
             content_changes: vec![TextDocumentContentChangeEvent {

sway-lsp/src/core/document.rs

@@ -5,6 +5,7 @@ use crate::{
 };
 use lsp_types::{Position, Range, TextDocumentContentChangeEvent, Url};
 use ropey::Rope;
+use tokio::fs::File;
 
 #[derive(Debug, Clone)]
 pub struct TextDocument {
@@ -17,8 +18,9 @@ pub struct TextDocument {
 }
 
 impl TextDocument {
-    pub fn build_from_path(path: &str) -> Result<Self, DocumentError> {
-        std::fs::read_to_string(path)
+    pub async fn build_from_path(path: &str) -> Result<Self, DocumentError> {
+        tokio::fs::read_to_string(path)
+            .await
             .map(|content| Self {
                 language_id: "sway".into(),
                 version: 1,
@@ -109,33 +111,39 @@ impl TextDocument {
 /// Marks the specified file as "dirty" by creating a corresponding flag file.
 ///
 /// This function ensures the necessary directory structure exists before creating the flag file.
-pub fn mark_file_as_dirty(uri: &Url) -> Result<(), LanguageServerError> {
+pub async fn mark_file_as_dirty(uri: &Url) -> Result<(), LanguageServerError> {
     let path = document::get_path_from_url(uri)?;
     let dirty_file_path = forc_util::is_dirty_path(&path);
     if let Some(dir) = dirty_file_path.parent() {
         // Ensure the directory exists
-        std::fs::create_dir_all(dir).map_err(|_| DirectoryError::LspLocksDirFailed)?;
+        tokio::fs::create_dir_all(dir)
+            .await
+            .map_err(|_| DirectoryError::LspLocksDirFailed)?;
     }
     // Create an empty "dirty" file
-    std::fs::File::create(&dirty_file_path).map_err(|err| DocumentError::UnableToCreateFile {
-        path: uri.path().to_string(),
-        err: err.to_string(),
-    })?;
+    File::create(&dirty_file_path)
+        .await
+        .map_err(|err| DocumentError::UnableToCreateFile {
+            path: uri.path().to_string(),
+            err: err.to_string(),
+        })?;
     Ok(())
 }
 
 /// Removes the corresponding flag file for the specifed Url.
 ///
 /// If the flag file does not exist, this function will do nothing.
-pub fn remove_dirty_flag(uri: &Url) -> Result<(), LanguageServerError> {
+pub async fn remove_dirty_flag(uri: &Url) -> Result<(), LanguageServerError> {
     let path = document::get_path_from_url(uri)?;
     let dirty_file_path = forc_util::is_dirty_path(&path);
     if dirty_file_path.exists() {
         // Remove the "dirty" file
-        std::fs::remove_file(dirty_file_path).map_err(|err| DocumentError::UnableToRemoveFile {
-            path: uri.path().to_string(),
-            err: err.to_string(),
-        })?;
+        tokio::fs::remove_file(dirty_file_path)
+            .await
+            .map_err(|err| DocumentError::UnableToRemoveFile {
+                path: uri.path().to_string(),
+                err: err.to_string(),
+            })?;
     }
     Ok(())
 }
@@ -152,17 +160,19 @@ mod tests {
     use super::*;
     use sway_lsp_test_utils::get_absolute_path;
 
-    #[test]
-    fn build_from_path_returns_text_document() {
+    #[tokio::test]
+    async fn build_from_path_returns_text_document() {
         let path = get_absolute_path("sway-lsp/tests/fixtures/cats.txt");
-        let result = TextDocument::build_from_path(&path);
+        let result = TextDocument::build_from_path(&path).await;
         assert!(result.is_ok(), "result = {result:?}");
     }
 
-    #[test]
-    fn build_from_path_returns_document_not_found_error() {
+    #[tokio::test]
+    async fn build_from_path_returns_document_not_found_error() {
         let path = get_absolute_path("not/a/real/file/path");
-        let result = TextDocument::build_from_path(&path).expect_err("expected DocumentNotFound");
+        let result = TextDocument::build_from_path(&path)
+            .await
+            .expect_err("expected DocumentNotFound");
         assert_eq!(result, DocumentError::DocumentNotFound { path });
     }
 }

sway-lsp/src/core/session.rs

@@ -25,14 +25,7 @@ use lsp_types::{
 use parking_lot::RwLock;
 use pkg::{manifest::ManifestFile, BuildPlan};
 use rayon::iter::{ParallelBridge, ParallelIterator};
-use std::{
-    fs::File,
-    io::Write,
-    ops::Deref,
-    path::PathBuf,
-    sync::{atomic::Ordering, Arc},
-    vec,
-};
+use std::{ops::Deref, path::PathBuf, sync::Arc};
 use sway_core::{
     decl_engine::DeclEngine,
     language::{
@@ -46,7 +39,7 @@ use sway_core::{
 use sway_error::{error::CompileError, handler::Handler, warning::CompileWarning};
 use sway_types::{SourceEngine, SourceId, Spanned};
 use sway_utils::{helpers::get_sway_files, PerformanceData};
-use tokio::sync::Semaphore;
+use tokio::{fs::File, io::AsyncWriteExt, sync::Semaphore};
 
 pub type Documents = DashMap<String, TextDocument>;
 pub type ProjectDirectory = PathBuf;
@@ -111,26 +104,23 @@ impl Session {
         }
     }
 
-    pub fn init(&self, uri: &Url) -> Result<ProjectDirectory, LanguageServerError> {
+    pub async fn init(&self, uri: &Url) -> Result<ProjectDirectory, LanguageServerError> {
         let manifest_dir = PathBuf::from(uri.path());
         // Create a new temp dir that clones the current workspace
         // and store manifest and temp paths
         self.sync.create_temp_dir_from_workspace(&manifest_dir)?;
         self.sync.clone_manifest_dir_to_temp()?;
         // iterate over the project dir, parse all sway files
-        let _ = self.store_sway_files();
+        let _ = self.store_sway_files().await;
         self.sync.watch_and_sync_manifest();
         self.sync.manifest_dir().map_err(Into::into)
     }
 
     pub fn shutdown(&self) {
-        // Set the should_end flag to true
-        self.sync.should_end.store(true, Ordering::Relaxed);
-
-        // Wait for the thread to finish
-        let mut join_handle_option = self.sync.notify_join_handle.write();
-        if let Some(join_handle) = std::mem::take(&mut *join_handle_option) {
-            let _ = join_handle.join();
+        // shutdown the thread watching the manifest file
+        let handle = self.sync.notify_join_handle.read();
+        if let Some(join_handle) = &*handle {
+            join_handle.abort();
         }
 
         // Delete the temporary directory.
@@ -276,16 +266,16 @@ impl Session {
             .map(|page_text_edit| vec![page_text_edit])
     }
 
-    pub fn handle_open_file(&self, uri: &Url) {
+    pub async fn handle_open_file(&self, uri: &Url) {
         if !self.documents.contains_key(uri.path()) {
-            if let Ok(text_document) = TextDocument::build_from_path(uri.path()) {
+            if let Ok(text_document) = TextDocument::build_from_path(uri.path()).await {
                 let _ = self.store_document(text_document);
             }
         }
     }
 
-    /// Writes the changes to the file and updates the document.
-    pub fn write_changes_to_file(
+    /// Asynchronously writes the changes to the file and updates the document.
+    pub async fn write_changes_to_file(
         &self,
         uri: &Url,
         changes: Vec<TextDocumentContentChangeEvent>,
@@ -295,15 +285,22 @@ impl Session {
                 path: uri.path().to_string(),
             }
         })?;
+
         let mut file =
-            File::create(uri.path()).map_err(|err| DocumentError::UnableToCreateFile {
+            File::create(uri.path())
+                .await
+                .map_err(|err| DocumentError::UnableToCreateFile {
+                    path: uri.path().to_string(),
+                    err: err.to_string(),
+                })?;
+
+        file.write_all(src.as_bytes())
+            .await
+            .map_err(|err| DocumentError::UnableToWriteFile {
                 path: uri.path().to_string(),
                 err: err.to_string(),
             })?;
-        writeln!(&mut file, "{src}").map_err(|err| DocumentError::UnableToWriteFile {
-            path: uri.path().to_string(),
-            err: err.to_string(),
-        })?;
+
         Ok(())
     }
 
@@ -399,11 +396,11 @@ impl Session {
     }
 
     /// Populate [Documents] with sway files found in the workspace.
-    fn store_sway_files(&self) -> Result<(), LanguageServerError> {
+    async fn store_sway_files(&self) -> Result<(), LanguageServerError> {
         let temp_dir = self.sync.temp_dir()?;
         // Store the documents.
         for path in get_sway_files(temp_dir).iter().filter_map(|fp| fp.to_str()) {
-            self.store_document(TextDocument::build_from_path(path)?)?;
+            self.store_document(TextDocument::build_from_path(path).await?)?;
         }
         Ok(())
     }
@@ -594,22 +591,22 @@ mod tests {
     use super::*;
     use sway_lsp_test_utils::{get_absolute_path, get_url};
 
-    #[test]
-    fn store_document_returns_empty_tuple() {
+    #[tokio::test]
+    async fn store_document_returns_empty_tuple() {
         let session = Session::new();
         let path = get_absolute_path("sway-lsp/tests/fixtures/cats.txt");
-        let document = TextDocument::build_from_path(&path).unwrap();
+        let document = TextDocument::build_from_path(&path).await.unwrap();
         let result = Session::store_document(&session, document);
         assert!(result.is_ok());
     }
 
-    #[test]
-    fn store_document_returns_document_already_stored_error() {
+    #[tokio::test]
+    async fn store_document_returns_document_already_stored_error() {
         let session = Session::new();
         let path = get_absolute_path("sway-lsp/tests/fixtures/cats.txt");
-        let document = TextDocument::build_from_path(&path).unwrap();
+        let document = TextDocument::build_from_path(&path).await.unwrap();
         Session::store_document(&session, document).expect("expected successfully stored");
-        let document = TextDocument::build_from_path(&path).unwrap();
+        let document = TextDocument::build_from_path(&path).await.unwrap();
         let result = Session::store_document(&session, document)
             .expect_err("expected DocumentAlreadyStored");
         assert_eq!(result, DocumentError::DocumentAlreadyStored { path });