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