| 1 | //! Background discovery for composer `@`-mention completions. |
| 2 | //! |
| 3 | //! Filesystem traversal never belongs on the TUI thread. This module owns one |
| 4 | //! serialized worker per composer, a single coalescing request slot, and a |
| 5 | //! generation token that prevents a superseded scan from publishing results. |
| 6 | |
| 7 | use std::path::PathBuf; |
| 8 | use std::sync::{ |
| 9 | Arc, |
| 10 | atomic::{AtomicBool, AtomicU64, Ordering}, |
| 11 | }; |
| 12 | use std::thread; |
| 13 | use std::time::{Duration, Instant}; |
| 14 | |
| 15 | use parking_lot::{Condvar, Mutex}; |
| 16 | |
| 17 | use crate::working_set::Workspace; |
| 18 | |
| 19 | /// Keep discovery memory and background work bounded even when the workspace |
| 20 | /// is an unignored drive root. The popup itself renders far fewer rows, but a |
| 21 | /// larger cached pool preserves useful fuzzy matching across keystrokes. |
| 22 | pub(crate) const MAX_MENTION_DISCOVERY_CANDIDATES: usize = 20_000; |
| 23 | |
| 24 | const MENTION_DISCOVERY_CACHE_TTL: Duration = Duration::from_secs(4); |
| 25 | |
| 26 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 27 | pub(crate) enum MentionDiscoveryBehavior { |
| 28 | Fuzzy, |
| 29 | Browser { partial: String }, |
| 30 | } |
| 31 | |
| 32 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 33 | pub(crate) struct MentionDiscoveryKey { |
| 34 | pub workspace: PathBuf, |
| 35 | pub cwd: Option<PathBuf>, |
| 36 | pub walk_depth: usize, |
| 37 | pub follow_links: bool, |
| 38 | pub behavior: MentionDiscoveryBehavior, |
| 39 | } |
| 40 | |
| 41 | impl MentionDiscoveryKey { |
| 42 | pub(crate) fn fuzzy( |
| 43 | workspace: PathBuf, |
| 44 | cwd: Option<PathBuf>, |
| 45 | walk_depth: usize, |
| 46 | follow_links: bool, |
| 47 | ) -> Self { |
| 48 | Self { |
| 49 | workspace, |
| 50 | cwd, |
| 51 | walk_depth, |
| 52 | follow_links, |
| 53 | behavior: MentionDiscoveryBehavior::Fuzzy, |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | pub(crate) fn browser( |
| 58 | workspace: PathBuf, |
| 59 | cwd: Option<PathBuf>, |
| 60 | walk_depth: usize, |
| 61 | follow_links: bool, |
| 62 | partial: String, |
| 63 | ) -> Self { |
| 64 | Self { |
| 65 | workspace, |
| 66 | cwd, |
| 67 | walk_depth, |
| 68 | follow_links, |
| 69 | behavior: MentionDiscoveryBehavior::Browser { partial }, |
| 70 | } |
| 71 | } |
| 72 | } |
| 73 | |
| 74 | #[derive(Debug)] |
| 75 | struct MentionDiscoveryRequest { |
| 76 | generation: u64, |
| 77 | key: MentionDiscoveryKey, |
| 78 | } |
| 79 | |
| 80 | #[derive(Debug)] |
| 81 | struct MentionDiscoveryResult { |
| 82 | generation: u64, |
| 83 | key: MentionDiscoveryKey, |
| 84 | entries: Vec<String>, |
| 85 | collected_at: Instant, |
| 86 | } |
| 87 | |
| 88 | struct WorkerShared { |
| 89 | pending: Mutex<Option<MentionDiscoveryRequest>>, |
| 90 | wake: Condvar, |
| 91 | result: Mutex<Option<MentionDiscoveryResult>>, |
| 92 | latest_generation: AtomicU64, |
| 93 | closed: AtomicBool, |
| 94 | } |
| 95 | |
| 96 | type MentionScanner = dyn Fn(&MentionDiscoveryKey, &dyn Fn() -> bool) -> Vec<String> + Send + Sync; |
| 97 | |
| 98 | struct MentionDiscoveryWorker { |
| 99 | shared: Arc<WorkerShared>, |
| 100 | } |
| 101 | |
| 102 | impl MentionDiscoveryWorker { |
| 103 | fn spawn(scanner: Arc<MentionScanner>) -> std::io::Result<Self> { |
| 104 | let shared = Arc::new(WorkerShared { |
| 105 | pending: Mutex::new(None), |
| 106 | wake: Condvar::new(), |
| 107 | result: Mutex::new(None), |
| 108 | latest_generation: AtomicU64::new(0), |
| 109 | closed: AtomicBool::new(false), |
| 110 | }); |
| 111 | let thread_shared = Arc::clone(&shared); |
| 112 | thread::Builder::new() |
| 113 | .name("codewhale-mention-discovery".to_string()) |
| 114 | .spawn(move || worker_loop(&thread_shared, &scanner))?; |
| 115 | Ok(Self { shared }) |
| 116 | } |
| 117 | |
| 118 | /// Replace the pending slot without waiting on filesystem work. The worker |
| 119 | /// only holds this mutex long enough to take one request; scanning happens |
| 120 | /// after it is released, so brief lock contention must not drop the request. |
| 121 | fn submit(&self, request: MentionDiscoveryRequest) -> bool { |
| 122 | self.shared |
| 123 | .latest_generation |
| 124 | .store(request.generation, Ordering::Release); |
| 125 | let mut pending = self.shared.pending.lock(); |
| 126 | *pending = Some(request); |
| 127 | drop(pending); |
| 128 | self.shared.wake.notify_one(); |
| 129 | true |
| 130 | } |
| 131 | |
| 132 | fn take_result(&self) -> Option<MentionDiscoveryResult> { |
| 133 | self.shared |
| 134 | .result |
| 135 | .try_lock() |
| 136 | .and_then(|mut result| result.take()) |
| 137 | } |
| 138 | |
| 139 | fn cancel(&self, generation: u64) { |
| 140 | self.shared |
| 141 | .latest_generation |
| 142 | .store(generation, Ordering::Release); |
| 143 | if let Some(mut pending) = self.shared.pending.try_lock() { |
| 144 | *pending = None; |
| 145 | } |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | impl Drop for MentionDiscoveryWorker { |
| 150 | fn drop(&mut self) { |
| 151 | self.shared.closed.store(true, Ordering::Release); |
| 152 | self.shared.latest_generation.fetch_add(1, Ordering::AcqRel); |
| 153 | self.shared.wake.notify_all(); |
| 154 | // Never join here: a filesystem call already in progress may be slow. |
| 155 | // Dropping the JoinHandle detached it at spawn time, and the worker |
| 156 | // exits as soon as that call returns and observes `closed`. |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | fn worker_loop(shared: &WorkerShared, scanner: &Arc<MentionScanner>) { |
| 161 | loop { |
| 162 | let request = { |
| 163 | let mut pending = shared.pending.lock(); |
| 164 | while pending.is_none() && !shared.closed.load(Ordering::Acquire) { |
| 165 | shared.wake.wait(&mut pending); |
| 166 | } |
| 167 | if shared.closed.load(Ordering::Acquire) { |
| 168 | return; |
| 169 | } |
| 170 | pending.take() |
| 171 | }; |
| 172 | let Some(request) = request else { |
| 173 | continue; |
| 174 | }; |
| 175 | let generation = request.generation; |
| 176 | let cancelled = || { |
| 177 | shared.closed.load(Ordering::Acquire) |
| 178 | || shared.latest_generation.load(Ordering::Acquire) != generation |
| 179 | }; |
| 180 | if cancelled() { |
| 181 | continue; |
| 182 | } |
| 183 | let entries = scanner(&request.key, &cancelled); |
| 184 | if cancelled() { |
| 185 | continue; |
| 186 | } |
| 187 | *shared.result.lock() = Some(MentionDiscoveryResult { |
| 188 | generation, |
| 189 | key: request.key, |
| 190 | entries, |
| 191 | collected_at: Instant::now(), |
| 192 | }); |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | fn filesystem_scanner(key: &MentionDiscoveryKey, cancelled: &dyn Fn() -> bool) -> Vec<String> { |
| 197 | let workspace = Workspace::with_cwd_depth_and_follow_links( |
| 198 | key.workspace.clone(), |
| 199 | key.cwd.clone(), |
| 200 | key.walk_depth, |
| 201 | key.follow_links, |
| 202 | ); |
| 203 | match &key.behavior { |
| 204 | MentionDiscoveryBehavior::Fuzzy => { |
| 205 | workspace.completion_discovery_candidates(MAX_MENTION_DISCOVERY_CANDIDATES, cancelled) |
| 206 | } |
| 207 | MentionDiscoveryBehavior::Browser { partial } => workspace |
| 208 | .browser_completion_discovery_candidates( |
| 209 | partial, |
| 210 | MAX_MENTION_DISCOVERY_CANDIDATES, |
| 211 | cancelled, |
| 212 | ), |
| 213 | } |
| 214 | } |
| 215 | |
| 216 | #[derive(Debug)] |
| 217 | struct CachedMentionDiscovery { |
| 218 | key: MentionDiscoveryKey, |
| 219 | entries: Vec<String>, |
| 220 | collected_at: Instant, |
| 221 | } |
| 222 | |
| 223 | /// UI-owned handle for the serialized mention-discovery worker. |
| 224 | /// |
| 225 | /// `ensure_requested`, `poll`, and `cached_entries` never wait on filesystem |
| 226 | /// discovery. The only synchronous work on the UI thread is key comparison, |
| 227 | /// cloning an already-bounded in-memory result, and briefly locking a request |
| 228 | /// slot that the worker never holds while scanning. |
| 229 | pub(crate) struct MentionDiscovery { |
| 230 | scanner: Arc<MentionScanner>, |
| 231 | worker: Option<MentionDiscoveryWorker>, |
| 232 | generation: u64, |
| 233 | in_flight: Option<(u64, MentionDiscoveryKey)>, |
| 234 | cached: Option<CachedMentionDiscovery>, |
| 235 | } |
| 236 | |
| 237 | impl Default for MentionDiscovery { |
| 238 | fn default() -> Self { |
| 239 | Self::new(Arc::new(filesystem_scanner)) |
| 240 | } |
| 241 | } |
| 242 | |
| 243 | impl MentionDiscovery { |
| 244 | fn new(scanner: Arc<MentionScanner>) -> Self { |
| 245 | Self { |
| 246 | scanner, |
| 247 | worker: None, |
| 248 | generation: 0, |
| 249 | in_flight: None, |
| 250 | cached: None, |
| 251 | } |
| 252 | } |
| 253 | |
| 254 | #[cfg(test)] |
| 255 | pub(crate) fn with_scanner<F>(scanner: F) -> Self |
| 256 | where |
| 257 | F: Fn(&MentionDiscoveryKey, &dyn Fn() -> bool) -> Vec<String> + Send + Sync + 'static, |
| 258 | { |
| 259 | Self::new(Arc::new(scanner)) |
| 260 | } |
| 261 | |
| 262 | /// Start or refresh discovery for `key`. Returns immediately even if a |
| 263 | /// previous generation is stalled inside one filesystem read. |
| 264 | pub(crate) fn ensure_requested(&mut self, key: MentionDiscoveryKey) { |
| 265 | let cache_is_fresh = self.cached.as_ref().is_some_and(|cached| { |
| 266 | cached.key == key && cached.collected_at.elapsed() < MENTION_DISCOVERY_CACHE_TTL |
| 267 | }); |
| 268 | if cache_is_fresh { |
| 269 | if self |
| 270 | .in_flight |
| 271 | .as_ref() |
| 272 | .is_some_and(|(_, pending_key)| *pending_key != key) |
| 273 | { |
| 274 | self.cancel(); |
| 275 | } |
| 276 | return; |
| 277 | } |
| 278 | if self |
| 279 | .in_flight |
| 280 | .as_ref() |
| 281 | .is_some_and(|(_, pending_key)| *pending_key == key) |
| 282 | { |
| 283 | return; |
| 284 | } |
| 285 | |
| 286 | if self.worker.is_none() { |
| 287 | match MentionDiscoveryWorker::spawn(Arc::clone(&self.scanner)) { |
| 288 | Ok(worker) => self.worker = Some(worker), |
| 289 | Err(err) => { |
| 290 | tracing::warn!(error = %err, "failed to start @-mention discovery worker"); |
| 291 | return; |
| 292 | } |
| 293 | } |
| 294 | } |
| 295 | |
| 296 | let generation = self.next_generation(); |
| 297 | let request = MentionDiscoveryRequest { |
| 298 | generation, |
| 299 | key: key.clone(), |
| 300 | }; |
| 301 | let submitted = self |
| 302 | .worker |
| 303 | .as_ref() |
| 304 | .is_some_and(|worker| worker.submit(request)); |
| 305 | if submitted { |
| 306 | self.in_flight = Some((generation, key)); |
| 307 | } |
| 308 | } |
| 309 | |
| 310 | /// Apply one completed result if it still belongs to the current |
| 311 | /// generation. Returns `true` when the visible cache changed. |
| 312 | pub(crate) fn poll(&mut self) -> bool { |
| 313 | let Some(result) = self |
| 314 | .worker |
| 315 | .as_ref() |
| 316 | .and_then(MentionDiscoveryWorker::take_result) |
| 317 | else { |
| 318 | return false; |
| 319 | }; |
| 320 | let is_current = result.generation == self.generation |
| 321 | && self.in_flight.as_ref().is_some_and(|(generation, key)| { |
| 322 | *generation == result.generation && *key == result.key |
| 323 | }); |
| 324 | if !is_current { |
| 325 | return false; |
| 326 | } |
| 327 | self.in_flight = None; |
| 328 | self.cached = Some(CachedMentionDiscovery { |
| 329 | key: result.key, |
| 330 | entries: result.entries, |
| 331 | collected_at: result.collected_at, |
| 332 | }); |
| 333 | true |
| 334 | } |
| 335 | |
| 336 | pub(crate) fn cached_entries(&self, key: &MentionDiscoveryKey) -> Option<&[String]> { |
| 337 | self.cached |
| 338 | .as_ref() |
| 339 | .filter(|cached| &cached.key == key) |
| 340 | .map(|cached| cached.entries.as_slice()) |
| 341 | } |
| 342 | |
| 343 | /// Cached fuzzy-scan candidates for send-time `@`-mention fallback |
| 344 | /// resolution. Returns `Some` only when the cache holds a completed fuzzy |
| 345 | /// scan for exactly this workspace/cwd/depth/symlink configuration — |
| 346 | /// browser-mode caches cover a single directory and are not a usable |
| 347 | /// resolution index. |
| 348 | pub(crate) fn fuzzy_candidates( |
| 349 | &self, |
| 350 | workspace: &std::path::Path, |
| 351 | cwd: &Option<PathBuf>, |
| 352 | walk_depth: usize, |
| 353 | follow_links: bool, |
| 354 | ) -> Option<&[String]> { |
| 355 | let key = MentionDiscoveryKey::fuzzy( |
| 356 | workspace.to_path_buf(), |
| 357 | cwd.clone(), |
| 358 | walk_depth, |
| 359 | follow_links, |
| 360 | ); |
| 361 | self.cached_entries(&key) |
| 362 | } |
| 363 | |
| 364 | /// Cancel the active generation while keeping a same-key cache available |
| 365 | /// for the next mention. |
| 366 | pub(crate) fn cancel(&mut self) { |
| 367 | if self.in_flight.is_none() { |
| 368 | return; |
| 369 | } |
| 370 | let generation = self.next_generation(); |
| 371 | if let Some(worker) = &self.worker { |
| 372 | worker.cancel(generation); |
| 373 | } |
| 374 | self.in_flight = None; |
| 375 | } |
| 376 | |
| 377 | /// Drop cached and in-flight state after workspace/completion settings |
| 378 | /// change. Any late worker result is rejected by the new generation. |
| 379 | pub(crate) fn invalidate(&mut self) { |
| 380 | let generation = self.next_generation(); |
| 381 | if let Some(worker) = &self.worker { |
| 382 | worker.cancel(generation); |
| 383 | } |
| 384 | self.in_flight = None; |
| 385 | self.cached = None; |
| 386 | } |
| 387 | |
| 388 | fn next_generation(&mut self) -> u64 { |
| 389 | self.generation = self.generation.wrapping_add(1).max(1); |
| 390 | self.generation |
| 391 | } |
| 392 | } |
| 393 | |
| 394 | #[cfg(test)] |
| 395 | mod tests { |
| 396 | use super::*; |
| 397 | use std::sync::mpsc; |
| 398 | |
| 399 | const TEST_WORKER_TIMEOUT: Duration = Duration::from_secs(10); |
| 400 | |
| 401 | fn key(name: &str) -> MentionDiscoveryKey { |
| 402 | MentionDiscoveryKey::fuzzy(PathBuf::from(name), None, 10, false) |
| 403 | } |
| 404 | |
| 405 | fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) { |
| 406 | let started = Instant::now(); |
| 407 | while !predicate() { |
| 408 | assert!(started.elapsed() < timeout, "timed out waiting for worker"); |
| 409 | thread::sleep(Duration::from_millis(2)); |
| 410 | } |
| 411 | } |
| 412 | |
| 413 | #[test] |
| 414 | fn request_path_stays_immediate_while_scanner_is_blocked() { |
| 415 | let (started_tx, started_rx) = mpsc::channel(); |
| 416 | let (release_tx, release_rx) = mpsc::channel(); |
| 417 | let release_rx = Mutex::new(release_rx); |
| 418 | let mut discovery = MentionDiscovery::with_scanner(move |_, _| { |
| 419 | let _ = started_tx.send(()); |
| 420 | let _ = release_rx.lock().recv(); |
| 421 | vec!["ready.rs".to_string()] |
| 422 | }); |
| 423 | |
| 424 | let started = Instant::now(); |
| 425 | discovery.ensure_requested(key("slow")); |
| 426 | assert!( |
| 427 | started.elapsed() < Duration::from_millis(50), |
| 428 | "request submission waited on the blocked scanner" |
| 429 | ); |
| 430 | started_rx |
| 431 | .recv_timeout(TEST_WORKER_TIMEOUT) |
| 432 | .expect("scanner should start in the background"); |
| 433 | |
| 434 | let second_started = Instant::now(); |
| 435 | discovery.ensure_requested(key("newer")); |
| 436 | assert!( |
| 437 | second_started.elapsed() < Duration::from_millis(50), |
| 438 | "superseding request waited on the blocked scanner" |
| 439 | ); |
| 440 | release_tx.send(()).unwrap(); |
| 441 | } |
| 442 | |
| 443 | #[test] |
| 444 | fn contended_pending_slot_does_not_drop_request() { |
| 445 | let (scan_started_tx, scan_started_rx) = mpsc::channel(); |
| 446 | let worker = MentionDiscoveryWorker::spawn(Arc::new(move |_, _| { |
| 447 | let _ = scan_started_tx.send(()); |
| 448 | Vec::new() |
| 449 | })) |
| 450 | .expect("worker should start"); |
| 451 | let pending_guard = worker.shared.pending.lock(); |
| 452 | let (attempting_tx, attempting_rx) = mpsc::channel(); |
| 453 | let (submitted_tx, submitted_rx) = mpsc::channel(); |
| 454 | let request = MentionDiscoveryRequest { |
| 455 | generation: 1, |
| 456 | key: key("contended"), |
| 457 | }; |
| 458 | |
| 459 | thread::scope(|scope| { |
| 460 | let worker = &worker; |
| 461 | scope.spawn(move || { |
| 462 | attempting_tx.send(()).unwrap(); |
| 463 | submitted_tx.send(worker.submit(request)).unwrap(); |
| 464 | }); |
| 465 | attempting_rx |
| 466 | .recv_timeout(TEST_WORKER_TIMEOUT) |
| 467 | .expect("submission should be attempted"); |
| 468 | let submitted_while_contended = submitted_rx.recv_timeout(Duration::from_millis(50)); |
| 469 | drop(pending_guard); |
| 470 | let submitted = match submitted_while_contended { |
| 471 | Ok(submitted) => submitted, |
| 472 | Err(mpsc::RecvTimeoutError::Timeout) => submitted_rx |
| 473 | .recv_timeout(TEST_WORKER_TIMEOUT) |
| 474 | .expect("submission should finish after contention clears"), |
| 475 | Err(mpsc::RecvTimeoutError::Disconnected) => { |
| 476 | panic!("submission thread disconnected") |
| 477 | } |
| 478 | }; |
| 479 | assert!(submitted, "a contended request must not be dropped"); |
| 480 | }); |
| 481 | |
| 482 | scan_started_rx |
| 483 | .recv_timeout(TEST_WORKER_TIMEOUT) |
| 484 | .expect("the contended request should reach the scanner"); |
| 485 | } |
| 486 | |
| 487 | #[test] |
| 488 | fn late_result_cannot_replace_new_generation() { |
| 489 | let (first_started_tx, first_started_rx) = mpsc::channel(); |
| 490 | let (release_first_tx, release_first_rx) = mpsc::channel(); |
| 491 | let release_first_rx = Mutex::new(release_first_rx); |
| 492 | let mut discovery = MentionDiscovery::with_scanner(move |key, cancelled| { |
| 493 | if key.workspace == std::path::Path::new("old") { |
| 494 | let _ = first_started_tx.send(()); |
| 495 | let _ = release_first_rx.lock().recv(); |
| 496 | // Simulate an uninterruptible filesystem call returning late. |
| 497 | assert!(cancelled()); |
| 498 | vec!["stale.rs".to_string()] |
| 499 | } else { |
| 500 | vec!["current.rs".to_string()] |
| 501 | } |
| 502 | }); |
| 503 | |
| 504 | let old_key = key("old"); |
| 505 | let new_key = key("new"); |
| 506 | discovery.ensure_requested(old_key.clone()); |
| 507 | first_started_rx |
| 508 | .recv_timeout(TEST_WORKER_TIMEOUT) |
| 509 | .expect("old scan should start"); |
| 510 | discovery.ensure_requested(new_key.clone()); |
| 511 | release_first_tx.send(()).unwrap(); |
| 512 | |
| 513 | wait_until(TEST_WORKER_TIMEOUT, || discovery.poll()); |
| 514 | assert!(discovery.cached_entries(&old_key).is_none()); |
| 515 | assert_eq!( |
| 516 | discovery.cached_entries(&new_key), |
| 517 | Some(["current.rs".to_string()].as_slice()) |
| 518 | ); |
| 519 | } |
| 520 | } |
| 521 |