From c3a58cdadd8be85b79d773122e807862a3da3a2f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bla=C5=BE=20Hrastnik?= <blaz@mxxn.io>
Date: Tue, 31 Aug 2021 16:03:06 +0900
Subject: [PATCH] lsp: Refactor capabilities as an async OnceCell

First step in making LSP init asynchronous
---
 helix-lsp/src/client.rs | 35 ++++++++++++++---------------------
 helix-lsp/src/lib.rs    | 29 ++++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/helix-lsp/src/client.rs b/helix-lsp/src/client.rs
index d0a8183f..87078c69 100644
--- a/helix-lsp/src/client.rs
+++ b/helix-lsp/src/client.rs
@@ -13,7 +13,10 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use tokio::{
     io::{BufReader, BufWriter},
     process::{Child, Command},
-    sync::mpsc::{channel, UnboundedReceiver, UnboundedSender},
+    sync::{
+        mpsc::{channel, UnboundedReceiver, UnboundedSender},
+        OnceCell,
+    },
 };
 
 #[derive(Debug)]
@@ -22,7 +25,7 @@ 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>,
 }
@@ -57,14 +60,11 @@ impl Client {
             _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))
     }
 
@@ -90,7 +90,7 @@ impl Client {
 
     pub fn capabilities(&self) -> &lsp::ServerCapabilities {
         self.capabilities
-            .as_ref()
+            .get()
             .expect("language server not yet initialized!")
     }
 
@@ -151,7 +151,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 +213,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 +281,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 +438,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))
@@ -496,7 +489,7 @@ impl Client {
         text_document: lsp::TextDocumentIdentifier,
         text: &Rope,
     ) -> Result<()> {
-        let capabilities = self.capabilities.as_ref().unwrap();
+        let capabilities = self.capabilities.get().unwrap();
 
         let include_text = match &capabilities.text_document_sync {
             Some(lsp::TextDocumentSyncCapability::Options(lsp::TextDocumentSyncOptions {
@@ -590,7 +583,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_formatting_provider {
@@ -618,7 +611,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..a118239f 100644
--- a/helix-lsp/src/lib.rs
+++ b/helix-lsp/src/lib.rs
@@ -312,17 +312,40 @@ impl Registry {
                 Entry::Vacant(entry) => {
                     // initialize a new client
                     let id = self.counter.fetch_add(1, Ordering::Relaxed);
-                    let (mut client, incoming) = Client::start(
+                    let (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);
 
+                    let _client = client.clone();
+                    let initialize = 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();
+                    });
+
+                    // TODO: remove this block
+                    futures_executor::block_on(initialize).map_err(|_| anyhow::anyhow!("bail"))?;
+
                     entry.insert((id, client.clone()));
                     Ok(client)
                 }