spawn separate task for LSP server loop, use MPSC channel

Reviewed By: josephsavona

Differential Revision: D21311187

fbshipit-source-id: 6c0117aaf616098903c403cc2dfc27d96425b891

compiler/crates/relay-compiler/src/compiler.rs

@@ -69,8 +69,6 @@ impl Compiler {
         &self,
         optional_serialized_state_path: Option<PathBuf>,
     ) -> Result<CompilerState> {
-        use env_logger::Env;
-        env_logger::from_env(Env::default().default_filter_or("info, error, warn")).init();
         let file_source = FileSource::connect(&self.config).await?;
         let (mut compiler_state, _) = self
             .create_compiler_state_and_file_source_result(

compiler/crates/relay-lsp/src/main.rs

@@ -12,11 +12,17 @@ mod server;
 use lsp_server::Connection;
 use std::error::Error;
 
+use env_logger::Env;
+use log::info;
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error + Sync + Send>> {
+    env_logger::from_env(Env::default().default_filter_or("info, warn, error, debug")).init();
     let (connection, io_handles) = Connection::stdio();
+    info!("Initialized stdio transport layer");
     let params = server::initialize(&connection)?;
-    server::run(&connection, params).await?;
+    info!("JSON-RPC handshake completed");
+    server::run(connection, params).await?;
     io_handles.join()?;
     Ok(())
 }

compiler/crates/relay-lsp/src/server.rs

@@ -8,8 +8,9 @@
 use std::error::Error;
 
 use lsp_types::{
-    notification::{Notification, ShowMessage},
+    notification::{DidOpenTextDocument, Notification, ShowMessage},
     InitializeParams, MessageType, ServerCapabilities, ShowMessageParams,
+    TextDocumentSyncCapability, TextDocumentSyncKind,
 };
 
 use lsp_server::{Connection, Message, Notification as ServerNotification};
@@ -21,60 +22,102 @@ use std::path::PathBuf;
 
 use log::info;
 
+use tokio::sync::mpsc;
+
 /// Initializes an LSP connection, handling the `initize` message and `initialized` notification
 /// handshake.
 pub fn initialize(
     connection: &Connection,
 ) -> Result<InitializeParams, Box<dyn Error + Sync + Send>> {
-    let server_capabilities = serde_json::to_value(&ServerCapabilities::default()).unwrap();
+    let mut server_capabilities = ServerCapabilities::default();
+    // Enable text document syncing so we can know when files are opened/changed/saved/closed
+    server_capabilities.text_document_sync =
+        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));
+
+    let server_capabilities = serde_json::to_value(&server_capabilities).unwrap();
     let params = connection.initialize(server_capabilities)?;
     let params: InitializeParams = serde_json::from_value(params).unwrap();
     Ok(params)
 }
 
+#[derive(Debug)]
+pub enum CompilerMessage {
+    Initialize,
+}
+
 /// Run the main server loop
 pub async fn run(
-    connection: &Connection,
+    connection: Connection,
     _params: InitializeParams,
 ) -> Result<(), Box<dyn Error + Sync + Send>> {
-    show_info_message("Relay Language Server Started", connection)?;
-
-    // TODO(brandondail) don't hardcode the test project config here
-    let home = std::env::var("HOME").unwrap();
-    let config_path = PathBuf::from(format!(
-        "{}/fbsource/fbcode/relay/config/config.test.json",
-        home
-    ));
-
-    let root_dir = PathBuf::from(format!("{}/fbsource", home));
-    let mut config = Config::load(root_dir, config_path).unwrap();
-
-    // Don't write artifacts by default
-    config.write_artifacts = false;
-
-    info!("Compiler config: {:#?}", config);
-
-    let compiler = Compiler::new(config);
-
-    match compiler.compile(None).await {
-        Ok(_) => {
-            info!("Compiled project successfully");
-        }
-        Err(_) => {
-            info!("Failed to compile project");
-        }
-    }
-
-    for msg in &connection.receiver {
-        match msg {
-            Message::Request(req) => {
-                info!("Request: {:?}", req);
+    show_info_message("Relay Language Server Started!", &connection)?;
+    info!("Running language server");
+
+    // We use an MPSC channel to communicate between the LSP server loop and
+    // the compiler. That way running compiler doesn't block us from receiving
+    // LSP messages.
+    let (mut tx, mut rx) = mpsc::channel::<CompilerMessage>(100);
+
+    let receiver = connection.receiver.clone();
+
+    // Thread for the LSP message loop
+    tokio::spawn(async move {
+        for msg in receiver {
+            match msg {
+                Message::Request(req) => {
+                    info!("Request: {:#?}", req);
+                }
+                Message::Response(resp) => {
+                    info!("Request: {:#?}", resp);
+                }
+                Message::Notification(notif) => {
+                    info!("Got a notification: {:#?}", notif);
+
+                    if notif.method == DidOpenTextDocument::METHOD {
+                        // Lazily start the compiler once a relevant file is opened
+                        if tx.send(CompilerMessage::Initialize).await.is_err() {
+                            return;
+                        }
+                    }
+                }
             }
-            Message::Response(resp) => {
-                info!("Request: {:?}", resp);
-            }
-            Message::Notification(notif) => {
-                info!("Notification: {:?}", notif);
+        }
+    });
+
+    info!("Starting to wait for receiver messages");
+
+    let mut is_compiler_running = false;
+
+    while let Some(res) = rx.recv().await {
+        match res {
+            CompilerMessage::Initialize => {
+                if !is_compiler_running {
+                    is_compiler_running = true;
+                    info!("Initializing compiler...");
+
+                    // TODO(brandondail) don't hardcode the test project config here
+                    let home = std::env::var("HOME").unwrap();
+                    let config_path = PathBuf::from(format!(
+                        "{}/fbsource/fbcode/relay/config/config.test.json",
+                        home
+                    ));
+
+                    let root_dir = PathBuf::from(format!("{}/fbsource", home));
+                    let mut config = Config::load(root_dir, config_path).unwrap();
+
+                    // Don't write artifacts by default
+                    config.write_artifacts = false;
+
+                    let compiler = Compiler::new(config);
+
+                    match compiler.compile(None).await {
+                        Ok(_) => info!("Compiled successfully"),
+                        Err(_) => {
+                            show_info_message("Failed to compile", &connection)?;
+                            info!("Error!");
+                        }
+                    }
+                }
             }
         }
     }