diff --git a/Cargo.lock b/Cargo.lock index 6498bee..eb5d51b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -504,9 +504,11 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" name = "java-lsp-proxy" version = "6.8.23" dependencies = [ + "hex", "proxy-common", "serde", "serde_json", + "sha1", ] [[package]] @@ -640,6 +642,7 @@ name = "proxy-common" version = "6.8.23" dependencies = [ "libc", + "percent-encoding", "serde", "serde_json", "tokio", diff --git a/proxy-common/Cargo.toml b/proxy-common/Cargo.toml index 08d80ce..4de15ac 100644 --- a/proxy-common/Cargo.toml +++ b/proxy-common/Cargo.toml @@ -14,6 +14,7 @@ path = "src/lib.rs" tokio = ["dep:tokio"] [dependencies] +percent-encoding = "2.3" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", default-features = false, features = ["io-util"], optional = true } diff --git a/proxy-common/src/uri.rs b/proxy-common/src/uri.rs index ffa0864..fead948 100644 --- a/proxy-common/src/uri.rs +++ b/proxy-common/src/uri.rs @@ -1,21 +1,84 @@ +use percent_encoding::{utf8_percent_encode, AsciiSet, NON_ALPHANUMERIC}; use std::path::Path; -/// Convert a filesystem path to a `file://` URI, matching how language servers' -/// `publishDiagnostics` and the editor key documents. -/// -/// On Unix the path already starts with `/`, so `file://` + path gives the -/// correct `file:///…` form with no extra work. -/// -/// On Windows the backslashes are replaced with `/` and an extra `/` is -/// prepended before the drive letter, so we get `file:///C:/…` rather than -/// `file://C:\…`. -#[cfg(unix)] +const PATH_ENCODE_SET: AsciiSet = NON_ALPHANUMERIC + .remove(b'/') + .remove(b':') + .remove(b'-') + .remove(b'.') + .remove(b'_') + .remove(b'~') + .remove(b'@'); + +/// Convert a filesystem path to an RFC 3986 percent-encoded `file://` URI. pub fn path_to_file_uri(path: &Path) -> String { - format!("file://{}", path.display()) + file_uri_from_path_string(&path.to_string_lossy(), cfg!(windows)) } -#[cfg(windows)] -pub fn path_to_file_uri(path: &Path) -> String { - let s = path.display().to_string().replace('\\', "/"); - format!("file:///{s}") +fn file_uri_from_path_string(path: &str, windows: bool) -> String { + let mut normalized = if windows { + path.replace('\\', "/") + } else { + path.to_string() + }; + if windows { + if let Some(unc) = normalized.strip_prefix("//?/UNC/") { + normalized = format!("//{unc}"); + } else if let Some(verbatim) = normalized.strip_prefix("//?/") { + normalized = verbatim.to_string(); + } + if let Some(unc) = normalized.strip_prefix("//") { + return format!("file://{}", utf8_percent_encode(unc, &PATH_ENCODE_SET)); + } + } + let prefix = if normalized.starts_with('/') { + "file://" + } else { + "file:///" + }; + format!( + "{prefix}{}", + utf8_percent_encode(&normalized, &PATH_ENCODE_SET) + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encodes_reserved_characters_in_unix_paths() { + assert_eq!( + file_uri_from_path_string("/tmp/Java Sources/#1%?.java", false), + "file:///tmp/Java%20Sources/%231%25%3F.java" + ); + } + + #[test] + fn normalizes_and_encodes_windows_paths() { + assert_eq!( + file_uri_from_path_string(r"C:\Users\Jane Doe\A#1.java", true), + "file:///C:/Users/Jane%20Doe/A%231.java" + ); + } + + #[test] + fn preserves_windows_unc_authority() { + assert_eq!( + file_uri_from_path_string(r"\\server\share\A File.java", true), + "file://server/share/A%20File.java" + ); + } + + #[test] + fn normalizes_windows_verbatim_paths() { + assert_eq!( + file_uri_from_path_string(r"\\?\C:\Users\Jane Doe\A.java", true), + "file:///C:/Users/Jane%20Doe/A.java" + ); + assert_eq!( + file_uri_from_path_string(r"\\?\UNC\server\share\A.java", true), + "file://server/share/A.java" + ); + } } diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index dad6fea..50a81fe 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -12,5 +12,7 @@ path = "src/main.rs" [dependencies] proxy-common.workspace = true +hex = "0.4" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +sha1 = "0.10" diff --git a/proxy/src/completions.rs b/proxy/src/completions.rs index cdd1a79..3dad734 100644 --- a/proxy/src/completions.rs +++ b/proxy/src/completions.rs @@ -1,19 +1,15 @@ use serde_json::Value; -/// Returns true if the message contains a completion response with items. -pub fn is_completion_response(msg: &Value) -> bool { - msg.get("result").is_some_and(|result| { - result.get("items").is_some_and(|v| v.is_array()) || result.is_array() - }) -} - /// Single-pass processing of completion items: /// - Sorts methods/functions by parameter count (prepends count to sortText) /// - Strips unsupported VS Code snippet variables ($TM_SELECTED_TEXT) from snippets pub fn process_completions(msg: &mut Value) { + let default_insert_text_format = msg + .pointer("/result/itemDefaults/insertTextFormat") + .and_then(Value::as_u64); let items = match msg.get_mut("result") { Some(result) if result.is_array() => result.as_array_mut(), - Some(result) => result.get_mut("items").and_then(|v| v.as_array_mut()), + Some(result) => result.get_mut("items").and_then(Value::as_array_mut), None => None, }; @@ -33,20 +29,29 @@ pub fn process_completions(msg: &mut Value) { let existing = item.get("sortText").and_then(|v| v.as_str()).unwrap_or(""); item["sortText"] = Value::String(format!("{count:02}{existing}")); } - // Snippet (15): strip $TM_SELECTED_TEXT - 15 => { - strip_tm_selected_text(item, "textEditText"); - strip_tm_selected_text(item, "insertText"); - } _ => {} } + + let insert_text_format = item + .get("insertTextFormat") + .and_then(Value::as_u64) + .or(default_insert_text_format); + if kind == 15 || insert_text_format == Some(2) { + sanitize_completion_item(item); + } } } -fn strip_tm_selected_text(item: &mut Value, key: &str) { - if let Some(text) = item.get(key).and_then(|v| v.as_str()) { +fn sanitize_completion_item(item: &mut Value) { + strip_tm_selected_text(item, "/textEditText"); + strip_tm_selected_text(item, "/insertText"); + strip_tm_selected_text(item, "/textEdit/newText"); +} + +fn strip_tm_selected_text(item: &mut Value, pointer: &str) { + if let Some(Value::String(text)) = item.pointer_mut(pointer) { if text.contains("$TM_SELECTED_TEXT") { - item[key] = Value::String(text.replace("$TM_SELECTED_TEXT", "")); + *text = text.replace("$TM_SELECTED_TEXT", ""); } } } @@ -56,15 +61,7 @@ pub fn sanitize_resolved_completion(msg: &mut Value) { let Some(result) = msg.get_mut("result") else { return; }; - strip_tm_selected_text(result, "textEditText"); - strip_tm_selected_text(result, "insertText"); - // Also check inside textEdit.newText - if let Some(new_text) = result.pointer("/textEdit/newText").and_then(|v| v.as_str()) { - if new_text.contains("$TM_SELECTED_TEXT") { - result["textEdit"]["newText"] = - Value::String(new_text.replace("$TM_SELECTED_TEXT", "")); - } - } + sanitize_completion_item(result); } fn count_params(detail: &str) -> usize { @@ -87,3 +84,133 @@ fn count_params(detail: &str) -> usize { } count } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn processes_array_completion_results() { + let mut response = json!({ + "result": [{ + "kind": 2, + "labelDetails": { "detail": "(String, List)" }, + "sortText": "method" + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["sortText"], json!("02method")); + } + + #[test] + fn applies_list_item_default_snippet_format() { + let mut response = json!({ + "result": { + "itemDefaults": { "insertTextFormat": 2 }, + "items": [{ + "kind": 2, + "insertText": "$TM_SELECTED_TEXT.trim()" + }] + } + }); + + process_completions(&mut response); + + assert_eq!( + response["result"]["items"][0]["insertText"], + json!(".trim()") + ); + } + + #[test] + fn sanitizes_text_edit_text() { + let mut response = json!({ + "result": [{ + "kind": 15, + "textEditText": "$TM_SELECTED_TEXT.field" + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["textEditText"], json!(".field")); + } + + #[test] + fn sanitizes_insert_replace_edit_text() { + let mut response = json!({ + "result": [{ + "kind": 15, + "textEdit": { + "newText": "$TM_SELECTED_TEXT.var", + "insert": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 0 } + }, + "replace": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 3 } + } + } + }] + }); + + process_completions(&mut response); + + assert_eq!(response["result"][0]["textEdit"]["newText"], json!(".var")); + } + + #[test] + fn leaves_plain_text_completion_unchanged() { + let mut response = json!({ + "result": [{ + "kind": 1, + "insertTextFormat": 1, + "insertText": "$TM_SELECTED_TEXT" + }] + }); + + process_completions(&mut response); + + assert_eq!( + response["result"][0]["insertText"], + json!("$TM_SELECTED_TEXT") + ); + } + + #[test] + fn postfix_var_completion_keeps_jdtls_spacing() { + let insertion = "var name = \"hello world\";"; + let mut response = json!({ + "result": [{ + "label": ".var", + "kind": 15, + "insertTextFormat": 2, + "textEdit": { + "newText": insertion, + "range": { + "start": { "line": 0, "character": 0 }, + "end": { "line": 0, "character": 17 } + } + } + }] + }); + + process_completions(&mut response); + + assert_eq!( + response["result"][0]["textEdit"]["newText"], + json!(insertion) + ); + } + + #[test] + fn counts_nested_generic_parameters() { + assert_eq!(count_params("(Map>, int)"), 2); + assert_eq!(count_params("()"), 0); + assert_eq!(count_params("not-a-signature"), 0); + } +} diff --git a/proxy/src/decompile.rs b/proxy/src/decompile.rs index f8e079c..2529c7e 100644 --- a/proxy/src/decompile.rs +++ b/proxy/src/decompile.rs @@ -1,203 +1,1457 @@ +use crate::{lsp_error, pending::PendingResponses}; +use proxy_common::{encode_lsp, path_to_file_uri}; use serde_json::{json, Value}; +use sha1::{Digest, Sha1}; use std::{ - collections::{hash_map::DefaultHasher, HashMap}, - env, fs, - hash::{Hash, Hasher}, + collections::{HashMap, HashSet, VecDeque}, + env, + fs::{self, OpenOptions}, io::Write, - path::PathBuf, - sync::{mpsc, Arc, Mutex}, + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Condvar, Mutex, + }, + thread::{self, JoinHandle}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; -use crate::{lsp_error, lsp_warn}; -use proxy_common::{encode_lsp, path_to_file_uri}; +const DECOMPILED_DIR: &str = "jdtls_decompiled"; +const FETCH_WORKERS: usize = 2; +const JOB_WORKERS: usize = 2; +const MAX_QUEUED_JOBS: usize = 64; +const FETCH_TIMEOUT: Duration = Duration::from_secs(5); +const NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(2); + +pub type SharedWriter = Arc>>; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RewriteMode { + Locations, + Strings, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Priority { + Interactive, + Bulk, +} + +pub struct RewriteJob { + pub token: u64, + pub message: Value, + pub mode: RewriteMode, + pub priority: Priority, + pub deadline: Instant, + pub complete: Box, +} + +#[derive(Clone)] +pub struct DecompileCoordinator { + inner: Arc, +} + +struct Inner { + fetcher: Arc, + owned_id_prefix: String, + cache_dir: PathBuf, + request_counter: AtomicU64, + max_bulk_jobs: usize, + state: Mutex, + work_available: Condvar, + state_changed: Condvar, + workers: Mutex>>, +} + +impl Drop for Inner { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.cache_dir); + } +} + +#[derive(Default)] +struct State { + uris: HashMap, + interactive_uris: VecDeque, + bulk_uris: VecDeque, + interactive_jobs: VecDeque, + bulk_jobs: VecDeque, + canceled_jobs: HashSet, + bulk_jobs_active: usize, + bulk_fetches: usize, + shutdown: bool, +} + +struct UriEntry { + status: UriStatus, + waiters: HashSet, +} + +enum UriStatus { + Queued(Priority), + InFlight { request_id: Value }, + Ready(String), + Failed(Instant), +} + +trait ClassContentFetcher: Send + Sync { + fn fetch(&self, uri: &str, request_id: Value) -> Option; + fn cancel(&self, request_id: &Value); +} + +struct JdtlsFetcher { + writer: SharedWriter, + pending: Arc, +} + +impl ClassContentFetcher for JdtlsFetcher { + fn fetch(&self, uri: &str, request_id: Value) -> Option { + let receiver = self.pending.register(request_id.clone()); + let request = encode_lsp(&json!({ + "jsonrpc": "2.0", + "id": request_id, + "method": "java/classFileContents", + "params": { "uri": uri } + })); + + let write_succeeded = { + let mut writer = self.writer.lock().unwrap(); + writer.write_all(request.as_bytes()).is_ok() && writer.flush().is_ok() + }; + if !write_succeeded { + self.pending.remove(&request_id); + return None; + } + + match receiver.recv_timeout(FETCH_TIMEOUT) { + Ok(response) => response + .get("result") + .and_then(Value::as_str) + .filter(|content| !content.is_empty()) + .map(str::to_string), + Err(_) => { + self.pending.remove(&request_id); + None + } + } + } + + fn cancel(&self, request_id: &Value) { + self.pending.remove(request_id); + let cancel = encode_lsp(&json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": request_id } + })); + let mut writer = self.writer.lock().unwrap(); + let _ = writer.write_all(cancel.as_bytes()); + let _ = writer.flush(); + } +} + +impl DecompileCoordinator { + pub fn new( + writer: SharedWriter, + pending: Arc, + owned_id_prefix: String, + ) -> Self { + Self::with_fetcher( + Arc::new(JdtlsFetcher { writer, pending }), + owned_id_prefix, + FETCH_WORKERS, + JOB_WORKERS, + ) + } + + fn with_fetcher( + fetcher: Arc, + owned_id_prefix: String, + fetch_workers: usize, + job_workers: usize, + ) -> Self { + let cache_dir = session_cache_dir(&owned_id_prefix); + let coordinator = Self { + inner: Arc::new(Inner { + fetcher, + owned_id_prefix, + cache_dir, + request_counter: AtomicU64::new(1), + max_bulk_jobs: job_workers.saturating_sub(1).max(1), + state: Mutex::new(State::default()), + work_available: Condvar::new(), + state_changed: Condvar::new(), + workers: Mutex::new(Vec::new()), + }), + }; + + let mut workers = Vec::with_capacity(fetch_workers + job_workers + 1); + for _ in 0..fetch_workers { + let inner = Arc::clone(&coordinator.inner); + workers.push(thread::spawn(move || fetch_worker(inner))); + } + for _ in 0..job_workers { + let inner = Arc::clone(&coordinator.inner); + workers.push(thread::spawn(move || job_worker(inner))); + } + let inner = Arc::clone(&coordinator.inner); + workers.push(thread::spawn(move || deadline_worker(inner))); + *coordinator.inner.workers.lock().unwrap() = workers; + + coordinator + } + + /// Enqueues without blocking the JDTLS stdout router. On saturation the + /// original job is returned so the caller can forward it unchanged. + pub fn submit(&self, job: RewriteJob) -> Result<(), RewriteJob> { + let mut state = self.inner.state.lock().unwrap(); + if state.shutdown || state.interactive_jobs.len() + state.bulk_jobs.len() >= MAX_QUEUED_JOBS + { + return Err(job); + } + + if job.priority == Priority::Bulk { + state.bulk_jobs.push_back(job); + } else { + state.interactive_jobs.push_back(job); + } + self.inner.work_available.notify_all(); + Ok(()) + } + + pub fn cancel(&self, token: u64) { + let request_ids = { + let mut state = self.inner.state.lock().unwrap(); + if cancel_job_locked(&mut state, token) { + state.canceled_jobs.remove(&token); + } + detach_orphaned_fetches(&mut state) + }; + for request_id in request_ids { + self.inner.fetcher.cancel(&request_id); + } + self.inner.state_changed.notify_all(); + self.inner.work_available.notify_all(); + } + + pub fn consume_cancellation(&self, token: u64) -> bool { + self.inner + .state + .lock() + .unwrap() + .canceled_jobs + .remove(&token) + } + + pub fn is_canceled(&self, token: u64) -> bool { + self.inner + .state + .lock() + .unwrap() + .canceled_jobs + .contains(&token) + } + + pub fn shutdown(&self) { + let queued = { + let mut state = self.inner.state.lock().unwrap(); + if state.shutdown { + Vec::new() + } else { + state.shutdown = true; + let mut queued: Vec<_> = state.interactive_jobs.drain(..).collect(); + queued.extend(state.bulk_jobs.drain(..)); + queued + } + }; + self.inner.work_available.notify_all(); + self.inner.state_changed.notify_all(); + + for job in queued { + if !self.consume_cancellation(job.token) { + (job.complete)(job.message); + } + } + + let workers = std::mem::take(&mut *self.inner.workers.lock().unwrap()); + for worker in workers { + let _ = worker.join(); + } + } + + pub fn cleanup_cache(&self) { + let _ = fs::remove_dir_all(&self.inner.cache_dir); + } +} + +fn cancel_job_locked(state: &mut State, token: u64) -> bool { + state.canceled_jobs.insert(token); + let queued_jobs = state.interactive_jobs.len() + state.bulk_jobs.len(); + state + .interactive_jobs + .retain(|queued| queued.token != token); + state.bulk_jobs.retain(|queued| queued.token != token); + for entry in state.uris.values_mut() { + entry.waiters.remove(&token); + } + queued_jobs != state.interactive_jobs.len() + state.bulk_jobs.len() +} + +fn detach_orphaned_fetches(state: &mut State) -> Vec { + let now = Instant::now(); + let orphaned: Vec<_> = state + .uris + .iter() + .filter(|(_, entry)| entry.waiters.is_empty()) + .filter(|(_, entry)| { + !matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) + }) + .map(|(uri, _)| uri.clone()) + .collect(); + let mut request_ids = Vec::new(); + for uri in orphaned { + if let Some(UriEntry { + status: UriStatus::InFlight { request_id }, + .. + }) = state.uris.remove(&uri) + { + request_ids.push(request_id); + } + } + request_ids +} + +fn job_worker(inner: Arc) { + loop { + let job = { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return; + } + if let Some(job) = state.interactive_jobs.pop_front() { + break job; + } + if state.bulk_jobs_active < inner.max_bulk_jobs { + if let Some(job) = state.bulk_jobs.pop_front() { + state.bulk_jobs_active += 1; + break job; + } + } + state = inner.work_available.wait(state).unwrap(); + } + }; + + if is_canceled(&inner, job.token) { + finish_job(&inner, job.token, job.priority); + continue; + } + + let mut message = job.message; + if Instant::now() < job.deadline { + rewrite_message( + &inner, + job.token, + &mut message, + job.mode, + job.priority, + job.deadline, + ); + } + let canceled = is_canceled(&inner, job.token); + finish_job(&inner, job.token, job.priority); + if !canceled { + (job.complete)(message); + } + } +} + +fn deadline_worker(inner: Arc) { + loop { + let expired = { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return; + } + + let now = Instant::now(); + let mut expired = take_expired_jobs(&mut state.interactive_jobs, now); + expired.extend(take_expired_jobs(&mut state.bulk_jobs, now)); + if !expired.is_empty() { + break expired; + } + + let next_deadline = state + .interactive_jobs + .iter() + .chain(state.bulk_jobs.iter()) + .map(|job| job.deadline) + .min(); + state = if let Some(deadline) = next_deadline { + let timeout = deadline.saturating_duration_since(Instant::now()); + inner.work_available.wait_timeout(state, timeout).unwrap().0 + } else { + inner.work_available.wait(state).unwrap() + }; + } + }; + + for job in expired { + (job.complete)(job.message); + } + } +} + +fn take_expired_jobs(queue: &mut VecDeque, now: Instant) -> Vec { + let mut expired = Vec::new(); + let queued = queue.len(); + for _ in 0..queued { + let job = queue.pop_front().unwrap(); + if job.deadline <= now { + expired.push(job); + } else { + queue.push_back(job); + } + } + expired +} + +fn is_canceled(inner: &Inner, token: u64) -> bool { + inner.state.lock().unwrap().canceled_jobs.contains(&token) +} + +fn finish_job(inner: &Inner, token: u64, priority: Priority) { + let mut state = inner.state.lock().unwrap(); + if priority == Priority::Bulk { + state.bulk_jobs_active = state.bulk_jobs_active.saturating_sub(1); + } + state.canceled_jobs.remove(&token); + for entry in state.uris.values_mut() { + entry.waiters.remove(&token); + } + let now = Instant::now(); + state.uris.retain(|_, entry| { + !entry.waiters.is_empty() + || matches!(entry.status, UriStatus::InFlight { .. }) + || matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) + }); + inner.work_available.notify_all(); +} + +fn rewrite_message( + inner: &Inner, + token: u64, + message: &mut Value, + mode: RewriteMode, + priority: Priority, + deadline: Instant, +) { + let Some(result) = message.get_mut("result") else { + return; + }; + + let mut uris = Vec::new(); + match mode { + RewriteMode::Locations => collect_jdt_location_uris(result, &mut uris, &mut HashSet::new()), + RewriteMode::Strings => collect_jdt_uris(result, &mut uris, &mut HashSet::new()), + } + if uris.is_empty() { + return; + } + + let replacements = resolve_uris(inner, token, &uris, priority, deadline); + match mode { + RewriteMode::Locations => { + replace_jdt_location_uris(result, &replacements); + } + RewriteMode::Strings => replace_in_strings(result, &replacements), + } +} + +fn resolve_uris( + inner: &Inner, + token: u64, + uris: &[String], + priority: Priority, + deadline: Instant, +) -> HashMap { + let mut replacements = HashMap::new(); + let mut unresolved = Vec::new(); + + for uri in uris { + let path = cache_path_in(&inner.cache_dir, uri); + if path.is_file() { + replacements.insert(uri.clone(), path_to_file_uri(&path)); + } else { + unresolved.push(uri.clone()); + } + } + if unresolved.is_empty() { + return replacements; + } + + let mut state = inner.state.lock().unwrap(); + let now = Instant::now(); + state.uris.retain(|_, entry| { + !entry.waiters.is_empty() + || matches!(entry.status, UriStatus::InFlight { .. }) + || matches!(entry.status, UriStatus::Failed(retry_after) if retry_after > now) + }); + if state.canceled_jobs.contains(&token) { + return replacements; + } + + for uri in &unresolved { + let mut enqueue = None; + match state.uris.get_mut(uri) { + Some(entry) => { + entry.waiters.insert(token); + match entry.status { + UriStatus::Ready(ref file_uri) => { + replacements.insert(uri.clone(), file_uri.clone()); + } + UriStatus::Failed(retry_after) if retry_after <= Instant::now() => { + entry.status = UriStatus::Queued(priority); + enqueue = Some(priority); + } + UriStatus::Queued(Priority::Bulk) if priority == Priority::Interactive => { + entry.status = UriStatus::Queued(Priority::Interactive); + enqueue = Some(Priority::Interactive); + } + _ => {} + } + } + None => { + state.uris.insert( + uri.clone(), + UriEntry { + status: UriStatus::Queued(priority), + waiters: HashSet::from([token]), + }, + ); + enqueue = Some(priority); + } + } + if let Some(priority) = enqueue { + queue_uri(&mut state, uri.clone(), priority); + } + } + inner.work_available.notify_all(); + + loop { + let mut waiting = false; + for uri in &unresolved { + match state.uris.get(uri).map(|entry| &entry.status) { + Some(UriStatus::Ready(file_uri)) => { + replacements.insert(uri.clone(), file_uri.clone()); + } + Some(UriStatus::Queued(_) | UriStatus::InFlight { .. }) => waiting = true, + Some(UriStatus::Failed(_)) | None => {} + } + } + + if !waiting + || state.shutdown + || state.canceled_jobs.contains(&token) + || Instant::now() >= deadline + { + break; + } + + let remaining = deadline.saturating_duration_since(Instant::now()); + let (next_state, _) = inner.state_changed.wait_timeout(state, remaining).unwrap(); + state = next_state; + } + + for uri in &unresolved { + if let Some(entry) = state.uris.get_mut(uri) { + entry.waiters.remove(&token); + } + } + let orphaned = detach_orphaned_fetches(&mut state); + drop(state); + for request_id in orphaned { + inner.fetcher.cancel(&request_id); + } + + replacements +} + +fn queue_uri(state: &mut State, uri: String, priority: Priority) { + match priority { + Priority::Interactive => state.interactive_uris.push_back(uri), + Priority::Bulk => state.bulk_uris.push_back(uri), + } +} + +fn fetch_worker(inner: Arc) { + loop { + let Some((uri, priority, request_id)) = take_uri_work(&inner) else { + return; + }; + + let content = inner.fetcher.fetch(&uri, request_id.clone()); + let resolved = content + .and_then(|content| write_cached_source(&inner.cache_dir, &uri, content.as_bytes())); + + let mut state = inner.state.lock().unwrap(); + if priority == Priority::Bulk { + state.bulk_fetches = state.bulk_fetches.saturating_sub(1); + } + if let Some(entry) = state.uris.get_mut(&uri) { + let is_current = matches!( + &entry.status, + UriStatus::InFlight { + request_id: current, + .. + } if current == &request_id + ); + if is_current { + entry.status = match resolved { + Some(file_uri) => UriStatus::Ready(file_uri), + None => UriStatus::Failed(Instant::now() + NEGATIVE_CACHE_TTL), + }; + } + } + let remove_ready = state.uris.get(&uri).is_some_and(|entry| { + entry.waiters.is_empty() && matches!(entry.status, UriStatus::Ready(_)) + }); + if remove_ready { + state.uris.remove(&uri); + } + inner.state_changed.notify_all(); + inner.work_available.notify_all(); + } +} -const DECOMPILED_DIR: &str = "jdtls-decompiled"; +fn take_uri_work(inner: &Inner) -> Option<(String, Priority, Value)> { + let mut state = inner.state.lock().unwrap(); + loop { + if state.shutdown { + return None; + } + + let candidate = pop_valid_uri(&mut state, Priority::Interactive).or_else(|| { + if state.bulk_fetches == 0 { + pop_valid_uri(&mut state, Priority::Bulk) + } else { + None + } + }); -fn cache_dir() -> PathBuf { - env::temp_dir().join(DECOMPILED_DIR) + if let Some((uri, priority)) = candidate { + let sequence = inner.request_counter.fetch_add(1, Ordering::Relaxed); + let request_id = + Value::String(format!("{}decompile-{sequence}", inner.owned_id_prefix)); + if priority == Priority::Bulk { + state.bulk_fetches += 1; + } + if let Some(entry) = state.uris.get_mut(&uri) { + entry.status = UriStatus::InFlight { + request_id: request_id.clone(), + }; + } + return Some((uri, priority, request_id)); + } + + state = inner.work_available.wait(state).unwrap(); + } } -fn cache_path(uri: &str) -> PathBuf { - let mut hasher = DefaultHasher::new(); - uri.hash(&mut hasher); - let hex = format!("{:016x}", hasher.finish()); +fn pop_valid_uri(state: &mut State, priority: Priority) -> Option<(String, Priority)> { + let queue = match priority { + Priority::Interactive => &mut state.interactive_uris, + Priority::Bulk => &mut state.bulk_uris, + }; + while let Some(uri) = queue.pop_front() { + let valid = state.uris.get(&uri).is_some_and(|entry| { + !entry.waiters.is_empty() + && matches!(entry.status, UriStatus::Queued(current) if current == priority) + }); + if valid { + return Some((uri, priority)); + } + } + None +} - // jdt://contents/java.base/java.util/ArrayList.java?=.../%3Cjava.util%28ArrayList.class - // The class name is between the last %28 (URL-encoded '(') and .class at the end - let name = uri +fn session_cache_dir(scope: &str) -> PathBuf { + static SESSION_COUNTER: AtomicU64 = AtomicU64::new(1); + let mut sequence = SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); + let scope = hex::encode(Sha1::digest(scope.as_bytes())); + let started_at = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let root = env::temp_dir().join(DECOMPILED_DIR); + let _ = fs::create_dir_all(&root); + loop { + let directory = root.join(format!( + "session_{scope}_{}_{started_at}_{sequence}", + std::process::id() + )); + match fs::create_dir(&directory) { + Ok(()) => return directory, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + sequence = SESSION_COUNTER.fetch_add(1, Ordering::Relaxed); + } + Err(_) => return directory, + } + } +} + +fn cache_path_in(directory: &Path, uri: &str) -> PathBuf { + let digest = hex::encode(Sha1::digest(uri.as_bytes())); + let raw_name = uri .rsplit_once("%28") .and_then(|(_, rest)| rest.strip_suffix(".class")) .or_else(|| { uri.split('?') .next() .and_then(|path| path.rsplit('/').next()) - .and_then(|seg| seg.strip_suffix(".java").or(seg.strip_suffix(".class"))) + .and_then(|segment| { + segment + .strip_suffix(".java") + .or(segment.strip_suffix(".class")) + }) }) .unwrap_or("Decompiled"); - - cache_dir().join(format!("{name}-{hex}.java")) -} - -/// Send `java/classFileContents` to JDTLS and wait for the response. -fn fetch_class_contents( - uri: &str, - writer: &Arc>, - pending: &Arc>>>, - request_id: Value, -) -> Option { - let (tx, rx) = mpsc::channel(); - pending.lock().unwrap().insert(request_id.clone(), tx); - - let req = encode_lsp(&json!({ - "jsonrpc": "2.0", - "id": request_id, - "method": "java/classFileContents", - "params": { "uri": uri } - })); - { - let mut w = writer.lock().unwrap(); - let _ = w.write_all(req.as_bytes()); - let _ = w.flush(); - } - - match rx.recv_timeout(std::time::Duration::from_secs(10)) { - Ok(resp) => { - let content = resp.get("result")?.as_str()?; - Some(content.to_string()) - } - Err(_) => { - lsp_warn!("[decompile] Timed out fetching class contents for {uri}"); - None - } + let mut name: String = raw_name + .chars() + .enumerate() + .map(|(index, character)| { + let valid = if index == 0 { + character.is_ascii_alphabetic() || matches!(character, '_' | '$') + } else { + character.is_ascii_alphanumeric() || matches!(character, '_' | '$') + }; + if valid { + character + } else { + '_' + } + }) + .collect(); + if name.is_empty() { + name.push_str("Decompiled"); } + + directory + .join(format!("uri_{digest}")) + .join(format!("{name}.java")) } -fn resolve_jdt_uri( - uri: &str, - writer: &Arc>, - pending: &Arc>>>, - request_id: Value, -) -> Option { - let path = cache_path(uri); - if path.exists() { - return Some(path_to_file_uri(&path)); +fn write_cached_source(directory: &Path, uri: &str, content: &[u8]) -> Option { + if content.is_empty() { + return None; + } + + if let Err(error) = fs::create_dir_all(directory) { + lsp_error!( + "[decompile] Failed to create {}: {error}", + directory.display() + ); + return None; + } + let target = cache_path_in(directory, uri); + if target.is_file() { + return Some(path_to_file_uri(&target)); + } + if let Some(parent) = target.parent() { + if let Err(error) = fs::create_dir_all(parent) { + lsp_error!("[decompile] Failed to create {}: {error}", parent.display()); + return None; + } } - let content = fetch_class_contents(uri, writer, pending, request_id)?; - let _ = fs::create_dir_all(cache_dir()); - match fs::write(&path, &content) { - Ok(_) => Some(path_to_file_uri(&path)), - Err(e) => { - lsp_error!("[decompile] Failed to write {}: {e}", path.display()); + static TEMP_COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed); + let temporary = target.with_extension(format!("tmp-{}-{sequence}", std::process::id())); + let result = (|| { + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary)?; + file.write_all(content)?; + file.flush()?; + fs::rename(&temporary, &target) + })(); + + match result { + Ok(()) => Some(path_to_file_uri(&target)), + Err(_error) if target.is_file() => { + let _ = fs::remove_file(&temporary); + Some(path_to_file_uri(&target)) + } + Err(error) => { + let _ = fs::remove_file(&temporary); + lsp_error!( + "[decompile] Failed to atomically write {}: {error}", + target.display() + ); None } } } -/// Rewrite any `jdt://` URIs in a definition/typeDefinition/implementation response. -/// Returns `true` if any URI was rewritten. -pub fn rewrite_jdt_locations( - msg: &mut Value, - writer: &Arc>, - pending: &Arc>>>, - next_id: &mut impl FnMut() -> Value, -) -> bool { - let results = match msg.get_mut("result") { - Some(Value::Array(arr)) => arr.iter_mut().collect::>(), - Some(obj @ Value::Object(_)) => vec![obj], - _ => return false, - }; +fn collect_jdt_location_uris(value: &Value, uris: &mut Vec, seen: &mut HashSet) { + match value { + Value::Array(values) => { + for value in values { + collect_jdt_location_uris(value, uris, seen); + } + } + Value::Object(object) => { + for (key, value) in object { + if matches!(key.as_str(), "uri" | "targetUri") { + if let Value::String(uri) = value { + if uri.starts_with("jdt://") && seen.insert(uri.clone()) { + uris.push(uri.clone()); + } + } + continue; + } + collect_jdt_location_uris(value, uris, seen); + } + } + _ => {} + } +} - let mut rewritten = false; - for loc in results { - for key in &["uri", "targetUri"] { - if let Some(Value::String(uri)) = loc.get(key) { - if uri.starts_with("jdt://") { - let jdt_uri = uri.clone(); - if let Some(file_uri) = resolve_jdt_uri(&jdt_uri, writer, pending, next_id()) { - loc[*key] = Value::String(file_uri); - rewritten = true; +fn replace_jdt_location_uris(value: &mut Value, replacements: &HashMap) -> bool { + match value { + Value::Array(values) => { + let mut rewritten = false; + for value in values { + rewritten |= replace_jdt_location_uris(value, replacements); + } + rewritten + } + Value::Object(object) => { + let mut rewritten = false; + for (key, value) in object { + if matches!(key.as_str(), "uri" | "targetUri") { + if let Value::String(uri) = value { + if let Some(file_uri) = replacements.get(uri) { + *uri = file_uri.clone(); + rewritten = true; + } } + continue; } + rewritten |= replace_jdt_location_uris(value, replacements); } + rewritten + } + _ => false, + } +} + +fn jdt_uri_end(value: &str) -> usize { + if let Some(class_end) = value.find(".class") { + let class_end = class_end + ".class".len(); + if value[class_end..].starts_with('?') { + return value[class_end..] + .find(is_jdt_uri_terminator) + .map(|query_end| class_end + query_end) + .unwrap_or(value.len()); } + return class_end; } - rewritten + value.find(is_jdt_uri_terminator).unwrap_or(value.len()) } -/// A jdt:// URI in embedded markdown/text terminates at whitespace or any of these -/// delimiters commonly used in markdown links and JSON strings. The URI itself only -/// contains URL-encoded forms of these characters, so scanning until we hit one of -/// them is safe. -fn jdt_uri_end(s: &str) -> usize { - s.find(|c: char| c.is_whitespace() || matches!(c, ')' | ']' | '"' | '>' | '`' | '\'')) - .unwrap_or(s.len()) +fn is_jdt_uri_terminator(character: char) -> bool { + character.is_whitespace() + || matches!( + character, + ')' | ']' | '}' | '"' | '>' | '`' | '\'' | ',' | ';' + ) } -/// Extract all unique `jdt://` URIs appearing inside any string in `value`. -fn collect_jdt_uris(value: &Value, out: &mut Vec) { +fn collect_jdt_uris(value: &Value, uris: &mut Vec, seen: &mut HashSet) { match value { - Value::String(s) => { - let mut rest = s.as_str(); - while let Some(pos) = rest.find("jdt://") { - let tail = &rest[pos..]; + Value::String(string) => { + let mut rest = string.as_str(); + while let Some(position) = rest.find("jdt://") { + let tail = &rest[position..]; let end = jdt_uri_end(tail); let uri = tail[..end].to_string(); - if !out.contains(&uri) { - out.push(uri); + if seen.insert(uri.clone()) { + uris.push(uri); } rest = &tail[end..]; } } - Value::Array(arr) => arr.iter().for_each(|v| collect_jdt_uris(v, out)), - Value::Object(obj) => obj.values().for_each(|v| collect_jdt_uris(v, out)), + Value::Array(values) => { + for value in values { + collect_jdt_uris(value, uris, seen); + } + } + Value::Object(object) => { + for value in object.values() { + collect_jdt_uris(value, uris, seen); + } + } _ => {} } } -/// Replace all occurrences of any key in `map` with its value, inside every string -/// contained in `value` (recursively). -fn replace_in_strings(value: &mut Value, map: &HashMap) { +fn replace_in_strings(value: &mut Value, replacements: &HashMap) { match value { - Value::String(s) => { - for (from, to) in map { - if s.contains(from.as_str()) { - *s = s.replace(from.as_str(), to); + Value::String(string) => { + for (from, to) in replacements { + if string.contains(from) { + *string = string.replace(from, to); } } } - Value::Array(arr) => arr.iter_mut().for_each(|v| replace_in_strings(v, map)), - Value::Object(obj) => obj.values_mut().for_each(|v| replace_in_strings(v, map)), + Value::Array(values) => { + for value in values { + replace_in_strings(value, replacements); + } + } + Value::Object(object) => { + for value in object.values_mut() { + replace_in_strings(value, replacements); + } + } _ => {} } } -/// Scan a documentation response (hover, signatureHelp, completionItem/resolve, …) -/// for embedded `jdt://` URIs, resolve each one to a `file://` URI backed by a temp -/// file, and replace the URIs in-place in every string of `msg.result`. -pub fn rewrite_jdt_in_strings( - msg: &mut Value, - writer: &Arc>, - pending: &Arc>>>, - next_id: &mut impl FnMut() -> Value, -) { - let Some(result) = msg.get_mut("result") else { - return; +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, }; - let mut uris = Vec::new(); - collect_jdt_uris(result, &mut uris); - if uris.is_empty() { - return; + struct FakeFetcher { + calls: AtomicUsize, + cancellations: AtomicUsize, + active: AtomicUsize, + max_active: AtomicUsize, + delay: Duration, + succeeds: bool, } - let mut map = HashMap::new(); - for uri in uris { - if let Some(file_uri) = resolve_jdt_uri(&uri, writer, pending, next_id()) { - map.insert(uri, file_uri); + impl FakeFetcher { + fn new(delay: Duration) -> Self { + Self { + calls: AtomicUsize::new(0), + cancellations: AtomicUsize::new(0), + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + delay, + succeeds: true, + } + } + + fn failing(delay: Duration) -> Self { + Self { + succeeds: false, + ..Self::new(delay) + } + } + } + + impl ClassContentFetcher for FakeFetcher { + fn fetch(&self, uri: &str, _request_id: Value) -> Option { + self.calls.fetch_add(1, Ordering::Relaxed); + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active.fetch_max(active, Ordering::SeqCst); + thread::sleep(self.delay); + self.active.fetch_sub(1, Ordering::SeqCst); + self.succeeds.then(|| format!("class {} {{}}", uri.len())) + } + + fn cancel(&self, _request_id: &Value) { + self.cancellations.fetch_add(1, Ordering::Relaxed); } } - if !map.is_empty() { - replace_in_strings(result, &map); + + fn unique_uri(name: &str) -> String { + format!( + "jdt://contents/test/{name}-{}.class", + AtomicU64::new(1).fetch_add(1, Ordering::Relaxed) + std::process::id() as u64 + ) + } + + fn rewrite_for_test( + coordinator: &DecompileCoordinator, + token: u64, + uris: &[String], + priority: Priority, + ) -> HashMap { + resolve_uris( + &coordinator.inner, + token, + uris, + priority, + Instant::now() + Duration::from_secs(2), + ) + } + + fn cached_path(coordinator: &DecompileCoordinator, uri: &str) -> PathBuf { + cache_path_in(&coordinator.inner.cache_dir, uri) + } + + #[test] + fn rewrites_nested_workspace_symbol_locations() { + let mut result = json!([ + { + "location": { + "uri": "jdt://contents/java.base/java.lang/String.class" + } + }, + { + "location": { + "uri": "file:///workspace/ProjectString.java" + } + } + ]); + let replacements = HashMap::from([( + "jdt://contents/java.base/java.lang/String.class".to_string(), + "file:///tmp/String.java".to_string(), + )]); + + assert!(replace_jdt_location_uris(&mut result, &replacements)); + assert_eq!( + result[0]["location"]["uri"], + json!("file:///tmp/String.java") + ); + assert_eq!( + result[1]["location"]["uri"], + json!("file:///workspace/ProjectString.java") + ); + } + + #[test] + fn deduplicates_location_uris_before_resolution() { + let uri = "jdt://contents/java.base/java.lang/String.class"; + let result = json!([{ "uri": uri }, { "targetUri": uri }]); + let mut uris = Vec::new(); + + collect_jdt_location_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + + #[test] + fn ignores_jdt_uris_outside_location_fields() { + let result = json!({ + "documentation": "jdt://contents/java.base/java.lang/String.class", + "data": { + "sourceUri": "jdt://contents/java.base/java.lang/Object.class" + } + }); + let mut uris = Vec::new(); + + collect_jdt_location_uris(&result, &mut uris, &mut HashSet::new()); + + assert!(uris.is_empty()); + } + + #[test] + fn embedded_uri_extraction_stops_before_sentence_punctuation() { + let uri = "jdt://contents/java.base/java.lang/String.class"; + let result = json!(format!("See {uri}, then java.lang.Object;")); + let mut uris = Vec::new(); + + collect_jdt_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + + #[test] + fn embedded_uri_extraction_preserves_query_components() { + let uri = "jdt://contents/java.base/java.lang/String.class?=java.base/String"; + let result = json!(format!("See {uri}, then continue.")); + let mut uris = Vec::new(); + + collect_jdt_uris(&result, &mut uris, &mut HashSet::new()); + + assert_eq!(uris, vec![uri]); + } + + #[test] + fn concurrent_waiters_share_one_fetch() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(20))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 2, 0); + let uri = unique_uri("single-flight"); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + + let first = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 1, &[uri], Priority::Interactive)) + }; + let second = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 2, &[uri], Priority::Interactive)) + }; + + assert_eq!(first.join().unwrap().len(), 1); + assert_eq!(second.join().unwrap().len(), 1); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 1); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + coordinator.shutdown(); + } + + #[test] + fn interactive_fetches_use_bounded_parallelism() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(30))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 2, 0); + let uris: Vec<_> = (0..4) + .map(|index| unique_uri(&format!("parallel-{index}"))) + .collect(); + for uri in &uris { + let _ = fs::remove_file(cached_path(&coordinator, uri)); + } + + assert_eq!( + rewrite_for_test(&coordinator, 1, &uris, Priority::Interactive).len(), + uris.len() + ); + assert_eq!(fetcher.max_active.load(Ordering::SeqCst), 2); + for uri in &uris { + let _ = fs::remove_file(cached_path(&coordinator, uri)); + } + coordinator.shutdown(); + } + + #[test] + fn deadline_returns_completed_replacements_only() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(50))); + let coordinator = DecompileCoordinator::with_fetcher(fetcher, "test-".to_string(), 1, 0); + let uris = vec![unique_uri("deadline-a"), unique_uri("deadline-b")]; + for uri in &uris { + let _ = fs::remove_file(cached_path(&coordinator, uri)); + } + + let replacements = resolve_uris( + &coordinator.inner, + 1, + &uris, + Priority::Interactive, + Instant::now() + Duration::from_millis(70), + ); + + assert_eq!(replacements.len(), 1); + for uri in &uris { + let _ = fs::remove_file(cached_path(&coordinator, uri)); + } + coordinator.shutdown(); + } + + #[test] + fn queued_jobs_complete_when_their_deadline_expires() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(150))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "deadline-".to_string(), 1, 1); + let first_uri = unique_uri("blocking-job"); + let second_uri = unique_uri("queued-deadline"); + let (first_tx, _first_rx) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 1, + message: json!({ "result": [{ "uri": first_uri }] }), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(move |message| { + let _ = first_tx.send(message); + }), + }) + .is_ok()); + let worker_start_deadline = Instant::now() + Duration::from_secs(1); + while fetcher.active.load(Ordering::SeqCst) == 0 && Instant::now() < worker_start_deadline { + thread::sleep(Duration::from_millis(1)); + } + assert_eq!(fetcher.active.load(Ordering::SeqCst), 1); + + let original = json!({ "result": [{ "uri": second_uri }] }); + let (second_tx, second_rx) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 2, + message: original.clone(), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_millis(20), + complete: Box::new(move |message| { + let _ = second_tx.send(message); + }), + }) + .is_ok()); + + assert_eq!( + second_rx.recv_timeout(Duration::from_millis(500)).unwrap(), + original + ); + coordinator.shutdown(); + } + + #[test] + fn canceling_one_waiter_keeps_shared_fetch_alive() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(50))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); + let uri = unique_uri("shared-cancel"); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + + let first = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 1, &[uri], Priority::Interactive)) + }; + let second = { + let coordinator = coordinator.clone(); + let uri = uri.clone(); + thread::spawn(move || rewrite_for_test(&coordinator, 2, &[uri], Priority::Interactive)) + }; + while { + let state = coordinator.inner.state.lock().unwrap(); + state + .uris + .get(&uri) + .is_none_or(|entry| entry.waiters.len() < 2) + } { + thread::yield_now(); + } + coordinator.cancel(1); + + assert!(first.join().unwrap().is_empty()); + assert_eq!(second.join().unwrap().len(), 1); + assert_eq!(fetcher.cancellations.load(Ordering::Relaxed), 0); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + coordinator.shutdown(); + } + + #[test] + fn timed_out_fetch_does_not_poison_immediate_retry() { + let fetcher = Arc::new(FakeFetcher::new(Duration::from_millis(60))); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "retry-".to_string(), 2, 0); + let uri = unique_uri("retry-after-timeout"); + + let first = resolve_uris( + &coordinator.inner, + 1, + std::slice::from_ref(&uri), + Priority::Interactive, + Instant::now() + Duration::from_millis(10), + ); + let second = resolve_uris( + &coordinator.inner, + 2, + std::slice::from_ref(&uri), + Priority::Interactive, + Instant::now() + Duration::from_millis(250), + ); + + assert!(first.is_empty()); + assert_eq!(second.len(), 1); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 2); + assert_eq!(fetcher.cancellations.load(Ordering::Relaxed), 1); + coordinator.shutdown(); + } + + #[test] + fn negative_cache_prevents_immediate_retry_storms() { + let fetcher = Arc::new(FakeFetcher::failing(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher.clone(), "test-".to_string(), 1, 0); + let uri = unique_uri("negative-cache"); + let _ = fs::remove_file(cached_path(&coordinator, &uri)); + + assert!(rewrite_for_test( + &coordinator, + 1, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .is_empty()); + assert!(rewrite_for_test( + &coordinator, + 2, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .is_empty()); + assert_eq!(fetcher.calls.load(Ordering::Relaxed), 1); + coordinator.shutdown(); + } + + #[test] + fn coordinators_do_not_share_cache_across_server_sessions() { + let uri = unique_uri("session-cache"); + let first_fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let first = DecompileCoordinator::with_fetcher( + first_fetcher.clone(), + "same-server-scope-".to_string(), + 1, + 0, + ); + assert_eq!( + rewrite_for_test(&first, 1, std::slice::from_ref(&uri), Priority::Interactive).len(), + 1 + ); + let first_cache = first.inner.cache_dir.clone(); + first.shutdown(); + first.cleanup_cache(); + assert!(!first_cache.exists()); + + let second_fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let second = DecompileCoordinator::with_fetcher( + second_fetcher.clone(), + "same-server-scope-".to_string(), + 1, + 0, + ); + assert_ne!(first_cache, second.inner.cache_dir); + assert_eq!( + rewrite_for_test( + &second, + 2, + std::slice::from_ref(&uri), + Priority::Interactive + ) + .len(), + 1 + ); + assert_eq!(second_fetcher.calls.load(Ordering::Relaxed), 1); + second.shutdown(); + second.cleanup_cache(); + } + + #[test] + fn canceling_a_queued_bulk_job_retires_all_state() { + let fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher, "queued-cancel-".to_string(), 0, 0); + assert!(coordinator + .submit(RewriteJob { + token: 7, + message: json!({ "result": [] }), + mode: RewriteMode::Locations, + priority: Priority::Bulk, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(|_| {}), + }) + .is_ok()); + + coordinator.cancel(7); + + let state = coordinator.inner.state.lock().unwrap(); + assert!(!state.canceled_jobs.contains(&7)); + assert!(state.bulk_jobs.is_empty()); + drop(state); + coordinator.shutdown(); + } + + #[test] + fn shutdown_completes_queued_jobs_unchanged() { + let fetcher = Arc::new(FakeFetcher::new(Duration::ZERO)); + let coordinator = + DecompileCoordinator::with_fetcher(fetcher, "shutdown-".to_string(), 0, 0); + let original = json!({ "result": [{ "uri": "jdt://contents/test/A.class" }] }); + let (sender, receiver) = mpsc::channel(); + assert!(coordinator + .submit(RewriteJob { + token: 8, + message: original.clone(), + mode: RewriteMode::Locations, + priority: Priority::Interactive, + deadline: Instant::now() + Duration::from_secs(1), + complete: Box::new(move |message| { + let _ = sender.send(message); + }), + }) + .is_ok()); + + coordinator.shutdown(); + + assert_eq!(receiver.recv().unwrap(), original); + } + + #[test] + fn competing_cache_writes_never_expose_partial_content() { + let uri = unique_uri("atomic-cache"); + let directory = session_cache_dir("atomic-cache-test"); + let path = cache_path_in(&directory, &uri); + let _ = fs::remove_file(&path); + let first_content = vec![b'a'; 32 * 1024]; + let second_content = vec![b'b'; 32 * 1024]; + + let first = { + let uri = uri.clone(); + let content = first_content.clone(); + let directory = directory.clone(); + thread::spawn(move || write_cached_source(&directory, &uri, &content)) + }; + let second = { + let uri = uri.clone(); + let content = second_content.clone(); + let directory = directory.clone(); + thread::spawn(move || write_cached_source(&directory, &uri, &content)) + }; + assert!(first.join().unwrap().is_some()); + assert!(second.join().unwrap().is_some()); + + let cached = fs::read(&path).unwrap(); + assert!(cached == first_content || cached == second_content); + let _ = fs::remove_file(path); + } + + #[test] + fn cache_names_use_stable_digest_and_sanitized_name() { + let directory = Path::new("/tmp/cache"); + let first = cache_path_in(directory, "jdt://contents/a/../../Bad Name.class"); + let second = cache_path_in(directory, "jdt://contents/a/../../Bad Name.class"); + + assert_eq!(first, second); + assert_eq!(first.file_name().unwrap(), "Bad_Name.java"); + assert_eq!(first.extension().unwrap(), "java"); + assert!(first + .parent() + .unwrap() + .file_name() + .unwrap() + .to_string_lossy() + .starts_with("uri_")); + } + + #[test] + fn cache_file_stems_are_valid_java_identifiers() { + let directory = Path::new("/tmp/cache"); + let path = cache_path_in(directory, "jdt://contents/a/Products$Product$Request.class"); + let stem = path.file_stem().unwrap().to_string_lossy(); + + assert_eq!(stem, "Products$Product$Request"); + assert!(stem + .chars() + .enumerate() + .all(|(index, character)| if index == 0 { + character.is_ascii_alphabetic() || matches!(character, '_' | '$') + } else { + character.is_ascii_alphanumeric() || matches!(character, '_' | '$') + })); + } + + #[test] + fn empty_sources_are_not_cached() { + let uri = unique_uri("empty"); + let directory = session_cache_dir("empty-cache-test"); + let path = cache_path_in(&directory, &uri); + let _ = fs::remove_file(&path); + + assert_eq!(write_cached_source(&directory, &uri, b""), None); + assert!(!path.exists()); } } diff --git a/proxy/src/http.rs b/proxy/src/http.rs index 33b4d3d..f4446b7 100644 --- a/proxy/src/http.rs +++ b/proxy/src/http.rs @@ -1,15 +1,15 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::{ - collections::HashMap, io::{BufRead, BufReader, Read, Write}, sync::{ atomic::{AtomicU64, Ordering}, - mpsc, Arc, Mutex, + Arc, Mutex, }, time::Duration, }; +use crate::pending::PendingResponses; use proxy_common::encode_lsp; pub const TIMEOUT: Duration = Duration::from_secs(5); @@ -31,9 +31,9 @@ struct LspRequest { pub fn handle_http( mut stream: std::net::TcpStream, writer: Arc>, - pending: Arc>>>, + pending: Arc, counter: Arc, - proxy_id: &str, + owned_id_prefix: &str, ) { let mut reader = BufReader::new(&stream); @@ -78,10 +78,8 @@ pub fn handle_http( }; let seq = counter.fetch_add(1, Ordering::Relaxed); - let id = Value::String(format!("{proxy_id}-{seq}")); - - let (tx, rx) = mpsc::channel(); - pending.lock().unwrap().insert(id.clone(), tx); + let id = Value::String(format!("{owned_id_prefix}http-{seq}")); + let rx = pending.register(id.clone()); let lsp_req = LspRequest { jsonrpc: "2.0", @@ -90,34 +88,45 @@ pub fn handle_http( params: req.params, }; let encoded = encode_lsp(&lsp_req); - { + let request_sent = { let mut w = writer.lock().unwrap(); - let _ = w.write_all(encoded.as_bytes()); - let _ = w.flush(); - } - - let response = match rx.recv_timeout(TIMEOUT) { - Ok(resp) => resp, - Err(_) => { - pending.lock().unwrap().remove(&id); - let cancel = encode_lsp(&serde_json::json!({ - "jsonrpc": "2.0", - "method": "$/cancelRequest", - "params": { "id": id } - })); - let mut w = writer.lock().unwrap(); - let _ = w.write_all(cancel.as_bytes()); - let _ = w.flush(); + w.write_all(encoded.as_bytes()).is_ok() && w.flush().is_ok() + }; - serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "error": { - "code": -32803, - "message": "Request to language server timed out after 5000ms." - } - }) + let response = if request_sent { + match rx.recv_timeout(TIMEOUT) { + Ok(resp) => resp, + Err(_) => { + pending.remove(&id); + let cancel = encode_lsp(&serde_json::json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": id } + })); + let mut w = writer.lock().unwrap(); + let _ = w.write_all(cancel.as_bytes()); + let _ = w.flush(); + + serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32803, + "message": "Request to language server timed out after 5000ms." + } + }) + } } + } else { + pending.remove(&id); + serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32603, + "message": "Failed to write request to language server." + } + }) }; let resp_body = serde_json::to_vec(&response).unwrap(); diff --git a/proxy/src/log.rs b/proxy/src/log.rs index a93ae58..29b676d 100644 --- a/proxy/src/log.rs +++ b/proxy/src/log.rs @@ -1,7 +1,9 @@ use serde::Serialize; -use std::io::{self, Write}; +use std::sync::OnceLock; -use proxy_common::encode_lsp; +use crate::output::Output; + +static OUTPUT: OnceLock = OnceLock::new(); /// LSP `MessageType` constants as defined in the specification. /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#messageType @@ -30,8 +32,11 @@ struct LogMessageParams<'a> { /// Sends a `window/logMessage` LSP notification to stdout so that Zed /// displays the message in its Server Logs panel. /// -/// This locks stdout for the duration of the write to ensure the LSP -/// framing is not interleaved with other output. +/// The shared output broker serializes this with every other LSP frame. +pub fn init(output: Output) { + let _ = OUTPUT.set(output); +} + fn send_log_message(level: u8, message: &str) { let notification = LogMessageNotification { jsonrpc: "2.0", @@ -42,12 +47,9 @@ fn send_log_message(level: u8, message: &str) { }, }; - let encoded = encode_lsp(¬ification); - - let stdout = io::stdout(); - let mut w = stdout.lock(); - let _ = w.write_all(encoded.as_bytes()); - let _ = w.flush(); + if let Some(output) = OUTPUT.get() { + output.send_value(¬ification); + } } #[allow(dead_code)] diff --git a/proxy/src/main.rs b/proxy/src/main.rs index 1022cfe..7d5191e 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -2,16 +2,20 @@ mod completions; mod decompile; mod http; mod log; +mod output; +mod pending; -use completions::{is_completion_response, process_completions, sanitize_resolved_completion}; -use decompile::{rewrite_jdt_in_strings, rewrite_jdt_locations}; +use completions::{process_completions, sanitize_resolved_completion}; +use decompile::{DecompileCoordinator, Priority, RewriteJob, RewriteMode, SharedWriter}; use http::handle_http; +use output::Output; +use pending::PendingResponses; use proxy_common::{ - parse_lsp_content, raw_has_id, spawn_parent_monitor, write_raw, write_to_stdout, LspReader, + contains_subslice, encode_lsp, parse_lsp_content, raw_has_id, spawn_parent_monitor, LspReader, }; -use serde_json::Value; +use serde_json::{json, Value}; use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, env, fs, io::{self, BufReader, Write}, net::TcpListener, @@ -19,18 +23,101 @@ use std::{ process::{self, Command, Stdio}, sync::{ atomic::{AtomicBool, AtomicU64, Ordering}, - mpsc, Arc, Mutex, + Arc, Mutex, OnceLock, }, thread, }; -#[derive(Clone, Copy)] -enum TrackedKind { - Definition, - Doc, +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RewriteKind { + Locations, + Documentation, + Completion, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct TrackedRequest { + token: u64, + method: String, + rewrite: RewriteKind, + original_params: Option, +} + +impl TrackedRequest { + fn new(token: u64, method: &str, rewrite: RewriteKind, original_params: Option) -> Self { + Self { + token, + method: method.to_string(), + rewrite, + original_params, + } + } +} + +#[derive(Default)] +struct SuppressedResponses { + ids: HashSet, +} + +impl SuppressedResponses { + fn insert(&mut self, id: Value) { + self.ids.insert(id); + } + + fn contains(&self, id: &Value) -> bool { + self.ids.contains(id) + } + + fn take(&mut self, id: &Value) -> bool { + self.ids.remove(id) + } +} + +type TrackedRequests = Arc>>; +type ActiveRewrites = Arc>>; +type SharedSuppressedResponses = Arc>; +type LatestWorkspaceJob = Arc>>; + +struct StdinContext { + writer: SharedWriter, + alive: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + jobs: Arc, + decompile: DecompileCoordinator, + output: Output, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, +} + +struct StdoutContext { + pending: Arc, + alive: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + decompile: DecompileCoordinator, + output: Output, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum InputRoute { + Forward, + Consumed, +} + +#[derive(Debug, PartialEq)] +enum OutputRoute { + Raw(Vec), + Value(Value), + Consumed, } fn main() { + let output = Output::start(); + log::init(output.clone()); + let args: Vec = env::args().skip(1).collect(); if args.len() < 2 { @@ -78,12 +165,17 @@ fn main() { lsp_info!("JDTLS process spawned (pid={})", child.id()); - let child_stdin = Arc::new(Mutex::new(child.stdin.take().unwrap())); + let child_stdin: SharedWriter = Arc::new(Mutex::new(Box::new(child.stdin.take().unwrap()))); let child_stdout = child.stdout.take().unwrap(); let alive = Arc::new(AtomicBool::new(true)); - let pending: Arc>>> = - Arc::new(Mutex::new(HashMap::new())); + let owned_id_prefix = format!("{proxy_id}-proxy-"); + let pending = Arc::new(PendingResponses::new()); + let decompile = DecompileCoordinator::new( + Arc::clone(&child_stdin), + Arc::clone(&pending), + owned_id_prefix.clone(), + ); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let port = listener.local_addr().unwrap().port(); @@ -95,151 +187,48 @@ fn main() { lsp_info!("HTTP server listening on 127.0.0.1:{port}"); let id_counter = Arc::new(AtomicU64::new(1)); + let job_counter = Arc::new(AtomicU64::new(1)); - // Track definition/typeDefinition/implementation and documentation request IDs - // so their responses can be intercepted and rewritten. - let tracked_ids: Arc>> = Arc::new(Mutex::new(HashMap::new())); + // Track requests whose responses may contain jdt:// URIs so they can be + // intercepted and rewritten. + let tracked_ids: TrackedRequests = Arc::new(Mutex::new(HashMap::new())); + let active_rewrites: ActiveRewrites = Arc::new(Mutex::new(HashMap::new())); + let suppressed_responses = Arc::new(Mutex::new(SuppressedResponses::default())); + let latest_workspace_job = Arc::new(Mutex::new(None::)); - // --- Thread 1: Zed stdin -> JDTLS stdin (track definition requests) --- - let stdin_writer = Arc::clone(&child_stdin); - let alive_stdin = Arc::clone(&alive); - let tracked_in = Arc::clone(&tracked_ids); - thread::spawn(move || { - let stdin = io::stdin().lock(); - let mut reader = LspReader::new(BufReader::new(stdin)); - while alive_stdin.load(Ordering::Relaxed) { - match reader.read_message() { - Ok(Some(raw)) => { - // Only requests (not notifications) carry an `id`; skip the - // JSON parse entirely for high-volume notifications like - // textDocument/didChange. - if raw_has_id(&raw) { - if let Some(msg) = parse_lsp_content(&raw) { - if let Some(method) = msg.get("method").and_then(|m| m.as_str()) { - let kind = match method { - "textDocument/definition" - | "textDocument/typeDefinition" - | "textDocument/implementation" => { - Some(TrackedKind::Definition) - } - "textDocument/hover" - | "textDocument/signatureHelp" - | "completionItem/resolve" => Some(TrackedKind::Doc), - _ => None, - }; - if let Some(kind) = kind { - if let Some(id) = msg.get("id").cloned() { - tracked_in.lock().unwrap().insert(id, kind); - } - } - } - } - } - let mut w = stdin_writer.lock().unwrap(); - if w.write_all(&raw).is_err() || w.flush().is_err() { - break; - } - } - Ok(None) | Err(_) => break, - } - } - alive_stdin.store(false, Ordering::Relaxed); - }); + // --- Thread 1: Zed stdin -> JDTLS stdin --- + let stdin_context = StdinContext { + writer: Arc::clone(&child_stdin), + alive: Arc::clone(&alive), + tracked: Arc::clone(&tracked_ids), + active: Arc::clone(&active_rewrites), + jobs: Arc::clone(&job_counter), + decompile: decompile.clone(), + output: output.clone(), + suppressed: Arc::clone(&suppressed_responses), + latest_workspace: Arc::clone(&latest_workspace_job), + }; + thread::spawn(move || run_zed_input(stdin_context)); - // --- Thread 2: JDTLS stdout -> rewrite jdt:// URIs, modify completions -> Zed stdout / resolve pending --- - let pending_out = Arc::clone(&pending); - let alive_out = Arc::clone(&alive); - let tracked_out = Arc::clone(&tracked_ids); - let decompile_writer = Arc::clone(&child_stdin); - let decompile_pending = Arc::clone(&pending); - let decompile_counter = Arc::clone(&id_counter); - let decompile_proxy_id = proxy_id.clone(); - thread::spawn(move || { - let mut reader = LspReader::new(BufReader::new(child_stdout)); - while alive_out.load(Ordering::Relaxed) { - match reader.read_message() { - Ok(Some(raw)) => { - // Fast path: notifications (no `id`) can't be responses we - // need to intercept. Forward the raw bytes without parsing. - if !raw_has_id(&raw) { - write_raw(&mut io::stdout().lock(), &raw); - continue; - } - - let Some(mut msg) = parse_lsp_content(&raw) else { - write_raw(&mut io::stdout().lock(), &raw); - continue; - }; - - // Route responses to pending HTTP requests - if let Some(id) = msg.get("id") { - if let Some(tx) = pending_out.lock().unwrap().remove(id) { - let _ = tx.send(msg); - continue; - } - } - - // Rewrite jdt:// URIs in definition or documentation responses. - // Spawns a thread so this loop stays unblocked and can route - // the java/classFileContents response back via `pending`. - if let Some(id) = msg.get("id").cloned() { - if let Some(kind) = tracked_out.lock().unwrap().remove(&id) { - let writer = Arc::clone(&decompile_writer); - let pending = Arc::clone(&decompile_pending); - let pid = decompile_proxy_id.clone(); - let counter = Arc::clone(&decompile_counter); - thread::spawn(move || { - let mut next_id = move || { - let seq = counter.fetch_add(1, Ordering::Relaxed); - Value::String(format!("{pid}-decompile-{seq}")) - }; - match kind { - TrackedKind::Definition => { - rewrite_jdt_locations( - &mut msg, - &writer, - &pending, - &mut next_id, - ); - } - TrackedKind::Doc => { - rewrite_jdt_in_strings( - &mut msg, - &writer, - &pending, - &mut next_id, - ); - sanitize_resolved_completion(&mut msg); - } - } - write_to_stdout(&msg); - }); - continue; - } - } - - // Process completion responses (sort + sanitize) in a single pass - if is_completion_response(&msg) { - process_completions(&mut msg); - write_to_stdout(&msg); - continue; - } - - // Passthrough - write_raw(&mut io::stdout().lock(), &raw); - } - Ok(None) | Err(_) => break, - } - } - alive_out.store(false, Ordering::Relaxed); - }); + // --- Thread 2: JDTLS stdout -> Zed stdout --- + let stdout_context = StdoutContext { + pending: Arc::clone(&pending), + alive: Arc::clone(&alive), + tracked: Arc::clone(&tracked_ids), + active: Arc::clone(&active_rewrites), + decompile: decompile.clone(), + output: output.clone(), + suppressed: Arc::clone(&suppressed_responses), + latest_workspace: Arc::clone(&latest_workspace_job), + }; + let stdout_thread = thread::spawn(move || run_jdtls_output(child_stdout, stdout_context)); // --- Thread 3: HTTP server for extension requests --- let http_writer = Arc::clone(&child_stdin); let http_pending = Arc::clone(&pending); let http_alive = Arc::clone(&alive); let http_id_counter = Arc::clone(&id_counter); - let http_proxy_id = proxy_id.clone(); + let http_owned_id_prefix = owned_id_prefix.clone(); thread::spawn(move || { for stream in listener.incoming() { if !http_alive.load(Ordering::Relaxed) { @@ -249,10 +238,10 @@ fn main() { let writer = Arc::clone(&http_writer); let pend = Arc::clone(&http_pending); let counter = Arc::clone(&http_id_counter); - let pid = http_proxy_id.clone(); + let owned_id_prefix = http_owned_id_prefix.clone(); thread::spawn(move || { - handle_http(stream, writer, pend, counter, &pid); + handle_http(stream, writer, pend, counter, &owned_id_prefix); }); } }); @@ -260,15 +249,854 @@ fn main() { // --- Thread 4: Parent process monitor --- spawn_parent_monitor(Arc::clone(&alive), child.id()); - // Wait for child to exit - let status = child.wait(); - lsp_info!("JDTLS process exited: {status:?}"); + // Poll so broken editor/JDTLS transport can terminate the child and release + // all pending proxy jobs instead of waiting indefinitely in `Child::wait`. + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break Ok(status), + Ok(None) if alive.load(Ordering::Relaxed) && !output.failed() => { + thread::sleep(std::time::Duration::from_millis(50)); + } + Ok(None) => { + let _ = child.kill(); + break child.wait(); + } + Err(error) => break Err(error), + } + }; alive.store(false, Ordering::Relaxed); + let _ = stdout_thread.join(); + lsp_info!("JDTLS process exited: {status:?}"); + pending.clear(); + decompile.shutdown(); + output.shutdown(); + decompile.cleanup_cache(); let _ = fs::remove_file(&port_file); } +fn run_zed_input(context: StdinContext) { + let stdin = io::stdin().lock(); + let mut reader = LspReader::new(BufReader::new(stdin)); + while context.alive.load(Ordering::Relaxed) { + let raw = match reader.read_message() { + Ok(Some(raw)) => raw, + Ok(None) | Err(_) => break, + }; + if route_zed_message(&context, &raw) == InputRoute::Forward + && !write_to_jdtls(&context.writer, &raw) + { + break; + } + } + context.alive.store(false, Ordering::Relaxed); +} + +fn route_zed_message(context: &StdinContext, raw: &[u8]) -> InputRoute { + let has_id = raw_has_id(raw); + if !has_id && !contains_subslice(raw, b"$/cancelRequest") { + return InputRoute::Forward; + } + let Some(message) = parse_lsp_content(raw) else { + return InputRoute::Forward; + }; + if message.get("method").and_then(Value::as_str) == Some("$/cancelRequest") { + return route_zed_cancellation(context, &message); + } + if has_id { + track_zed_request(context, &message); + } + InputRoute::Forward +} + +fn route_zed_cancellation(context: &StdinContext, message: &Value) -> InputRoute { + let Some(id) = message.pointer("/params/id") else { + return InputRoute::Forward; + }; + let already_suppressed = context.suppressed.lock().unwrap().contains(id); + let (tracked, active_token) = + take_request_for_cancellation(&context.tracked, &context.active, id); + if let Some(request) = tracked { + clear_latest_workspace(&context.latest_workspace, request.token); + } + let handled_locally = active_token.is_some(); + if let Some(token) = active_token { + clear_latest_workspace(&context.latest_workspace, token); + context.decompile.cancel(token); + context.output.send_value(&request_canceled(id)); + } + + if handled_locally || already_suppressed { + InputRoute::Consumed + } else { + InputRoute::Forward + } +} + +fn track_zed_request(context: &StdinContext, message: &Value) { + let token = context.jobs.fetch_add(1, Ordering::Relaxed); + let Some((id, request)) = tracked_request_for(message, token) else { + return; + }; + if request.method == "workspace/symbol" { + supersede_workspace_job(context, token); + } + + let (previous, active_token) = { + let mut tracked = context.tracked.lock().unwrap(); + let mut active = context.active.lock().unwrap(); + let previous = tracked.insert(id.clone(), request); + let active_token = active.remove(&id); + (previous, active_token) + }; + if let Some(previous) = previous { + clear_latest_workspace(&context.latest_workspace, previous.token); + } + if let Some(active_token) = active_token { + context.decompile.cancel(active_token); + } +} + +fn supersede_workspace_job(context: &StdinContext, token: u64) { + let previous = context.latest_workspace.lock().unwrap().replace(token); + let Some(previous) = previous else { + return; + }; + let (suppressed, active) = retire_request_token(&context.tracked, &context.active, previous); + if let Some(id) = suppressed { + context.suppressed.lock().unwrap().insert(id.clone()); + cancel_jdtls_request(&context.writer, &id); + } + if active { + context.decompile.cancel(previous); + } +} + +fn write_to_jdtls(writer: &SharedWriter, raw: &[u8]) -> bool { + let mut writer = writer.lock().unwrap(); + writer.write_all(raw).is_ok() && writer.flush().is_ok() +} + +fn run_jdtls_output(reader: impl io::Read, context: StdoutContext) { + let mut reader = LspReader::new(BufReader::new(reader)); + while let Ok(Some(raw)) = reader.read_message() { + match route_jdtls_message(&context, raw) { + OutputRoute::Raw(raw) => { + context.output.send_raw(raw); + } + OutputRoute::Value(message) => { + context.output.send_value(&message); + } + OutputRoute::Consumed => {} + } + } + context.alive.store(false, Ordering::Relaxed); +} + +fn route_jdtls_message(context: &StdoutContext, raw: Vec) -> OutputRoute { + // Notifications cannot be responses that the proxy needs to intercept. + if !raw_has_id(&raw) { + return OutputRoute::Raw(raw); + } + let Some(message) = parse_lsp_content(&raw) else { + return OutputRoute::Raw(raw); + }; + if context.pending.route(&message) { + return OutputRoute::Consumed; + } + if message + .get("id") + .is_some_and(|id| context.suppressed.lock().unwrap().take(id)) + { + return OutputRoute::Consumed; + } + if message.get("method").is_some() { + return OutputRoute::Raw(raw); + } + let Some(id) = message.get("id").cloned() else { + return OutputRoute::Raw(raw); + }; + let request = context.tracked.lock().unwrap().get(&id).cloned(); + let Some(request) = request else { + return OutputRoute::Raw(raw); + }; + + route_tracked_response(context, raw, message, id, request) +} + +fn route_tracked_response( + context: &StdoutContext, + raw: Vec, + mut message: Value, + id: Value, + request: TrackedRequest, +) -> OutputRoute { + if let Some(fallback) = completion_resolve_fallback(&message, &request) { + remove_tracked_request(&context.tracked, &id, request.token); + if should_log_completion_fallback() { + lsp_warn!( + "JDTLS completion resolution failed with -32603; \ + using the unresolved item, so documentation, imports, \ + commands, or additional edits may be missing" + ); + } + return OutputRoute::Value(fallback); + } + if message.get("error").is_some() { + remove_tracked_request(&context.tracked, &id, request.token); + clear_latest_workspace(&context.latest_workspace, request.token); + return OutputRoute::Raw(raw); + } + if request.rewrite == RewriteKind::Completion { + remove_tracked_request(&context.tracked, &id, request.token); + process_completions(&mut message); + return OutputRoute::Value(message); + } + + queue_rewrite_response(context, raw, message, id, request) +} + +fn queue_rewrite_response( + context: &StdoutContext, + raw: Vec, + message: Value, + id: Value, + request: TrackedRequest, +) -> OutputRoute { + let sanitize_completion = request.method == "completionItem/resolve"; + let mode = match request.rewrite { + RewriteKind::Locations => RewriteMode::Locations, + RewriteKind::Documentation => RewriteMode::Strings, + RewriteKind::Completion => unreachable!(), + }; + let priority = if request.method == "workspace/symbol" { + Priority::Bulk + } else { + Priority::Interactive + }; + let deadline = std::time::Instant::now() + rewrite_timeout(&request.method); + let workspace_job = (request.method == "workspace/symbol") + .then(|| (Arc::clone(&context.latest_workspace), request.token)); + if !activate_rewrite(&context.tracked, &context.active, &id, request.token) { + return OutputRoute::Raw(raw); + } + if context.decompile.is_canceled(request.token) { + context.active.lock().unwrap().remove(&id); + context.decompile.consume_cancellation(request.token); + clear_latest_workspace(&context.latest_workspace, request.token); + return OutputRoute::Consumed; + } + + let token = request.token; + let output = context.output.clone(); + let active = Arc::clone(&context.active); + let completion_coordinator = context.decompile.clone(); + let job = RewriteJob { + token, + message, + mode, + priority, + deadline, + complete: Box::new(move |mut message| { + if sanitize_completion { + sanitize_resolved_completion(&mut message); + } + let owns_response = active.lock().unwrap().remove(&id) == Some(token); + completion_coordinator.consume_cancellation(token); + if owns_response { + output.send_value(&message); + } + if let Some((latest, token)) = workspace_job { + clear_latest_workspace(&latest, token); + } + }), + }; + if let Err(job) = context.decompile.submit(job) { + let RewriteJob { + message, complete, .. + } = job; + complete(message); + } + OutputRoute::Consumed +} + // --- Utilities --- fn hex_encode(s: &str) -> String { s.as_bytes().iter().map(|b| format!("{b:02x}")).collect() } + +fn clear_latest_workspace(latest: &Mutex>, token: u64) { + let mut latest = latest.lock().unwrap(); + if *latest == Some(token) { + *latest = None; + } +} + +fn request_canceled(id: &Value) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32800, + "message": "Request cancelled." + } + }) +} + +fn cancel_jdtls_request(writer: &SharedWriter, id: &Value) { + let cancel = encode_lsp(&json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": id } + })); + let mut writer = writer.lock().unwrap(); + let _ = writer.write_all(cancel.as_bytes()); + let _ = writer.flush(); +} + +fn take_request_for_cancellation( + tracked: &Mutex>, + active: &Mutex>, + id: &Value, +) -> (Option, Option) { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + (tracked.remove(id), active.remove(id)) +} + +fn retire_request_token( + tracked: &Mutex>, + active: &Mutex>, + token: u64, +) -> (Option, bool) { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + let tracked_id = tracked + .iter() + .find_map(|(id, request)| (request.token == token).then(|| id.clone())); + if let Some(id) = &tracked_id { + tracked.remove(id); + } + let active_id = active + .iter() + .find_map(|(id, active_token)| (*active_token == token).then(|| id.clone())); + if let Some(id) = &active_id { + active.remove(id); + } + (tracked_id, active_id.is_some()) +} + +fn activate_rewrite( + tracked: &Mutex>, + active: &Mutex>, + id: &Value, + token: u64, +) -> bool { + let mut tracked = tracked.lock().unwrap(); + let mut active = active.lock().unwrap(); + if tracked.get(id).map(|request| request.token) != Some(token) { + return false; + } + tracked.remove(id); + active.insert(id.clone(), token); + true +} + +fn remove_tracked_request(tracked: &Mutex>, id: &Value, token: u64) { + let mut tracked = tracked.lock().unwrap(); + if tracked.get(id).map(|request| request.token) == Some(token) { + tracked.remove(id); + } +} + +fn should_log_completion_fallback() -> bool { + static LAST_WARNING: OnceLock>> = OnceLock::new(); + let now = std::time::Instant::now(); + let mut last = LAST_WARNING + .get_or_init(|| Mutex::new(None)) + .lock() + .unwrap(); + if last + .is_some_and(|previous| now.duration_since(previous) < std::time::Duration::from_secs(60)) + { + return false; + } + *last = Some(now); + true +} + +fn tracked_request_for(msg: &Value, token: u64) -> Option<(Value, TrackedRequest)> { + let method = msg.get("method")?.as_str()?; + let id = msg.get("id")?.clone(); + let rewrite = match method { + "textDocument/completion" => RewriteKind::Completion, + "textDocument/definition" + | "textDocument/declaration" + | "textDocument/typeDefinition" + | "textDocument/implementation" + | "textDocument/references" + | "textDocument/prepareCallHierarchy" + | "callHierarchy/incomingCalls" + | "callHierarchy/outgoingCalls" + | "textDocument/prepareTypeHierarchy" + | "typeHierarchy/supertypes" + | "typeHierarchy/subtypes" + | "workspace/symbol" => RewriteKind::Locations, + "textDocument/hover" | "textDocument/signatureHelp" | "completionItem/resolve" => { + RewriteKind::Documentation + } + _ => return None, + }; + + let original_params = if method == "completionItem/resolve" { + msg.get("params") + .filter(|params| params.is_object()) + .cloned() + } else { + None + }; + + Some(( + id, + TrackedRequest::new(token, method, rewrite, original_params), + )) +} + +fn completion_resolve_fallback(message: &Value, request: &TrackedRequest) -> Option { + if request.method != "completionItem/resolve" + || message.pointer("/error/code").and_then(Value::as_i64) != Some(-32603) + { + return None; + } + let item = request.original_params.as_ref()?; + Some(json!({ + "jsonrpc": "2.0", + "id": message.get("id")?.clone(), + "result": item, + })) +} + +fn rewrite_timeout(method: &str) -> std::time::Duration { + match method { + "textDocument/hover" | "textDocument/signatureHelp" | "completionItem/resolve" => { + std::time::Duration::from_millis(500) + } + "workspace/symbol" => std::time::Duration::from_secs(1), + _ => std::time::Duration::from_secs(2), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + struct RoutingFixture { + writer: SharedWriter, + alive: Arc, + jobs: Arc, + pending: Arc, + tracked: TrackedRequests, + active: ActiveRewrites, + suppressed: SharedSuppressedResponses, + latest_workspace: LatestWorkspaceJob, + decompile: DecompileCoordinator, + output: Output, + } + + impl RoutingFixture { + fn new() -> Self { + let writer: SharedWriter = Arc::new(Mutex::new(Box::new(Vec::::new()))); + let pending = Arc::new(PendingResponses::new()); + let decompile = DecompileCoordinator::new( + Arc::clone(&writer), + Arc::clone(&pending), + "routing-test-".to_string(), + ); + Self { + writer, + alive: Arc::new(AtomicBool::new(true)), + jobs: Arc::new(AtomicU64::new(1)), + pending, + tracked: Arc::new(Mutex::new(HashMap::new())), + active: Arc::new(Mutex::new(HashMap::new())), + suppressed: Arc::new(Mutex::new(SuppressedResponses::default())), + latest_workspace: Arc::new(Mutex::new(None)), + decompile, + output: Output::start(), + } + } + + fn stdin_context(&self) -> StdinContext { + StdinContext { + writer: Arc::clone(&self.writer), + alive: Arc::clone(&self.alive), + tracked: Arc::clone(&self.tracked), + active: Arc::clone(&self.active), + jobs: Arc::clone(&self.jobs), + decompile: self.decompile.clone(), + output: self.output.clone(), + suppressed: Arc::clone(&self.suppressed), + latest_workspace: Arc::clone(&self.latest_workspace), + } + } + + fn stdout_context(&self) -> StdoutContext { + StdoutContext { + pending: Arc::clone(&self.pending), + alive: Arc::clone(&self.alive), + tracked: Arc::clone(&self.tracked), + active: Arc::clone(&self.active), + decompile: self.decompile.clone(), + output: self.output.clone(), + suppressed: Arc::clone(&self.suppressed), + latest_workspace: Arc::clone(&self.latest_workspace), + } + } + } + + impl Drop for RoutingFixture { + fn drop(&mut self) { + self.decompile.shutdown(); + self.output.shutdown(); + } + } + + fn frame(value: &Value) -> Vec { + encode_lsp(value).into_bytes() + } + + #[test] + fn stdin_router_tracks_requests_and_forwards_cancellation_before_rewrite() { + let fixture = RoutingFixture::new(); + let context = fixture.stdin_context(); + let request = json!({ + "jsonrpc": "2.0", + "id": 7, + "method": "textDocument/definition", + "params": {} + }); + + assert_eq!( + route_zed_message(&context, &frame(&request)), + InputRoute::Forward + ); + assert_eq!( + fixture.tracked.lock().unwrap().get(&json!(7)), + Some(&TrackedRequest::new( + 1, + "textDocument/definition", + RewriteKind::Locations, + None + )) + ); + + let cancellation = json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": { "id": 7 } + }); + assert_eq!( + route_zed_message(&context, &frame(&cancellation)), + InputRoute::Forward + ); + assert!(fixture.tracked.lock().unwrap().is_empty()); + } + + #[test] + fn stdout_router_returns_processed_completion_values() { + let fixture = RoutingFixture::new(); + let context = fixture.stdout_context(); + fixture.tracked.lock().unwrap().insert( + json!(9), + TrackedRequest::new(1, "textDocument/completion", RewriteKind::Completion, None), + ); + let response = json!({ + "jsonrpc": "2.0", + "id": 9, + "result": [{ + "kind": 15, + "textEditText": "$TM_SELECTED_TEXT.field" + }] + }); + + assert_eq!( + route_jdtls_message(&context, frame(&response)), + OutputRoute::Value(json!({ + "jsonrpc": "2.0", + "id": 9, + "result": [{ + "kind": 15, + "textEditText": ".field" + }] + })) + ); + assert!(fixture.tracked.lock().unwrap().is_empty()); + } + + #[test] + fn stdout_router_consumes_pending_and_suppressed_responses() { + let fixture = RoutingFixture::new(); + let context = fixture.stdout_context(); + let pending_id = json!("proxy-request"); + let receiver = fixture.pending.register(pending_id.clone()); + let pending_response = json!({ "jsonrpc": "2.0", "id": pending_id, "result": "ok" }); + + assert_eq!( + route_jdtls_message(&context, frame(&pending_response)), + OutputRoute::Consumed + ); + assert_eq!(receiver.recv().unwrap(), pending_response); + + let suppressed_id = json!("superseded-request"); + fixture + .suppressed + .lock() + .unwrap() + .insert(suppressed_id.clone()); + let suppressed_response = json!({ "jsonrpc": "2.0", "id": suppressed_id, "result": null }); + assert_eq!( + route_jdtls_message(&context, frame(&suppressed_response)), + OutputRoute::Consumed + ); + assert!(!fixture + .suppressed + .lock() + .unwrap() + .contains(&json!("superseded-request"))); + } + + #[test] + fn stdio_routers_preserve_unhandled_raw_frames() { + let fixture = RoutingFixture::new(); + let stdin_context = fixture.stdin_context(); + let stdout_context = fixture.stdout_context(); + let malformed = b"Content-Length: 6\r\n\r\n{\"id\":".to_vec(); + + assert_eq!( + route_zed_message(&stdin_context, &malformed), + InputRoute::Forward + ); + assert_eq!( + route_jdtls_message(&stdout_context, malformed.clone()), + OutputRoute::Raw(malformed) + ); + } + + #[test] + fn tracks_location_response_methods() { + let methods = [ + "textDocument/definition", + "textDocument/declaration", + "textDocument/typeDefinition", + "textDocument/implementation", + "textDocument/references", + "textDocument/prepareCallHierarchy", + "callHierarchy/incomingCalls", + "callHierarchy/outgoingCalls", + "textDocument/prepareTypeHierarchy", + "typeHierarchy/supertypes", + "typeHierarchy/subtypes", + "workspace/symbol", + ]; + + for method in methods { + let request = json!({ + "jsonrpc": "2.0", + "id": 7, + "method": method, + "params": {} + }); + + assert_eq!( + tracked_request_for(&request, 42), + Some(( + json!(7), + TrackedRequest::new(42, method, RewriteKind::Locations, None) + )), + "{method} should rewrite location URIs" + ); + } + } + + #[test] + fn tracks_completion_resolve_documentation() { + let request = json!({ + "jsonrpc": "2.0", + "id": "resolve-1", + "method": "completionItem/resolve", + "params": { "label": "String" } + }); + + assert_eq!( + tracked_request_for(&request, 43), + Some(( + json!("resolve-1"), + TrackedRequest::new( + 43, + "completionItem/resolve", + RewriteKind::Documentation, + Some(json!({ "label": "String" })) + ) + )) + ); + } + + #[test] + fn ignores_untracked_requests_and_notifications() { + let request = json!({ + "jsonrpc": "2.0", + "id": 8, + "method": "textDocument/rename", + "params": {} + }); + let notification = json!({ + "jsonrpc": "2.0", + "method": "textDocument/didChange", + "params": {} + }); + + assert_eq!(tracked_request_for(&request, 1), None); + assert_eq!(tracked_request_for(¬ification, 2), None); + } + + #[test] + fn tracks_completion_responses_by_method() { + let request = json!({ + "jsonrpc": "2.0", + "id": 9, + "method": "textDocument/completion", + "params": {} + }); + + assert_eq!( + tracked_request_for(&request, 44), + Some(( + json!(9), + TrackedRequest::new(44, "textDocument/completion", RewriteKind::Completion, None) + )) + ); + } + + #[test] + fn activating_rewrite_moves_request_atomically() { + let id = json!(9); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(44, "textDocument/definition", RewriteKind::Locations, None), + )])); + let active = Mutex::new(HashMap::new()); + + assert!(activate_rewrite(&tracked, &active, &id, 44)); + assert!(tracked.lock().unwrap().is_empty()); + assert_eq!(active.lock().unwrap().get(&id), Some(&44)); + } + + #[test] + fn retiring_tracked_request_returns_id_for_late_suppression() { + let id = json!("workspace-1"); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(45, "workspace/symbol", RewriteKind::Locations, None), + )])); + let active = Mutex::new(HashMap::new()); + + assert_eq!( + retire_request_token(&tracked, &active, 45), + (Some(id), false) + ); + assert!(tracked.lock().unwrap().is_empty()); + } + + #[test] + fn suppressed_response_ids_remain_owned_until_the_response_arrives() { + let mut suppressed = SuppressedResponses::default(); + for id in 0..=1024 { + suppressed.insert(json!(id)); + } + + assert!(suppressed.contains(&json!(0))); + assert!(suppressed.take(&json!(0))); + assert!(!suppressed.contains(&json!(0))); + assert!(suppressed.take(&json!(1024))); + } + + #[test] + fn cancellation_response_preserves_request_id() { + assert_eq!( + request_canceled(&json!("request-1")), + json!({ + "jsonrpc": "2.0", + "id": "request-1", + "error": { + "code": -32800, + "message": "Request cancelled." + } + }) + ); + } + + #[test] + fn cancellation_removes_unanswered_tracking_state() { + let id = json!("hover-1"); + let tracked = Mutex::new(HashMap::from([( + id.clone(), + TrackedRequest::new(46, "textDocument/hover", RewriteKind::Documentation, None), + )])); + let active = Mutex::new(HashMap::new()); + + let (request, active_token) = take_request_for_cancellation(&tracked, &active, &id); + + assert_eq!(request.map(|request| request.token), Some(46)); + assert_eq!(active_token, None); + assert!(tracked.lock().unwrap().is_empty()); + } + + #[test] + fn falls_back_only_for_internal_completion_resolve_errors() { + let request = TrackedRequest::new( + 1, + "completionItem/resolve", + RewriteKind::Documentation, + Some(json!({ "label": "value", "kind": 6 })), + ); + let internal = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32603, "message": "Invalid completion proposal" } + }); + let canceled = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32800, "message": "Request cancelled" } + }); + + assert_eq!( + completion_resolve_fallback(&internal, &request), + Some(json!({ + "jsonrpc": "2.0", + "id": 5, + "result": { "label": "value", "kind": 6 } + })) + ); + assert_eq!(completion_resolve_fallback(&canceled, &request), None); + } + + #[test] + fn malformed_completion_resolve_params_do_not_fallback() { + let request = json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "completionItem/resolve", + "params": null + }); + let (_, tracked) = tracked_request_for(&request, 1).unwrap(); + let error = json!({ + "jsonrpc": "2.0", + "id": 5, + "error": { "code": -32603, "message": "Internal error" } + }); + + assert_eq!(completion_resolve_fallback(&error, &tracked), None); + } +} diff --git a/proxy/src/output.rs b/proxy/src/output.rs new file mode 100644 index 0000000..91118b0 --- /dev/null +++ b/proxy/src/output.rs @@ -0,0 +1,210 @@ +use proxy_common::encode_lsp; +use serde::Serialize; +use std::{ + io::{self, Write}, + sync::{ + atomic::{AtomicBool, Ordering}, + mpsc, Arc, Mutex, + }, + thread::{self, JoinHandle}, + time::Duration, +}; + +const SHUTDOWN_DRAIN_TIMEOUT: Duration = Duration::from_millis(250); + +#[derive(Clone)] +pub struct Output { + inner: Arc, +} + +struct Inner { + sender: Mutex>>, + worker: Mutex>>, + done: Mutex>>, + failed: Arc, +} + +enum Message { + Frame(Vec), + Shutdown, +} + +impl Output { + pub fn start() -> Self { + Self::start_with_writer(io::stdout()) + } + + fn start_with_writer(writer: impl Write + Send + 'static) -> Self { + let (sender, receiver) = mpsc::channel(); + let (done_sender, done) = mpsc::channel(); + let failed = Arc::new(AtomicBool::new(false)); + let writer_failed = Arc::clone(&failed); + let worker = thread::spawn(move || { + write_loop(writer, receiver, writer_failed); + let _ = done_sender.send(()); + }); + Self { + inner: Arc::new(Inner { + sender: Mutex::new(Some(sender)), + worker: Mutex::new(Some(worker)), + done: Mutex::new(Some(done)), + failed, + }), + } + } + + pub fn send_raw(&self, raw: Vec) -> bool { + let sent = self + .inner + .sender + .lock() + .unwrap() + .as_ref() + .is_some_and(|sender| sender.send(Message::Frame(raw)).is_ok()); + if !sent { + self.inner.failed.store(true, Ordering::Relaxed); + } + sent + } + + pub fn send_value(&self, value: &impl Serialize) -> bool { + self.send_raw(encode_lsp(value).into_bytes()) + } + + pub fn failed(&self) -> bool { + self.inner.failed.load(Ordering::Relaxed) + } + + /// Stops accepting frames and gives stdout a bounded interval to drain. + /// A blocked editor must not prevent the proxy process from terminating. + pub fn shutdown(&self) { + if let Some(sender) = self.inner.sender.lock().unwrap().take() { + let _ = sender.send(Message::Shutdown); + } + let drained = self + .inner + .done + .lock() + .unwrap() + .take() + .is_none_or(|done| done.recv_timeout(SHUTDOWN_DRAIN_TIMEOUT).is_ok()); + if let Some(worker) = self.inner.worker.lock().unwrap().take() { + if drained { + let _ = worker.join(); + } + } + } +} + +fn write_loop(mut writer: impl Write, receiver: mpsc::Receiver, failed: Arc) { + while let Ok(message) = receiver.recv() { + match message { + Message::Frame(frame) => { + if writer.write_all(&frame).is_err() || writer.flush().is_err() { + failed.store(true, Ordering::Relaxed); + break; + } + } + Message::Shutdown => break, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Clone, Default)] + struct SharedWriter(Arc>>); + + impl Write for SharedWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + struct FailingWriter; + + impl Write for FailingWriter { + fn write(&mut self, _buf: &[u8]) -> io::Result { + Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed")) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + struct BlockingWriter { + started: mpsc::Sender<()>, + release: mpsc::Receiver<()>, + } + + impl Write for BlockingWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let _ = self.started.send(()); + let _ = self.release.recv(); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn serializes_complete_frames_in_submission_order() { + let writer = SharedWriter::default(); + let bytes = Arc::clone(&writer.0); + let output = Output::start_with_writer(writer); + + assert!(output.send_raw(b"first".to_vec())); + assert!(output.send_raw(b"second".to_vec())); + output.shutdown(); + + assert_eq!(*bytes.lock().unwrap(), b"firstsecond"); + assert!(!output.send_raw(b"third".to_vec())); + } + + #[test] + fn reports_writer_failure() { + let output = Output::start_with_writer(FailingWriter); + assert!(output.send_raw(b"message".to_vec())); + + for _ in 0..100 { + if output.failed() { + break; + } + thread::sleep(std::time::Duration::from_millis(1)); + } + + assert!(output.failed()); + output.shutdown(); + } + + #[test] + fn shutdown_does_not_wait_forever_for_blocked_stdout() { + let (started_sender, started_receiver) = mpsc::channel(); + let (release_sender, release_receiver) = mpsc::channel(); + let output = Output::start_with_writer(BlockingWriter { + started: started_sender, + release: release_receiver, + }); + assert!(output.send_raw(b"blocked".to_vec())); + started_receiver + .recv_timeout(Duration::from_secs(1)) + .unwrap(); + + let started = std::time::Instant::now(); + output.shutdown(); + + assert!(started.elapsed() < Duration::from_secs(1)); + let _ = release_sender.send(()); + } +} diff --git a/proxy/src/pending.rs b/proxy/src/pending.rs new file mode 100644 index 0000000..308c8c1 --- /dev/null +++ b/proxy/src/pending.rs @@ -0,0 +1,138 @@ +use serde_json::Value; +use std::{ + collections::{HashMap, HashSet}, + sync::{mpsc, Mutex}, +}; + +pub struct PendingResponses { + state: Mutex, +} + +#[derive(Default)] +struct State { + pending: HashMap>, + owned: HashSet, +} + +impl PendingResponses { + pub fn new() -> Self { + Self { + state: Mutex::new(State::default()), + } + } + + pub fn register(&self, id: Value) -> mpsc::Receiver { + let (sender, receiver) = mpsc::channel(); + let mut state = self.state.lock().unwrap(); + state.owned.insert(id.clone()); + state.pending.insert(id, sender); + receiver + } + + pub fn remove(&self, id: &Value) { + let mut state = self.state.lock().unwrap(); + state.pending.remove(id); + } + + /// Returns true when the response belongs to the proxy and must not be + /// forwarded to the editor, including responses that arrive after timeout. + pub fn route(&self, message: &Value) -> bool { + if message.get("method").is_some() { + return false; + } + let Some(id) = message.get("id") else { + return false; + }; + + let (sender, owned) = { + let mut state = self.state.lock().unwrap(); + let sender = state.pending.remove(id); + let owned = state.owned.remove(id); + (sender, owned) + }; + if let Some(sender) = sender { + let _ = sender.send(message.clone()); + return true; + } + + owned + } + + pub fn clear(&self) { + let mut state = self.state.lock().unwrap(); + state.pending.clear(); + state.owned.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn routes_each_pending_response_once() { + let pending = PendingResponses::new(); + let id = json!("proxy-owned-1"); + let receiver = pending.register(id.clone()); + let response = json!({ "jsonrpc": "2.0", "id": id, "result": "ok" }); + + assert!(pending.route(&response)); + assert_eq!(receiver.recv().unwrap(), response); + assert!(!pending.route(&response)); + } + + #[test] + fn swallows_late_owned_responses_but_not_client_responses() { + let pending = PendingResponses::new(); + let retired = json!("proxy-owned-late"); + pending.register(retired.clone()); + pending.remove(&retired); + + let response = json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-late", + "result": null + }); + assert!(pending.route(&response)); + assert!(!pending.route(&response)); + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": 7, + "result": null + }))); + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-server-request", + "method": "workspace/applyEdit", + "params": {} + }))); + } + + #[test] + fn does_not_claim_unregistered_prefixed_ids() { + let pending = PendingResponses::new(); + + assert!(!pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-editor-request", + "result": null + }))); + } + + #[test] + fn ownership_is_retained_until_a_response_arrives() { + let pending = PendingResponses::new(); + for id in 0..=1024 { + let id = json!(format!("proxy-owned-{id}")); + pending.register(id.clone()); + pending.remove(&id); + } + + assert!(pending.route(&json!({ + "jsonrpc": "2.0", + "id": "proxy-owned-0", + "result": null + }))); + } +}