1use crate::error::{Result, SearchError};
2use crate::parse::TranslationEntry;
3use hashbrown::HashMap;
4use serde::{Deserialize, Serialize};
5use sled::Db;
6use std::fs;
7use std::io::{Read, Write};
8use std::net::{Shutdown, TcpListener, TcpStream};
9use std::path::{Path, PathBuf};
10use std::process::{Command, Stdio};
11use std::sync::Mutex;
12use std::time::{Duration, SystemTime};
13
14const CACHE_DIR_NAME: &str = "cs";
15const PORT_FILE: &str = "cache.port";
16const SERVER_FLAG: &str = "--cache-server";
17const FRONT_CACHE_CAP: usize = 512;
18const MAX_CACHE_SIZE: u64 = 1_000_000_000;
19const MAX_CACHE_AGE_SECS: u64 = 30 * 24 * 60 * 60;
20const CLEANUP_INTERVAL_SECS: u64 = 6 * 60 * 60;
21
22#[derive(Serialize, Deserialize, Clone)]
24struct CacheValue {
25 mtime_secs: u64,
26 file_size: u64,
27 last_accessed: u64,
28 results: Vec<TranslationEntry>,
29}
30
31pub struct SearchResultCache {
33 backend: CacheBackend,
34}
35
36enum CacheBackend {
37 Local(LocalCache),
38 Remote(RemoteCache),
39}
40
41#[derive(Serialize, Deserialize, Debug)]
42enum CacheRequest {
43 Get {
44 file: PathBuf,
45 query: String,
46 case_sensitive: bool,
47 mtime_secs: u64,
48 file_size: u64,
49 },
50 Set {
51 file: PathBuf,
52 query: String,
53 case_sensitive: bool,
54 mtime_secs: u64,
55 file_size: u64,
56 results: Vec<TranslationEntry>,
57 },
58 Clear,
59 Ping,
60}
61
62#[derive(Serialize, Deserialize, Debug)]
63enum CacheResponse {
64 Get(Option<Vec<TranslationEntry>>),
65 Ack(bool),
66}
67
68impl SearchResultCache {
69 pub fn new() -> Result<Self> {
71 if std::env::var("CS_DISABLE_CACHE_SERVER").is_ok() {
72 return Ok(Self {
73 backend: CacheBackend::Local(LocalCache::new()?),
74 });
75 }
76
77 if let Some(remote) = RemoteCache::connect_or_spawn()? {
78 return Ok(Self {
79 backend: CacheBackend::Remote(remote),
80 });
81 }
82
83 Ok(Self {
84 backend: CacheBackend::Local(LocalCache::new()?),
85 })
86 }
87
88 pub fn with_cache_dir(cache_dir: PathBuf) -> Result<Self> {
90 Ok(Self {
91 backend: CacheBackend::Local(LocalCache::with_cache_dir(cache_dir)?),
92 })
93 }
94
95 pub fn get(
96 &self,
97 file: &Path,
98 query: &str,
99 case_sensitive: bool,
100 current_mtime: SystemTime,
101 current_size: u64,
102 ) -> Option<Vec<TranslationEntry>> {
103 match &self.backend {
104 CacheBackend::Local(inner) => {
105 inner.get(file, query, case_sensitive, current_mtime, current_size)
106 }
107 CacheBackend::Remote(remote) => remote
108 .get(file, query, case_sensitive, current_mtime, current_size)
109 .ok()
110 .flatten(),
111 }
112 }
113
114 pub fn set(
115 &self,
116 file: &Path,
117 query: &str,
118 case_sensitive: bool,
119 mtime: SystemTime,
120 file_size: u64,
121 results: &[TranslationEntry],
122 ) -> Result<()> {
123 match &self.backend {
124 CacheBackend::Local(inner) => {
125 inner.set(file, query, case_sensitive, mtime, file_size, results)
126 }
127 CacheBackend::Remote(remote) => {
128 remote.set(file, query, case_sensitive, mtime, file_size, results)
129 }
130 }
131 }
132
133 pub fn clear(&self) -> Result<()> {
134 match &self.backend {
135 CacheBackend::Local(inner) => inner.clear(),
136 CacheBackend::Remote(remote) => remote.clear(),
137 }
138 }
139
140 pub fn start_server_blocking() -> Result<()> {
142 run_cache_server()
143 }
144}
145
146struct LocalCache {
147 db: Db,
148 last_cleanup: SystemTime,
149 front_cache: Mutex<HashMap<Vec<u8>, CacheValue>>,
150 cache_dir: PathBuf,
151}
152
153impl LocalCache {
154 fn cache_dir() -> PathBuf {
155 dirs::cache_dir()
156 .unwrap_or_else(|| PathBuf::from("."))
157 .join(CACHE_DIR_NAME)
158 }
159
160 fn with_cache_dir(cache_dir: PathBuf) -> Result<Self> {
161 fs::create_dir_all(&cache_dir)?;
162 let db = sled::open(cache_dir.join("db"))
163 .map_err(|e| SearchError::Generic(format!("Failed to open cache: {}", e)))?;
164
165 let last_cleanup = Self::read_last_cleanup_marker(&cache_dir)?;
166 let cache = Self {
167 db,
168 last_cleanup,
169 front_cache: Mutex::new(HashMap::new()),
170 cache_dir,
171 };
172 cache.maybe_cleanup_on_open()?;
173 Ok(cache)
174 }
175
176 fn new() -> Result<Self> {
177 Self::with_cache_dir(Self::cache_dir())
178 }
179
180 fn get(
181 &self,
182 file: &Path,
183 query: &str,
184 case_sensitive: bool,
185 current_mtime: SystemTime,
186 current_size: u64,
187 ) -> Option<Vec<TranslationEntry>> {
188 let key = self.make_key(file, query, case_sensitive);
189
190 if let Some(entries) = self.front_get(&key, current_mtime, current_size) {
191 return Some(entries);
192 }
193
194 let cached_bytes = self.db.get(&key).ok()??;
195 let mut cached: CacheValue = bincode::deserialize(&cached_bytes).ok()?;
196
197 let current_secs = current_mtime
198 .duration_since(SystemTime::UNIX_EPOCH)
199 .ok()?
200 .as_secs();
201
202 let now = SystemTime::now()
203 .duration_since(SystemTime::UNIX_EPOCH)
204 .ok()?
205 .as_secs();
206
207 if now.saturating_sub(cached.last_accessed) > MAX_CACHE_AGE_SECS {
209 let _ = self.db.remove(&key);
211 return None;
212 }
213
214 if cached.mtime_secs == current_secs && cached.file_size == current_size {
215 cached.last_accessed = now;
216
217 if let Ok(updated_bytes) = bincode::serialize(&cached) {
218 let _ = self.db.insert(&key, updated_bytes);
219 }
220
221 self.front_set(key.clone(), cached.clone());
222 Some(cached.results)
223 } else {
224 let _ = self.db.remove(&key);
226 None
227 }
228 }
229
230 fn set(
231 &self,
232 file: &Path,
233 query: &str,
234 case_sensitive: bool,
235 mtime: SystemTime,
236 file_size: u64,
237 results: &[TranslationEntry],
238 ) -> Result<()> {
239 let key = self.make_key(file, query, case_sensitive);
240
241 let mtime_secs = mtime
242 .duration_since(SystemTime::UNIX_EPOCH)
243 .map_err(|e| SearchError::Generic(format!("Invalid mtime: {}", e)))?
244 .as_secs();
245
246 let last_accessed = SystemTime::now()
247 .duration_since(SystemTime::UNIX_EPOCH)
248 .map_err(|e| SearchError::Generic(format!("Failed to get current time: {}", e)))?
249 .as_secs();
250
251 let value = CacheValue {
252 mtime_secs,
253 file_size,
254 last_accessed,
255 results: results.to_vec(),
256 };
257
258 let value_bytes = bincode::serialize(&value)
259 .map_err(|e| SearchError::Generic(format!("Failed to serialize cache: {}", e)))?;
260
261 self.front_set(key.clone(), value.clone());
262
263 self.db
264 .insert(key, value_bytes)
265 .map_err(|e| SearchError::Generic(format!("Failed to write cache: {}", e)))?;
266
267 Ok(())
268 }
269
270 fn clear(&self) -> Result<()> {
271 self.db
272 .clear()
273 .map_err(|e| SearchError::Generic(format!("Failed to clear cache: {}", e)))?;
274 if let Ok(mut map) = self.front_cache.lock() {
275 map.clear();
276 }
277 let _ = fs::remove_file(Self::meta_file_path(&self.cache_dir));
278 Ok(())
279 }
280
281 fn front_get(
282 &self,
283 key: &[u8],
284 current_mtime: SystemTime,
285 current_size: u64,
286 ) -> Option<Vec<TranslationEntry>> {
287 let guard = self.front_cache.lock().ok()?;
288 let entry = guard.get(key)?;
289 let current_secs = current_mtime
290 .duration_since(SystemTime::UNIX_EPOCH)
291 .ok()?
292 .as_secs();
293 if entry.mtime_secs == current_secs && entry.file_size == current_size {
294 Some(entry.results.clone())
295 } else {
296 None
297 }
298 }
299
300 fn front_set(&self, key: Vec<u8>, value: CacheValue) {
301 if let Ok(mut map) = self.front_cache.lock() {
302 if map.len() >= FRONT_CACHE_CAP {
303 if let Some(oldest_key) = map
304 .iter()
305 .min_by_key(|(_, v)| v.last_accessed)
306 .map(|(k, _)| k.clone())
307 {
308 map.remove(&oldest_key);
309 }
310 }
311 map.insert(key, value);
312 }
313 }
314
315 fn make_key(&self, file: &Path, query: &str, case_sensitive: bool) -> Vec<u8> {
316 let normalized_query = if case_sensitive {
317 query.to_string()
318 } else {
319 query.to_lowercase()
320 };
321 format!("{}|{}", file.display(), normalized_query).into_bytes()
322 }
323
324 fn maybe_cleanup_on_open(&self) -> Result<()> {
325 let now = SystemTime::now()
326 .duration_since(SystemTime::UNIX_EPOCH)
327 .map_err(|e| SearchError::Generic(format!("Failed to get current time: {}", e)))?
328 .as_secs();
329
330 let last = self
331 .last_cleanup
332 .duration_since(SystemTime::UNIX_EPOCH)
333 .unwrap_or_default()
334 .as_secs();
335
336 if now.saturating_sub(last) >= CLEANUP_INTERVAL_SECS {
337 self.cleanup_if_needed()?;
338 }
339
340 Ok(())
341 }
342
343 fn cleanup_if_needed(&self) -> Result<()> {
344 let size = self
346 .db
347 .size_on_disk()
348 .map_err(|e| SearchError::Generic(format!("Failed to get cache size: {}", e)))?;
349
350 if size <= MAX_CACHE_SIZE {
352 return Ok(());
353 }
354
355 let now = SystemTime::now()
357 .duration_since(SystemTime::UNIX_EPOCH)
358 .map_err(|e| SearchError::Generic(format!("Failed to get current time: {}", e)))?
359 .as_secs();
360
361 let mut entries: Vec<(Vec<u8>, u64)> = Vec::new();
362
363 for (key, value) in self.db.iter().flatten() {
364 if let Ok(cache_value) = bincode::deserialize::<CacheValue>(&value) {
365 if now.saturating_sub(cache_value.last_accessed) <= MAX_CACHE_AGE_SECS {
367 entries.push((key.to_vec(), cache_value.last_accessed));
368 }
369 }
370 }
371
372 entries.sort_by_key(|(_, last_accessed)| *last_accessed);
374
375 for (key, _) in entries.iter() {
377 if self
378 .db
379 .size_on_disk()
380 .ok()
381 .map(|s| s <= MAX_CACHE_SIZE)
382 .unwrap_or(true)
383 {
384 break;
385 }
386 let _ = self.db.remove(key);
387 }
388
389 let _ = self.db.flush();
390 self.write_last_cleanup_marker(&self.cache_dir);
391 Ok(())
392 }
393
394 fn meta_file_path(cache_dir: &Path) -> PathBuf {
395 cache_dir.join("meta.last")
396 }
397
398 fn write_last_cleanup_marker(&self, cache_dir: &Path) {
399 let _ = fs::write(
400 Self::meta_file_path(cache_dir),
401 SystemTime::now()
402 .duration_since(SystemTime::UNIX_EPOCH)
403 .map(|d| d.as_secs().to_string())
404 .unwrap_or_else(|_| "0".to_string()),
405 );
406 }
407
408 fn read_last_cleanup_marker(cache_dir: &Path) -> Result<SystemTime> {
409 let path = Self::meta_file_path(cache_dir);
410
411 let contents = fs::read_to_string(path).ok();
412 if let Some(s) = contents {
413 if let Ok(secs) = s.trim().parse::<u64>() {
414 return Ok(SystemTime::UNIX_EPOCH + Duration::from_secs(secs));
415 }
416 }
417
418 Ok(SystemTime::UNIX_EPOCH)
419 }
420}
421
422struct RemoteCache {
423 addr: String,
424}
425
426impl RemoteCache {
427 fn connect_or_spawn() -> Result<Option<Self>> {
428 if let Some(addr) = read_port_file() {
429 if Self::ping_addr(&addr).is_ok() {
430 return Ok(Some(Self { addr }));
431 }
432 }
433
434 spawn_server()?;
435
436 if let Some(addr) = read_port_file() {
437 if Self::ping_addr(&addr).is_ok() {
438 return Ok(Some(Self { addr }));
439 }
440 }
441
442 Ok(None)
443 }
444
445 fn ping_addr(addr: &str) -> Result<()> {
446 let client = Self {
447 addr: addr.to_string(),
448 };
449 match client.send_request(CacheRequest::Ping)? {
450 CacheResponse::Ack(true) => Ok(()),
451 _ => Err(SearchError::Generic(
452 "Cache server did not acknowledge ping".to_string(),
453 )),
454 }
455 }
456
457 fn get(
458 &self,
459 file: &Path,
460 query: &str,
461 case_sensitive: bool,
462 current_mtime: SystemTime,
463 current_size: u64,
464 ) -> Result<Option<Vec<TranslationEntry>>> {
465 let mtime_secs = current_mtime
466 .duration_since(SystemTime::UNIX_EPOCH)
467 .map_err(|e| SearchError::Generic(format!("Invalid mtime: {}", e)))?
468 .as_secs();
469
470 let req = CacheRequest::Get {
471 file: file.to_path_buf(),
472 query: query.to_string(),
473 case_sensitive,
474 mtime_secs,
475 file_size: current_size,
476 };
477
478 match self.send_request(req)? {
479 CacheResponse::Get(res) => Ok(res),
480 _ => Err(SearchError::Generic("Invalid cache response".to_string())),
481 }
482 }
483
484 fn set(
485 &self,
486 file: &Path,
487 query: &str,
488 case_sensitive: bool,
489 mtime: SystemTime,
490 file_size: u64,
491 results: &[TranslationEntry],
492 ) -> Result<()> {
493 let mtime_secs = mtime
494 .duration_since(SystemTime::UNIX_EPOCH)
495 .map_err(|e| SearchError::Generic(format!("Invalid mtime: {}", e)))?
496 .as_secs();
497
498 let req = CacheRequest::Set {
499 file: file.to_path_buf(),
500 query: query.to_string(),
501 case_sensitive,
502 mtime_secs,
503 file_size,
504 results: results.to_vec(),
505 };
506
507 match self.send_request(req)? {
508 CacheResponse::Ack(true) => Ok(()),
509 _ => Err(SearchError::Generic("Cache write failed".to_string())),
510 }
511 }
512
513 fn clear(&self) -> Result<()> {
514 match self.send_request(CacheRequest::Clear)? {
515 CacheResponse::Ack(true) => Ok(()),
516 _ => Err(SearchError::Generic("Failed to clear cache".to_string())),
517 }
518 }
519
520 fn send_request(&self, req: CacheRequest) -> Result<CacheResponse> {
521 let mut stream = TcpStream::connect(&self.addr)
522 .map_err(|e| SearchError::Generic(format!("Failed to connect cache server: {}", e)))?;
523
524 let bytes = bincode::serialize(&req)
525 .map_err(|e| SearchError::Generic(format!("Failed to encode cache request: {}", e)))?;
526
527 stream
528 .write_all(&bytes)
529 .map_err(|e| SearchError::Generic(format!("Failed to write cache request: {}", e)))?;
530 let _ = stream.shutdown(Shutdown::Write);
531
532 let mut buf = Vec::new();
533 stream
534 .read_to_end(&mut buf)
535 .map_err(|e| SearchError::Generic(format!("Failed to read cache response: {}", e)))?;
536
537 let resp: CacheResponse = bincode::deserialize(&buf)
538 .map_err(|e| SearchError::Generic(format!("Failed to decode cache response: {}", e)))?;
539 Ok(resp)
540 }
541}
542
543fn run_cache_server() -> Result<()> {
545 let cache_dir = LocalCache::cache_dir();
546 fs::create_dir_all(&cache_dir)?;
547
548 let listener = TcpListener::bind("127.0.0.1:0")
549 .map_err(|e| SearchError::Generic(format!("Failed to bind cache server: {}", e)))?;
550 let addr = listener
551 .local_addr()
552 .map_err(|e| SearchError::Generic(format!("Failed to get cache server address: {}", e)))?;
553 write_port_file(&cache_dir, &addr.to_string())?;
554
555 let local = LocalCache::with_cache_dir(cache_dir)?;
556 for stream in listener.incoming() {
557 match stream {
558 Ok(mut stream) => {
559 let _ = handle_connection(&local, &mut stream);
560 }
561 Err(_) => continue,
562 }
563 }
564 Ok(())
565}
566
567fn handle_connection(local: &LocalCache, stream: &mut TcpStream) -> Result<()> {
568 let mut buf = Vec::new();
569 stream.read_to_end(&mut buf)?;
570 let req: CacheRequest = bincode::deserialize(&buf)
571 .map_err(|e| SearchError::Generic(format!("Failed to decode cache request: {}", e)))?;
572
573 let resp = match req {
574 CacheRequest::Get {
575 file,
576 query,
577 case_sensitive,
578 mtime_secs,
579 file_size,
580 } => {
581 let ts = SystemTime::UNIX_EPOCH + Duration::from_secs(mtime_secs);
582 let hit = local.get(&file, &query, case_sensitive, ts, file_size);
583 CacheResponse::Get(hit)
584 }
585 CacheRequest::Set {
586 file,
587 query,
588 case_sensitive,
589 mtime_secs,
590 file_size,
591 results,
592 } => {
593 let ts = SystemTime::UNIX_EPOCH + Duration::from_secs(mtime_secs);
594 let res = local.set(&file, &query, case_sensitive, ts, file_size, &results);
595 CacheResponse::Ack(res.is_ok())
596 }
597 CacheRequest::Clear => {
598 let res = local.clear();
599 CacheResponse::Ack(res.is_ok())
600 }
601 CacheRequest::Ping => CacheResponse::Ack(true),
602 };
603
604 let resp_bytes = bincode::serialize(&resp)
605 .map_err(|e| SearchError::Generic(format!("Failed to encode cache response: {}", e)))?;
606 stream.write_all(&resp_bytes)?;
607 let _ = stream.shutdown(Shutdown::Write);
608 Ok(())
609}
610
611fn cache_port_path(cache_dir: &Path) -> PathBuf {
613 cache_dir.join(PORT_FILE)
614}
615
616fn write_port_file(cache_dir: &Path, addr: &str) -> Result<()> {
617 fs::write(cache_port_path(cache_dir), addr)
618 .map_err(|e| SearchError::Generic(format!("Failed to write cache port: {}", e)))
619}
620
621fn read_port_file() -> Option<String> {
622 let path = cache_port_path(&LocalCache::cache_dir());
623 fs::read_to_string(path).ok().map(|s| s.trim().to_string())
624}
625
626fn spawn_server() -> Result<()> {
627 let exe = resolve_server_binary()?;
628
629 Command::new(exe)
630 .arg(SERVER_FLAG)
631 .stdout(Stdio::null())
632 .stderr(Stdio::null())
633 .spawn()
634 .map_err(|e| SearchError::Generic(format!("Failed to spawn cache server: {}", e)))?;
635 std::thread::sleep(Duration::from_millis(150));
636 Ok(())
637}
638
639fn resolve_server_binary() -> Result<PathBuf> {
640 let exe = std::env::current_exe()
641 .map_err(|e| SearchError::Generic(format!("Failed to get current exe: {}", e)))?;
642
643 let bin_name = if cfg!(target_os = "windows") {
644 "cs.exe"
645 } else {
646 "cs"
647 };
648
649 let mut candidates = Vec::new();
652 if let Some(dir) = exe.parent() {
653 candidates.push(dir.join(bin_name));
654 if let Some(parent) = dir.parent() {
655 candidates.push(parent.join(bin_name));
656 }
657 }
658 candidates.push(exe.clone());
659
660 for path in candidates {
661 if path.exists() {
662 return Ok(path);
663 }
664 }
665
666 Err(SearchError::Generic(
667 "Could not locate cache server binary".to_string(),
668 ))
669}
670
671#[cfg(test)]
672mod tests {
673 use super::*;
674 use std::fs;
675 use tempfile::{NamedTempFile, TempDir};
676
677 #[test]
678 fn test_cache_hit_local() {
679 let cache_dir = TempDir::new().unwrap();
680 let cache = SearchResultCache::with_cache_dir(cache_dir.path().to_path_buf()).unwrap();
681 let file = NamedTempFile::new().unwrap();
682 fs::write(&file, "test content").unwrap();
683
684 let metadata = fs::metadata(file.path()).unwrap();
685 let mtime = metadata.modified().unwrap();
686 let size = metadata.len();
687
688 let results = vec![TranslationEntry {
689 key: "test.key".to_string(),
690 value: "test value".to_string(),
691 file: file.path().to_path_buf(),
692 line: 1,
693 }];
694
695 cache
696 .set(file.path(), "query", false, mtime, size, &results)
697 .unwrap();
698 let cached = cache.get(file.path(), "query", false, mtime, size);
699 assert!(cached.is_some());
700 assert_eq!(cached.unwrap().len(), 1);
701 }
702
703 #[test]
704 fn test_cache_invalidation_on_file_change_local() {
705 let cache_dir = TempDir::new().unwrap();
706 let cache = SearchResultCache::with_cache_dir(cache_dir.path().to_path_buf()).unwrap();
707 let file = NamedTempFile::new().unwrap();
708 fs::write(&file, "original content").unwrap();
709
710 let metadata = fs::metadata(file.path()).unwrap();
711 let mtime = metadata.modified().unwrap();
712 let size = metadata.len();
713
714 let results = vec![TranslationEntry {
715 key: "test.key".to_string(),
716 value: "test value".to_string(),
717 file: file.path().to_path_buf(),
718 line: 1,
719 }];
720
721 cache
722 .set(file.path(), "query", false, mtime, size, &results)
723 .unwrap();
724
725 std::thread::sleep(std::time::Duration::from_secs(1));
726 fs::write(&file, "modified content with different size").unwrap();
727
728 let new_metadata = fs::metadata(file.path()).unwrap();
729 let new_mtime = new_metadata.modified().unwrap();
730 let new_size = new_metadata.len();
731
732 assert!(new_size != size || new_mtime != mtime);
733
734 let cached = cache.get(file.path(), "query", false, new_mtime, new_size);
735 assert!(cached.is_none());
736 }
737
738 #[test]
739 fn test_case_insensitive_normalization_local() {
740 let cache_dir = TempDir::new().unwrap();
741 let cache = SearchResultCache::with_cache_dir(cache_dir.path().to_path_buf()).unwrap();
742 let file = NamedTempFile::new().unwrap();
743 fs::write(&file, "test content").unwrap();
744
745 let metadata = fs::metadata(file.path()).unwrap();
746 let mtime = metadata.modified().unwrap();
747 let size = metadata.len();
748
749 let results = vec![TranslationEntry {
750 key: "test.key".to_string(),
751 value: "test value".to_string(),
752 file: file.path().to_path_buf(),
753 line: 1,
754 }];
755
756 cache
757 .set(file.path(), "query", false, mtime, size, &results)
758 .unwrap();
759
760 let cached = cache.get(file.path(), "QUERY", false, mtime, size);
761 assert!(cached.is_some());
762 }
763}