feat: add retry times record. (#1476)

feat: add retry times record. for newline complete request, retry at least once to ensure document and sema info is updated.

Signed-off-by: he1pa <[email protected]>

kclvm/tools/src/LSP/src/dispatcher.rs

@@ -1,5 +1,5 @@
 use crossbeam_channel::Sender;
-use lsp_server::{ExtractError, Request};
+use lsp_server::{ExtractError, Request, RequestId};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::error::Error;
@@ -141,6 +141,58 @@ impl<'a> RequestDispatcher<'a> {
         Ok(self)
     }
 
+    /// Try to dispatch the event as the given Request type on the thread pool.
+    pub fn on_maybe_retry<R>(
+        &mut self,
+        compute_response_fn: fn(
+            LanguageServerSnapshot,
+            R::Params,
+            Sender<Task>,
+            RequestId,
+        ) -> anyhow::Result<R::Result>,
+    ) -> anyhow::Result<&mut Self>
+    where
+        R: lsp_types::request::Request + 'static,
+        R::Params: DeserializeOwned + 'static + Send,
+        R::Result: Serialize + 'static,
+    {
+        let (req, params) = match self.parse::<R>() {
+            Some(it) => it,
+            None => return Ok(self),
+        };
+
+        self.state.thread_pool.execute({
+            let snapshot = self.state.snapshot();
+            let sender = self.state.task_sender.clone();
+            let request_retry = self.state.request_retry.clone();
+            move || {
+                let result = compute_response_fn(snapshot, params, sender.clone(), req.id.clone());
+                match &result {
+                    Err(e)
+                        if e.downcast_ref::<LSPError>()
+                            .map_or(false, |lsp_err| matches!(lsp_err, LSPError::Retry)) =>
+                    {
+                        sender.send(Task::Retry(req.clone())).unwrap();
+                        let mut request_retry = request_retry.write();
+                        match request_retry.get_mut(&req.clone().id) {
+                            Some(t) => *t += 1,
+                            None => {
+                                request_retry.insert(req.id.clone(), 1);
+                            }
+                        }
+                    }
+                    _ => {
+                        sender
+                            .send(Task::Response(result_to_response::<R>(req.id, result)))
+                            .unwrap();
+                    }
+                }
+            }
+        });
+
+        Ok(self)
+    }
+
     /// Tries to parse the request as the specified type. If the request is of the specified type,
     /// the request is transferred and any subsequent call to this method will return None. If an
     /// error is encountered during parsing of the request parameters an error is send to the

kclvm/tools/src/LSP/src/request.rs

@@ -2,6 +2,7 @@ use anyhow::anyhow;
 use crossbeam_channel::Sender;
 
 use kclvm_sema::info::is_valid_kcl_name;
+use lsp_server::RequestId;
 use lsp_types::{Location, SemanticTokensResult, TextEdit};
 use ra_ap_vfs::{AbsPathBuf, VfsPath};
 use std::collections::HashMap;
@@ -52,7 +53,6 @@ impl LanguageServerState {
             })?
             .on::<lsp_types::request::GotoDefinition>(handle_goto_definition)?
             .on::<lsp_types::request::References>(handle_reference)?
-            .on::<lsp_types::request::Completion>(handle_completion)?
             .on::<lsp_types::request::HoverRequest>(handle_hover)?
             .on::<lsp_types::request::DocumentSymbolRequest>(handle_document_symbol)?
             .on::<lsp_types::request::CodeActionRequest>(handle_code_action)?
@@ -61,6 +61,7 @@ impl LanguageServerState {
             .on::<lsp_types::request::Rename>(handle_rename)?
             .on::<lsp_types::request::SemanticTokensFullRequest>(handle_semantic_tokens_full)?
             .on::<lsp_types::request::InlayHintRequest>(handle_inlay_hint)?
+            .on_maybe_retry::<lsp_types::request::Completion>(handle_completion)?
             .finish();
 
         Ok(())
@@ -298,6 +299,7 @@ pub(crate) fn handle_completion(
     snapshot: LanguageServerSnapshot,
     params: lsp_types::CompletionParams,
     sender: Sender<Task>,
+    id: RequestId,
 ) -> anyhow::Result<Option<lsp_types::CompletionResponse>> {
     let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
     let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
@@ -316,7 +318,9 @@ pub(crate) fn handle_completion(
         .and_then(|s| s.chars().next());
 
     if matches!(completion_trigger_character, Some('\n')) {
-        if !snapshot.verify_request_version(db.version, &path)? {
+        if snapshot.request_retry.read().get(&id).is_none() {
+            return Err(anyhow!(LSPError::Retry));
+        } else if !snapshot.verify_request_version(db.version, &path)? {
             return Err(anyhow!(LSPError::Retry));
         }
     }

kclvm/tools/src/LSP/src/state.rs

@@ -10,6 +10,7 @@ use kclvm_driver::CompileUnitOptions;
 use kclvm_parser::KCLModuleCache;
 use kclvm_sema::core::global_state::GlobalState;
 use kclvm_sema::resolver::scope::KCLScopeCache;
+use lsp_server::RequestId;
 use lsp_server::{ReqQueue, Request, Response};
 use lsp_types::Url;
 use lsp_types::{
@@ -76,6 +77,8 @@ pub(crate) struct LanguageServerState {
     pub opened_files: Arc<RwLock<HashMap<FileId, DocumentVersion>>>,
     /// The VFS loader
     pub loader: Handle<Box<dyn ra_ap_vfs::loader::Handle>, Receiver<ra_ap_vfs::loader::Message>>,
+    /// request retry time
+    pub request_retry: Arc<RwLock<HashMap<RequestId, i32>>>,
     /// The word index map
     pub word_index_map: KCLWordIndexMap,
     /// KCL parse cache
@@ -99,6 +102,8 @@ pub(crate) struct LanguageServerSnapshot {
     pub db: Arc<RwLock<HashMap<FileId, Arc<AnalysisDatabase>>>>,
     /// Documents that are currently kept in memory from the client
     pub opened_files: Arc<RwLock<HashMap<FileId, DocumentVersion>>>,
+    /// request retry time
+    pub request_retry: Arc<RwLock<HashMap<RequestId, i32>>>,
     /// The word index map
     pub word_index_map: KCLWordIndexMap,
     /// KCL parse cache
@@ -141,6 +146,7 @@ impl LanguageServerState {
             entry_cache: KCLEntryCache::default(),
             tool: Arc::new(RwLock::new(toolchain::default())),
             gs_cache: KCLGlobalStateCache::default(),
+            request_retry: Arc::new(RwLock::new(HashMap::new())),
         };
 
         let word_index_map = state.word_index_map.clone();
@@ -380,6 +386,7 @@ impl LanguageServerState {
             scope_cache: self.scope_cache.clone(),
             entry_cache: self.entry_cache.clone(),
             tool: self.tool.clone(),
+            request_retry: self.request_retry.clone(),
         }
     }