Auto merge of rust-lang#17157 - Veykril:retry, r=Veykril

fix: Don't retry position relient requests and version resolve data

Fixes rust-lang/rust-analyzer#15907
Fixes rust-lang/rust-analyzer#14839
Fixes rust-lang/rust-analyzer#16536

Author: bors

src/tools/rust-analyzer/crates/rust-analyzer/src/dispatch.rs

@@ -95,45 +95,8 @@ impl RequestDispatcher<'_> {
         self
     }
 
-    /// Dispatches a non-latency-sensitive request onto the thread pool
-    /// without retrying it if it panics.
-    pub(crate) fn on_no_retry<R>(
-        &mut self,
-        f: fn(GlobalStateSnapshot, R::Params) -> anyhow::Result<R::Result>,
-    ) -> &mut Self
-    where
-        R: lsp_types::request::Request + 'static,
-        R::Params: DeserializeOwned + panic::UnwindSafe + Send + fmt::Debug,
-        R::Result: Serialize,
-    {
-        let (req, params, panic_context) = match self.parse::<R>() {
-            Some(it) => it,
-            None => return self,
-        };
-
-        self.global_state.task_pool.handle.spawn(ThreadIntent::Worker, {
-            let world = self.global_state.snapshot();
-            move || {
-                let result = panic::catch_unwind(move || {
-                    let _pctx = stdx::panic_context::enter(panic_context);
-                    f(world, params)
-                });
-                match thread_result_to_response::<R>(req.id.clone(), result) {
-                    Ok(response) => Task::Response(response),
-                    Err(_) => Task::Response(lsp_server::Response::new_err(
-                        req.id,
-                        lsp_server::ErrorCode::ContentModified as i32,
-                        "content modified".to_owned(),
-                    )),
-                }
-            }
-        });
-
-        self
-    }
-
     /// Dispatches a non-latency-sensitive request onto the thread pool.
-    pub(crate) fn on<R>(
+    pub(crate) fn on<const ALLOW_RETRYING: bool, R>(
         &mut self,
         f: fn(GlobalStateSnapshot, R::Params) -> anyhow::Result<R::Result>,
     ) -> &mut Self
@@ -142,11 +105,11 @@ impl RequestDispatcher<'_> {
         R::Params: DeserializeOwned + panic::UnwindSafe + Send + fmt::Debug,
         R::Result: Serialize,
     {
-        self.on_with_thread_intent::<true, R>(ThreadIntent::Worker, f)
+        self.on_with_thread_intent::<true, ALLOW_RETRYING, R>(ThreadIntent::Worker, f)
     }
 
     /// Dispatches a latency-sensitive request onto the thread pool.
-    pub(crate) fn on_latency_sensitive<R>(
+    pub(crate) fn on_latency_sensitive<const ALLOW_RETRYING: bool, R>(
         &mut self,
         f: fn(GlobalStateSnapshot, R::Params) -> anyhow::Result<R::Result>,
     ) -> &mut Self
@@ -155,7 +118,7 @@ impl RequestDispatcher<'_> {
         R::Params: DeserializeOwned + panic::UnwindSafe + Send + fmt::Debug,
         R::Result: Serialize,
     {
-        self.on_with_thread_intent::<true, R>(ThreadIntent::LatencySensitive, f)
+        self.on_with_thread_intent::<true, ALLOW_RETRYING, R>(ThreadIntent::LatencySensitive, f)
     }
 
     /// Formatting requests should never block on waiting a for task thread to open up, editors will wait
@@ -170,7 +133,7 @@ impl RequestDispatcher<'_> {
         R::Params: DeserializeOwned + panic::UnwindSafe + Send + fmt::Debug,
         R::Result: Serialize,
     {
-        self.on_with_thread_intent::<false, R>(ThreadIntent::LatencySensitive, f)
+        self.on_with_thread_intent::<false, false, R>(ThreadIntent::LatencySensitive, f)
     }
 
     pub(crate) fn finish(&mut self) {
@@ -185,7 +148,7 @@ impl RequestDispatcher<'_> {
         }
     }
 
-    fn on_with_thread_intent<const MAIN_POOL: bool, R>(
+    fn on_with_thread_intent<const MAIN_POOL: bool, const ALLOW_RETRYING: bool, R>(
         &mut self,
         intent: ThreadIntent,
         f: fn(GlobalStateSnapshot, R::Params) -> anyhow::Result<R::Result>,
@@ -215,7 +178,12 @@ impl RequestDispatcher<'_> {
             });
             match thread_result_to_response::<R>(req.id.clone(), result) {
                 Ok(response) => Task::Response(response),
-                Err(_) => Task::Retry(req),
+                Err(_cancelled) if ALLOW_RETRYING => Task::Retry(req),
+                Err(_cancelled) => Task::Response(lsp_server::Response::new_err(
+                    req.id,
+                    lsp_server::ErrorCode::ContentModified as i32,
+                    "content modified".to_owned(),
+                )),
             }
         });
 

src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs

@@ -488,6 +488,10 @@ impl GlobalStateSnapshot {
         Ok(res)
     }
 
+    pub(crate) fn file_version(&self, file_id: FileId) -> Option<i32> {
+        Some(self.mem_docs.get(self.vfs_read().file_path(file_id))?.version)
+    }
+
     pub(crate) fn url_file_version(&self, url: &Url) -> Option<i32> {
         let path = from_proto::vfs_path(url).ok()?;
         Some(self.mem_docs.get(&path)?.version)

src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs

@@ -934,16 +934,18 @@ pub(crate) fn handle_related_tests(
 
 pub(crate) fn handle_completion(
     snap: GlobalStateSnapshot,
-    params: lsp_types::CompletionParams,
+    lsp_types::CompletionParams { text_document_position, context,.. }: lsp_types::CompletionParams,
 ) -> anyhow::Result<Option<lsp_types::CompletionResponse>> {
     let _p = tracing::span!(tracing::Level::INFO, "handle_completion").entered();
-    let text_document_position = params.text_document_position.clone();
-    let position = from_proto::file_position(&snap, params.text_document_position)?;
+    let mut position = from_proto::file_position(&snap, text_document_position.clone())?;
+    let line_index = snap.file_line_index(position.file_id)?;
     let completion_trigger_character =
-        params.context.and_then(|ctx| ctx.trigger_character).and_then(|s| s.chars().next());
+        context.and_then(|ctx| ctx.trigger_character).and_then(|s| s.chars().next());
 
     let source_root = snap.analysis.source_root(position.file_id)?;
     let completion_config = &snap.config.completion(Some(source_root));
+    // FIXME: We should fix up the position when retrying the cancelled request instead
+    position.offset = position.offset.min(line_index.index.len());
     let items = match snap.analysis.completions(
         completion_config,
         position,
@@ -952,10 +954,14 @@ pub(crate) fn handle_completion(
         None => return Ok(None),
         Some(items) => items,
     };
-    let line_index = snap.file_line_index(position.file_id)?;
 
-    let items =
-        to_proto::completion_items(&snap.config, &line_index, text_document_position, items);
+    let items = to_proto::completion_items(
+        &snap.config,
+        &line_index,
+        snap.file_version(position.file_id),
+        text_document_position,
+        items,
+    );
 
     let completion_list = lsp_types::CompletionList { is_incomplete: true, items };
     Ok(Some(completion_list.into()))
@@ -974,16 +980,16 @@ pub(crate) fn handle_completion_resolve(
         .into());
     }
 
-    let data = match original_completion.data.take() {
-        Some(it) => it,
-        None => return Ok(original_completion),
-    };
+    let Some(data) = original_completion.data.take() else { return Ok(original_completion) };
 
     let resolve_data: lsp_ext::CompletionResolveData = serde_json::from_value(data)?;
 
     let file_id = from_proto::file_id(&snap, &resolve_data.position.text_document.uri)?;
     let line_index = snap.file_line_index(file_id)?;
-    let offset = from_proto::offset(&line_index, resolve_data.position.position)?;
+    // FIXME: We should fix up the position when retrying the cancelled request instead
+    let Ok(offset) = from_proto::offset(&line_index, resolve_data.position.position) else {
+        return Ok(original_completion);
+    };
     let source_root = snap.analysis.source_root(file_id)?;
 
     let additional_edits = snap
@@ -1240,8 +1246,11 @@ pub(crate) fn handle_code_action(
         frange,
     )?;
     for (index, assist) in assists.into_iter().enumerate() {
-        let resolve_data =
-            if code_action_resolve_cap { Some((index, params.clone())) } else { None };
+        let resolve_data = if code_action_resolve_cap {
+            Some((index, params.clone(), snap.file_version(file_id)))
+        } else {
+            None
+        };
         let code_action = to_proto::code_action(&snap, assist, resolve_data)?;
 
         // Check if the client supports the necessary `ResourceOperation`s.
@@ -1280,12 +1289,14 @@ pub(crate) fn handle_code_action_resolve(
     mut code_action: lsp_ext::CodeAction,
 ) -> anyhow::Result<lsp_ext::CodeAction> {
     let _p = tracing::span!(tracing::Level::INFO, "handle_code_action_resolve").entered();
-    let params = match code_action.data.take() {
-        Some(it) => it,
-        None => return Err(invalid_params_error("code action without data".to_owned()).into()),
+    let Some(params) = code_action.data.take() else {
+        return Err(invalid_params_error("code action without data".to_owned()).into());
     };
 
     let file_id = from_proto::file_id(&snap, &params.code_action_params.text_document.uri)?;
+    if snap.file_version(file_id) != params.version {
+        return Err(invalid_params_error("stale code action".to_owned()).into());
+    }
     let line_index = snap.file_line_index(file_id)?;
     let range = from_proto::text_range(&line_index, params.code_action_params.range)?;
     let frange = FileRange { file_id, range };
@@ -1411,12 +1422,11 @@ pub(crate) fn handle_code_lens(
 
 pub(crate) fn handle_code_lens_resolve(
     snap: GlobalStateSnapshot,
-    code_lens: CodeLens,
+    mut code_lens: CodeLens,
 ) -> anyhow::Result<CodeLens> {
-    if code_lens.data.is_none() {
-        return Ok(code_lens);
-    }
-    let Some(annotation) = from_proto::annotation(&snap, code_lens.clone())? else {
+    let Some(data) = code_lens.data.take() else { return Ok(code_lens) };
+    let resolve = serde_json::from_value::<lsp_ext::CodeLensResolveData>(data)?;
+    let Some(annotation) = from_proto::annotation(&snap, code_lens.range, resolve)? else {
         return Ok(code_lens);
     };
     let annotation = snap.analysis.resolve_annotation(annotation)?;
@@ -1495,6 +1505,10 @@ pub(crate) fn handle_inlay_hints(
     )?;
     let line_index = snap.file_line_index(file_id)?;
     let source_root = snap.analysis.source_root(file_id)?;
+    let range = TextRange::new(
+        range.start().min(line_index.index.len()),
+        range.end().min(line_index.index.len()),
+    );
 
     let inlay_hints_config = snap.config.inlay_hints(Some(source_root));
     Ok(Some(
@@ -1522,8 +1536,12 @@ pub(crate) fn handle_inlay_hints_resolve(
 
     let Some(data) = original_hint.data.take() else { return Ok(original_hint) };
     let resolve_data: lsp_ext::InlayHintResolveData = serde_json::from_value(data)?;
-    let Some(hash) = resolve_data.hash.parse().ok() else { return Ok(original_hint) };
     let file_id = FileId::from_raw(resolve_data.file_id);
+    if resolve_data.version != snap.file_version(file_id) {
+        tracing::warn!("Inlay hint resolve data is outdated");
+        return Ok(original_hint);
+    }
+    let Some(hash) = resolve_data.hash.parse().ok() else { return Ok(original_hint) };
     anyhow::ensure!(snap.file_exists(file_id), "Invalid LSP resolve data");
 
     let line_index = snap.file_line_index(file_id)?;

src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/ext.rs

@@ -561,6 +561,7 @@ pub struct CodeAction {
 pub struct CodeActionData {
     pub code_action_params: lsp_types::CodeActionParams,
     pub id: String,
+    pub version: Option<i32>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Default, Deserialize, Serialize)]
@@ -802,13 +803,15 @@ impl Request for OnTypeFormatting {
 pub struct CompletionResolveData {
     pub position: lsp_types::TextDocumentPositionParams,
     pub imports: Vec<CompletionImport>,
+    pub version: Option<i32>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct InlayHintResolveData {
     pub file_id: u32,
     // This is a string instead of a u64 as javascript can't represent u64 fully
     pub hash: String,
+    pub version: Option<i32>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/from_proto.rs

@@ -9,10 +9,8 @@ use syntax::{TextRange, TextSize};
 use vfs::AbsPathBuf;
 
 use crate::{
-    from_json,
     global_state::GlobalStateSnapshot,
     line_index::{LineIndex, PositionEncoding},
-    lsp::utils::invalid_params_error,
     lsp_ext,
 };
 
@@ -105,16 +103,13 @@ pub(crate) fn assist_kind(kind: lsp_types::CodeActionKind) -> Option<AssistKind>
 
 pub(crate) fn annotation(
     snap: &GlobalStateSnapshot,
-    code_lens: lsp_types::CodeLens,
+    range: lsp_types::Range,
+    data: lsp_ext::CodeLensResolveData,
 ) -> anyhow::Result<Option<Annotation>> {
-    let data =
-        code_lens.data.ok_or_else(|| invalid_params_error("code lens without data".to_owned()))?;
-    let resolve = from_json::<lsp_ext::CodeLensResolveData>("CodeLensResolveData", &data)?;
-
-    match resolve.kind {
+    match data.kind {
         lsp_ext::CodeLensResolveDataKind::Impls(params) => {
             if snap.url_file_version(&params.text_document_position_params.text_document.uri)
-                != Some(resolve.version)
+                != Some(data.version)
             {
                 return Ok(None);
             }
@@ -123,19 +118,19 @@ pub(crate) fn annotation(
             let line_index = snap.file_line_index(file_id)?;
 
             Ok(Annotation {
-                range: text_range(&line_index, code_lens.range)?,
+                range: text_range(&line_index, range)?,
                 kind: AnnotationKind::HasImpls { pos, data: None },
             })
         }
         lsp_ext::CodeLensResolveDataKind::References(params) => {
-            if snap.url_file_version(&params.text_document.uri) != Some(resolve.version) {
+            if snap.url_file_version(&params.text_document.uri) != Some(data.version) {
                 return Ok(None);
             }
             let pos @ FilePosition { file_id, .. } = file_position(snap, params)?;
             let line_index = snap.file_line_index(file_id)?;
 
             Ok(Annotation {
-                range: text_range(&line_index, code_lens.range)?,
+                range: text_range(&line_index, range)?,
                 kind: AnnotationKind::HasReferences { pos, data: None },
             })
         }

src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs

@@ -225,13 +225,14 @@ pub(crate) fn snippet_text_edit_vec(
 pub(crate) fn completion_items(
     config: &Config,
     line_index: &LineIndex,
+    version: Option<i32>,
     tdpp: lsp_types::TextDocumentPositionParams,
     items: Vec<CompletionItem>,
 ) -> Vec<lsp_types::CompletionItem> {
     let max_relevance = items.iter().map(|it| it.relevance.score()).max().unwrap_or_default();
     let mut res = Vec::with_capacity(items.len());
     for item in items {
-        completion_item(&mut res, config, line_index, &tdpp, max_relevance, item);
+        completion_item(&mut res, config, line_index, version, &tdpp, max_relevance, item);
     }
 
     if let Some(limit) = config.completion(None).limit {
@@ -246,6 +247,7 @@ fn completion_item(
     acc: &mut Vec<lsp_types::CompletionItem>,
     config: &Config,
     line_index: &LineIndex,
+    version: Option<i32>,
     tdpp: &lsp_types::TextDocumentPositionParams,
     max_relevance: u32,
     item: CompletionItem,
@@ -328,7 +330,7 @@ fn completion_item(
             })
             .collect::<Vec<_>>();
         if !imports.is_empty() {
-            let data = lsp_ext::CompletionResolveData { position: tdpp.clone(), imports };
+            let data = lsp_ext::CompletionResolveData { position: tdpp.clone(), imports, version };
             lsp_item.data = Some(to_value(data).unwrap());
         }
     }
@@ -483,6 +485,7 @@ pub(crate) fn inlay_hint(
             to_value(lsp_ext::InlayHintResolveData {
                 file_id: file_id.index(),
                 hash: hash.to_string(),
+                version: snap.file_version(file_id),
             })
             .unwrap(),
         ),
@@ -1318,7 +1321,7 @@ pub(crate) fn code_action_kind(kind: AssistKind) -> lsp_types::CodeActionKind {
 pub(crate) fn code_action(
     snap: &GlobalStateSnapshot,
     assist: Assist,
-    resolve_data: Option<(usize, lsp_types::CodeActionParams)>,
+    resolve_data: Option<(usize, lsp_types::CodeActionParams, Option<i32>)>,
 ) -> Cancellable<lsp_ext::CodeAction> {
     let mut res = lsp_ext::CodeAction {
         title: assist.label.to_string(),
@@ -1336,10 +1339,11 @@ pub(crate) fn code_action(
 
     match (assist.source_change, resolve_data) {
         (Some(it), _) => res.edit = Some(snippet_workspace_edit(snap, it)?),
-        (None, Some((index, code_action_params))) => {
+        (None, Some((index, code_action_params, version))) => {
             res.data = Some(lsp_ext::CodeActionData {
                 id: format!("{}:{}:{index}", assist.id.0, assist.id.1.name()),
                 code_action_params,
+                version,
             });
         }
         (None, None) => {