Merge branch 'lsp-async-init'

3 files changed, 213 insertions, 116 deletions

diff --git a/helix-lsp/src/client.rs b/helix-lsp/src/client.rs
index d0a8183f..f2bb0059 100644
--- a/helix-lsp/src/client.rs
+++ b/helix-lsp/src/client.rs
@@ -9,11 +9,17 @@ use lsp_types as lsp;
 use serde_json::Value;
 use std::future::Future;
 use std::process::Stdio;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{
+    atomic::{AtomicU64, Ordering},
+    Arc,
+};
 use tokio::{
     io::{BufReader, BufWriter},
     process::{Child, Command},
-    sync::mpsc::{channel, UnboundedReceiver, UnboundedSender},
+    sync::{
+        mpsc::{channel, UnboundedReceiver, UnboundedSender},
+        Notify, OnceCell,
+    },
 };
 
 #[derive(Debug)]
@@ -22,18 +28,19 @@ pub struct Client {
     _process: Child,
     server_tx: UnboundedSender<Payload>,
     request_counter: AtomicU64,
-    capabilities: Option<lsp::ServerCapabilities>,
+    pub(crate) capabilities: OnceCell<lsp::ServerCapabilities>,
     offset_encoding: OffsetEncoding,
     config: Option<Value>,
 }
 
 impl Client {
+    #[allow(clippy::type_complexity)]
     pub fn start(
         cmd: &str,
         args: &[String],
         config: Option<Value>,
         id: usize,
-    ) -> Result<(Self, UnboundedReceiver<(usize, Call)>)> {
+    ) -> Result<(Self, UnboundedReceiver<(usize, Call)>, Arc<Notify>)> {
         let process = Command::new(cmd)
             .args(args)
             .stdin(Stdio::piped())
@@ -50,22 +57,20 @@ impl Client {
         let reader = BufReader::new(process.stdout.take().expect("Failed to open stdout"));
         let stderr = BufReader::new(process.stderr.take().expect("Failed to open stderr"));
 
-        let (server_rx, server_tx) = Transport::start(reader, writer, stderr, id);
+        let (server_rx, server_tx, initialize_notify) =
+            Transport::start(reader, writer, stderr, id);
 
         let client = Self {
             id,
             _process: process,
             server_tx,
             request_counter: AtomicU64::new(0),
-            capabilities: None,
+            capabilities: OnceCell::new(),
             offset_encoding: OffsetEncoding::Utf8,
             config,
         };
 
-        // TODO: async client.initialize()
-        // maybe use an arc<atomic> flag
-
-        Ok((client, server_rx))
+        Ok((client, server_rx, initialize_notify))
     }
 
     pub fn id(&self) -> usize {
@@ -88,9 +93,13 @@ impl Client {
         }
     }
 
+    pub fn is_initialized(&self) -> bool {
+        self.capabilities.get().is_some()
+    }
+
     pub fn capabilities(&self) -> &lsp::ServerCapabilities {
         self.capabilities
-            .as_ref()
+            .get()
             .expect("language server not yet initialized!")
     }
 
@@ -143,7 +152,8 @@ impl Client {
                 })
                 .map_err(|e| Error::Other(e.into()))?;
 
-            timeout(Duration::from_secs(2), rx.recv())
+            // TODO: specifiable timeout, delay other calls until initialize success
+            timeout(Duration::from_secs(20), rx.recv())
                 .await
                 .map_err(|_| Error::Timeout)? // return Timeout
                 .ok_or(Error::StreamClosed)?
@@ -151,7 +161,7 @@ impl Client {
     }
 
     /// Send a RPC notification to the language server.
-    fn notify<R: lsp::notification::Notification>(
+    pub fn notify<R: lsp::notification::Notification>(
         &self,
         params: R::Params,
     ) -> impl Future<Output = Result<()>>
@@ -213,7 +223,7 @@ impl Client {
     // General messages
     // -------------------------------------------------------------------------------------------
 
-    pub(crate) async fn initialize(&mut self) -> Result<()> {
+    pub(crate) async fn initialize(&self) -> Result<lsp::InitializeResult> {
         // TODO: delay any requests that are triggered prior to initialize
         let root = find_root(None).and_then(|root| lsp::Url::from_file_path(root).ok());
 
@@ -281,14 +291,7 @@ impl Client {
             locale: None, // TODO
         };
 
-        let response = self.request::<lsp::request::Initialize>(params).await?;
-        self.capabilities = Some(response.capabilities);
-
-        // next up, notify<initialized>
-        self.notify::<lsp::notification::Initialized>(lsp::InitializedParams {})
-            .await?;
-
-        Ok(())
+        self.request::<lsp::request::Initialize>(params).await
     }
 
     pub async fn shutdown(&self) -> Result<()> {
@@ -445,7 +448,7 @@ impl Client {
     ) -> Option<impl Future<Output = Result<()>>> {
         // figure out what kind of sync the server supports
 
-        let capabilities = self.capabilities.as_ref().unwrap();
+        let capabilities = self.capabilities.get().unwrap();
 
         let sync_capabilities = match capabilities.text_document_sync {
             Some(lsp::TextDocumentSyncCapability::Kind(kind))
@@ -463,7 +466,7 @@ impl Client {
                     // range = None -> whole document
                     range: None,        //Some(Range)
                     range_length: None, // u64 apparently deprecated
-                    text: "".to_string(),
+                    text: new_text.to_string(),
                 }]
             }
             lsp::TextDocumentSyncKind::Incremental => {
@@ -491,12 +494,12 @@ impl Client {
 
     // will_save / will_save_wait_until
 
-    pub async fn text_document_did_save(
+    pub fn text_document_did_save(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         text: &Rope,
-    ) -> Result<()> {
-        let capabilities = self.capabilities.as_ref().unwrap();
+    ) -> Option<impl Future<Output = Result<()>>> {
+        let capabilities = self.capabilities.get().unwrap();
 
         let include_text = match &capabilities.text_document_sync {
             Some(lsp::TextDocumentSyncCapability::Options(lsp::TextDocumentSyncOptions {
@@ -508,17 +511,18 @@ impl Client {
                     include_text,
                 }) => include_text.unwrap_or(false),
                 // Supported(false)
-                _ => return Ok(()),
+                _ => return None,
             },
             // unsupported
-            _ => return Ok(()),
+            _ => return None,
         };
 
-        self.notify::<lsp::notification::DidSaveTextDocument>(lsp::DidSaveTextDocumentParams {
-            text_document,
-            text: include_text.then(|| text.into()),
-        })
-        .await
+        Some(self.notify::<lsp::notification::DidSaveTextDocument>(
+            lsp::DidSaveTextDocumentParams {
+                text_document,
+                text: include_text.then(|| text.into()),
+            },
+        ))
     }
 
     pub fn completion(
@@ -584,19 +588,19 @@ impl Client {
 
     // formatting
 
-    pub async fn text_document_formatting(
+    pub fn text_document_formatting(
         &self,
         text_document: lsp::TextDocumentIdentifier,
         options: lsp::FormattingOptions,
         work_done_token: Option<lsp::ProgressToken>,
-    ) -> anyhow::Result<Vec<lsp::TextEdit>> {
-        let capabilities = self.capabilities.as_ref().unwrap();
+    ) -> Option<impl Future<Output = Result<Vec<lsp::TextEdit>>>> {
+        let capabilities = self.capabilities.get().unwrap();
 
         // check if we're able to format
         match capabilities.document_formatting_provider {
             Some(lsp::OneOf::Left(true)) | Some(lsp::OneOf::Right(_)) => (),
             // None | Some(false)
-            _ => return Ok(Vec::new()),
+            _ => return None,
         };
         // TODO: return err::unavailable so we can fall back to tree sitter formatting
 
@@ -606,9 +610,13 @@ impl Client {
             work_done_progress_params: lsp::WorkDoneProgressParams { work_done_token },
         };
 
-        let response = self.request::<lsp::request::Formatting>(params).await?;
+        let request = self.call::<lsp::request::Formatting>(params);
 
-        Ok(response.unwrap_or_default())
+        Some(async move {
+            let json = request.await?;
+            let response: Vec<lsp::TextEdit> = serde_json::from_value(json)?;
+            Ok(response)
+        })
     }
 
     pub async fn text_document_range_formatting(
@@ -618,7 +626,7 @@ impl Client {
         options: lsp::FormattingOptions,
         work_done_token: Option<lsp::ProgressToken>,
     ) -> anyhow::Result<Vec<lsp::TextEdit>> {
-        let capabilities = self.capabilities.as_ref().unwrap();
+        let capabilities = self.capabilities.get().unwrap();
 
         // check if we're able to format
         match capabilities.document_range_formatting_provider {
diff --git a/helix-lsp/src/lib.rs b/helix-lsp/src/lib.rs
index 72606b70..7357c885 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -226,6 +226,8 @@ impl MethodCall {
 
 #[derive(Debug, PartialEq, Clone)]
 pub enum Notification {
+    // we inject this notification to signal the LSP is ready
+    Initialized,
     PublishDiagnostics(lsp::PublishDiagnosticsParams),
     ShowMessage(lsp::ShowMessageParams),
     LogMessage(lsp::LogMessageParams),
@@ -237,6 +239,7 @@ impl Notification {
         use lsp::notification::Notification as _;
 
         let notification = match method {
+            lsp::notification::Initialized::METHOD => Self::Initialized,
             lsp::notification::PublishDiagnostics::METHOD => {
                 let params: lsp::PublishDiagnosticsParams = params
                     .parse()
@@ -294,7 +297,7 @@ impl Registry {
         }
     }
 
-    pub fn get_by_id(&mut self, id: usize) -> Option<&Client> {
+    pub fn get_by_id(&self, id: usize) -> Option<&Client> {
         self.inner
             .values()
             .find(|(client_id, _)| client_id == &id)
@@ -302,33 +305,52 @@ impl Registry {
     }
 
     pub fn get(&mut self, language_config: &LanguageConfiguration) -> Result<Arc<Client>> {
-        if let Some(config) = &language_config.language_server {
-            // avoid borrow issues
-            let inner = &mut self.inner;
-            let s_incoming = &mut self.incoming;
-
-            match inner.entry(language_config.scope.clone()) {
-                Entry::Occupied(entry) => Ok(entry.get().1.clone()),
-                Entry::Vacant(entry) => {
-                    // initialize a new client
-                    let id = self.counter.fetch_add(1, Ordering::Relaxed);
-                    let (mut client, incoming) = Client::start(
-                        &config.command,
-                        &config.args,
-                        serde_json::from_str(language_config.config.as_deref().unwrap_or("")).ok(),
-                        id,
-                    )?;
-                    // TODO: run this async without blocking
-                    futures_executor::block_on(client.initialize())?;
-                    s_incoming.push(UnboundedReceiverStream::new(incoming));
-                    let client = Arc::new(client);
-
-                    entry.insert((id, client.clone()));
-                    Ok(client)
-                }
+        let config = match &language_config.language_server {
+            Some(config) => config,
+            None => return Err(Error::LspNotDefined),
+        };
+
+        match self.inner.entry(language_config.scope.clone()) {
+            Entry::Occupied(entry) => Ok(entry.get().1.clone()),
+            Entry::Vacant(entry) => {
+                // initialize a new client
+                let id = self.counter.fetch_add(1, Ordering::Relaxed);
+                let (client, incoming, initialize_notify) = Client::start(
+                    &config.command,
+                    &config.args,
+                    serde_json::from_str(language_config.config.as_deref().unwrap_or("")).ok(),
+                    id,
+                )?;
+                self.incoming.push(UnboundedReceiverStream::new(incoming));
+                let client = Arc::new(client);
+
+                // Initialize the client asynchronously
+                let _client = client.clone();
+                tokio::spawn(async move {
+                    use futures_util::TryFutureExt;
+                    let value = _client
+                        .capabilities
+                        .get_or_try_init(|| {
+                            _client
+                                .initialize()
+                                .map_ok(|response| response.capabilities)
+                        })
+                        .await;
+
+                    value.expect("failed to initialize capabilities");
+
+                    // next up, notify<initialized>
+                    _client
+                        .notify::<lsp::notification::Initialized>(lsp::InitializedParams {})
+                        .await
+                        .unwrap();
+
+                    initialize_notify.notify_one();
+                });
+
+                entry.insert((id, client.clone()));
+                Ok(client)
             }
-        } else {
-            Err(Error::LspNotDefined)
         }
     }
 
@@ -415,32 +437,6 @@ impl LspProgressMap {
     }
 }
 
-// REGISTRY = HashMap<LanguageId, Lazy/OnceCell<Arc<RwLock<Client>>>
-// spawn one server per language type, need to spawn one per workspace if server doesn't support
-// workspaces
-//
-// could also be a client per root dir
-//
-// storing a copy of Option<Arc<RwLock<Client>>> on Document would make the LSP client easily
-// accessible during edit/save callbacks
-//
-// the event loop needs to process all incoming streams, maybe we can just have that be a separate
-// task that's continually running and store the state on the client, then use read lock to
-// retrieve data during render
-// -> PROBLEM: how do you trigger an update on the editor side when data updates?
-//
-// -> The data updates should pull all events until we run out so we don't frequently re-render
-//
-//
-// v2:
-//
-// there should be a registry of lsp clients, one per language type (or workspace).
-// the clients should lazy init on first access
-// the client.initialize() should be called async and we buffer any requests until that completes
-// there needs to be a way to process incoming lsp messages from all clients.
-//  -> notifications need to be dispatched to wherever
-//  -> requests need to generate a reply and travel back to the same lsp!
-
 #[cfg(test)]
 mod tests {
     use super::{lsp, util::*, OffsetEncoding};
diff --git a/helix-lsp/src/transport.rs b/helix-lsp/src/transport.rs
index 67b7b48f..6e28094d 100644
--- a/helix-lsp/src/transport.rs
+++ b/helix-lsp/src/transport.rs
@@ -1,7 +1,7 @@
-use crate::Result;
+use crate::{Error, Result};
 use anyhow::Context;
 use jsonrpc_core as jsonrpc;
-use log::{debug, error, info, warn};
+use log::{error, info};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
@@ -11,7 +11,7 @@ use tokio::{
     process::{ChildStderr, ChildStdin, ChildStdout},
     sync::{
         mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender},
-        Mutex,
+        Mutex, Notify,
     },
 };
 
@@ -51,9 +51,11 @@ impl Transport {
     ) -> (
         UnboundedReceiver<(usize, jsonrpc::Call)>,
         UnboundedSender<Payload>,
+        Arc<Notify>,
     ) {
         let (client_tx, rx) = unbounded_channel();
         let (tx, client_rx) = unbounded_channel();
+        let notify = Arc::new(Notify::new());
 
         let transport = Self {
             id,
@@ -62,11 +64,21 @@ impl Transport {
 
         let transport = Arc::new(transport);
 
-        tokio::spawn(Self::recv(transport.clone(), server_stdout, client_tx));
+        tokio::spawn(Self::recv(
+            transport.clone(),
+            server_stdout,
+            client_tx.clone(),
+        ));
         tokio::spawn(Self::err(transport.clone(), server_stderr));
-        tokio::spawn(Self::send(transport, server_stdin, client_rx));
-
-        (rx, tx)
+        tokio::spawn(Self::send(
+            transport,
+            server_stdin,
+            client_tx,
+            client_rx,
+            notify.clone(),
+        ));
+
+        (rx, tx, notify)
     }
 
     async fn recv_server_message(
@@ -76,14 +88,18 @@ impl Transport {
         let mut content_length = None;
         loop {
             buffer.truncate(0);
-            reader.read_line(buffer).await?;
-            let header = buffer.trim();
+            if reader.read_line(buffer).await? == 0 {
+                return Err(Error::StreamClosed);
+            };
+
+            // debug!("<- header {:?}", buffer);
 
-            if header.is_empty() {
+            if buffer == "\r\n" {
+                // look for an empty CRLF line
                 break;
             }
 
-            debug!("<- header {}", header);
+            let header = buffer.trim();
 
             let parts = header.split_once(": ");
 
@@ -96,7 +112,8 @@ impl Transport {
                     // Workaround: Some non-conformant language servers will output logging and other garbage
                     // into the same stream as JSON-RPC messages. This can also happen from shell scripts that spawn
                     // the server. Skip such lines and log a warning.
-                    warn!("Failed to parse header: {:?}", header);
+
+                    // warn!("Failed to parse header: {:?}", header);
                 }
             }
         }
@@ -121,8 +138,10 @@ impl Transport {
         buffer: &mut String,
     ) -> Result<()> {
         buffer.truncate(0);
-        err.read_line(buffer).await?;
-        error!("err <- {}", buffer);
+        if err.read_line(buffer).await? == 0 {
+            return Err(Error::StreamClosed);
+        };
+        error!("err <- {:?}", buffer);
 
         Ok(())
     }
@@ -255,16 +274,90 @@ impl Transport {
     async fn send(
         transport: Arc<Self>,
         mut server_stdin: BufWriter<ChildStdin>,
+        client_tx: UnboundedSender<(usize, jsonrpc::Call)>,
         mut client_rx: UnboundedReceiver<Payload>,
+        initialize_notify: Arc<Notify>,
     ) {
-        while let Some(msg) = client_rx.recv().await {
-            match transport
-                .send_payload_to_server(&mut server_stdin, msg)
-                .await
-            {
-                Ok(_) => {}
-                Err(err) => {
-                    error!("err: <- {:?}", err);
+        let mut pending_messages: Vec<Payload> = Vec::new();
+        let mut is_pending = true;
+
+        // Determine if a message is allowed to be sent early
+        fn is_initialize(payload: &Payload) -> bool {
+            use lsp_types::{
+                notification::{Initialized, Notification},
+                request::{Initialize, Request},
+            };
+            match payload {
+                Payload::Request {
+                    value: jsonrpc::MethodCall { method, .. },
+                    ..
+                } if method == Initialize::METHOD => true,
+                Payload::Notification(jsonrpc::Notification { method, .. })
+                    if method == Initialized::METHOD =>
+                {
+                    true
+                }
+                _ => false,
+            }
+        }
+
+        // TODO: events that use capabilities need to do the right thing
+
+        loop {
+            tokio::select! {
+                biased;
+                _ = initialize_notify.notified() => { // TODO: notified is technically not cancellation safe
+                    // server successfully initialized
+                    is_pending = false;
+
+                    use lsp_types::notification::Notification;
+                    // Hack: inject an initialized notification so we trigger code that needs to happen after init
+                    let notification = ServerMessage::Call(jsonrpc::Call::Notification(jsonrpc::Notification {
+                        jsonrpc: None,
+
+                        method: lsp_types::notification::Initialized::METHOD.to_string(),
+                        params: jsonrpc::Params::None,
+                    }));
+                    match transport.process_server_message(&client_tx, notification).await {
+                        Ok(_) => {}
+                        Err(err) => {
+                            error!("err: <- {:?}", err);
+                        }
+                    }
+
+                    // drain the pending queue and send payloads to server
+                    for msg in pending_messages.drain(..) {
+                        log::info!("Draining pending message {:?}", msg);
+                        match transport.send_payload_to_server(&mut server_stdin, msg).await {
+                            Ok(_) => {}
+                            Err(err) => {
+                                error!("err: <- {:?}", err);
+                            }
+                        }
+                    }
+                }
+                msg = client_rx.recv() => {
+                    if let Some(msg) = msg {
+                        if is_pending && !is_initialize(&msg) {
+                            // ignore notifications
+                            if let Payload::Notification(_) = msg {
+                                continue;
+                            }
+
+                            log::info!("Language server not initialized, delaying request");
+                            pending_messages.push(msg);
+                        } else {
+                            match transport.send_payload_to_server(&mut server_stdin, msg).await {
+                                Ok(_) => {}
+                                Err(err) => {
+                                    error!("err: <- {:?}", err);
+                                }
+                            }
+                        }
+                    } else {
+                        // channel closed
+                        break;
+                    }
                 }
             }
         }