git_worktree_manager/operations/
pr_cache.rs1use std::collections::HashMap;
10use std::path::{Path, PathBuf};
11use std::sync::atomic::{AtomicU64, Ordering};
12#[cfg(not(test))]
13use std::sync::OnceLock;
14use std::time::{SystemTime, UNIX_EPOCH};
15
16static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
19
20#[cfg(not(test))]
25static SWEEP_DONE: OnceLock<()> = OnceLock::new();
26
27use serde::{Deserialize, Serialize};
28use sha2::{Digest, Sha256};
29
30const CACHE_TTL_SECS: u64 = 60;
32
33const GH_FETCH_LIMIT: usize = 500;
39
40#[non_exhaustive]
45#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
46#[serde(rename_all = "UPPERCASE")]
47pub enum PrState {
48 Open,
49 Merged,
50 Closed,
51 #[serde(other)]
52 Other,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
56struct CacheFile {
57 fetched_at: u64,
58 repo: String,
59 prs: HashMap<String, PrState>,
60}
61
62#[derive(Debug, Default, Clone)]
63pub struct PrCache {
64 map: HashMap<String, PrState>,
65}
66
67impl PrCache {
68 pub fn state(&self, branch: &str) -> Option<&PrState> {
76 self.map.get(branch)
77 }
78
79 pub fn from_disk(repo: &Path) -> Option<Self> {
82 load_from_disk(repo).map(|map| PrCache { map })
83 }
84
85 pub fn fetch_and_persist(repo: &Path) -> Self {
91 match fetch_from_gh(repo) {
92 Some(map) => {
93 write_to_disk(repo, &map);
94 PrCache { map }
95 }
96 None => PrCache::default(),
97 }
98 }
99
100 pub fn load_or_fetch(repo: &Path, no_cache: bool) -> Self {
108 if !no_cache {
109 if let Some(c) = Self::from_disk(repo) {
110 return c;
111 }
112 }
113 Self::fetch_and_persist(repo)
114 }
115}
116
117fn repo_hash(repo: &Path) -> String {
125 let canon = repo.canonicalize().unwrap_or_else(|_| repo.to_path_buf());
126 let mut hasher = Sha256::new();
127 hasher.update(canon.to_string_lossy().as_bytes());
128 let digest = hasher.finalize();
129 hex_short(&digest[..8])
130}
131
132fn hex_short(bytes: &[u8]) -> String {
133 use std::fmt::Write;
134 let mut out = String::with_capacity(bytes.len() * 2);
135 for b in bytes {
136 let _ = write!(out, "{:02x}", b);
137 }
138 out
139}
140
141fn cache_path_for(repo: &Path) -> Option<PathBuf> {
144 #[cfg(test)]
145 if let Ok(dir) = std::env::var("GW_TEST_CACHE_DIR") {
146 return Some(
147 PathBuf::from(dir)
148 .join("gw")
149 .join(format!("pr-status-{}.json", repo_hash(repo))),
150 );
151 }
152
153 let base = dirs::cache_dir()?.join("gw");
154 Some(base.join(format!("pr-status-{}.json", repo_hash(repo))))
155}
156
157#[derive(Debug, Deserialize)]
158struct GhPr {
159 #[serde(rename = "headRefName")]
160 head_ref_name: String,
161 state: PrState,
162}
163
164fn fetch_from_gh(repo: &Path) -> Option<HashMap<String, PrState>> {
169 #[cfg(test)]
170 {
171 if std::env::var("GW_TEST_GH_FAIL").ok().as_deref() == Some("1") {
174 return None;
175 }
176 if let Ok(json) = std::env::var("GW_TEST_GH_JSON") {
177 let prs: Vec<GhPr> = serde_json::from_str(json.trim()).ok()?;
178 let mut map = HashMap::with_capacity(prs.len());
179 for pr in prs {
180 map.insert(pr.head_ref_name, pr.state);
181 }
182 return Some(map);
183 }
184 }
185
186 if !crate::git::has_command("gh") {
187 return None;
188 }
189 let limit = GH_FETCH_LIMIT.to_string();
190 let result = crate::git::run_command(
191 &[
192 "gh",
193 "pr",
194 "list",
195 "--state",
196 "all",
197 "--json",
198 "headRefName,state",
199 "--limit",
200 &limit,
201 ],
202 Some(repo),
203 false,
204 true,
205 )
206 .ok()?;
207 if result.returncode != 0 {
208 return None;
209 }
210
211 let prs: Vec<GhPr> = serde_json::from_str(result.stdout.trim()).ok()?;
212 let mut map = HashMap::with_capacity(prs.len());
213 for pr in prs {
214 map.insert(pr.head_ref_name, pr.state);
215 }
216 Some(map)
217}
218
219fn now_secs() -> Option<u64> {
223 SystemTime::now()
224 .duration_since(UNIX_EPOCH)
225 .ok()
226 .map(|d| d.as_secs())
227}
228
229fn load_from_disk(repo: &Path) -> Option<HashMap<String, PrState>> {
232 let path = cache_path_for(repo)?;
233 let data = std::fs::read_to_string(&path).ok()?;
234 let file: CacheFile = serde_json::from_str(&data).ok()?;
235 let now = now_secs()?;
237 if file.fetched_at > now {
239 return None;
240 }
241 let age = now.saturating_sub(file.fetched_at);
242 if age > CACHE_TTL_SECS {
243 return None;
244 }
245 Some(file.prs)
246}
247
248fn sweep_orphans(parent: &Path, cutoff: SystemTime) {
253 let Ok(entries) = std::fs::read_dir(parent) else {
254 return;
255 };
256 for entry in entries.flatten() {
257 let name = entry.file_name();
258 let name_str = name.to_string_lossy();
259 if !(name_str.starts_with("pr-status-") && name_str.contains(".tmp.")) {
260 continue;
261 }
262 let Ok(meta) = entry.metadata() else {
263 continue;
264 };
265 let Ok(modified) = meta.modified() else {
266 continue;
267 };
268 if modified < cutoff {
269 let _ = std::fs::remove_file(entry.path());
270 }
271 }
272}
273
274fn write_to_disk(repo: &Path, prs: &HashMap<String, PrState>) {
284 let Some(path) = cache_path_for(repo) else {
285 return;
286 };
287 if let Some(parent) = path.parent() {
288 let _ = std::fs::create_dir_all(parent);
289 }
290 let now_instant = SystemTime::now();
296 let dur = match now_instant.duration_since(UNIX_EPOCH) {
297 Ok(d) => d,
298 Err(_) => return, };
300 let now = dur.as_secs();
301 let nanos = dur.subsec_nanos();
302
303 let file = CacheFile {
304 fetched_at: now,
305 repo: repo.to_string_lossy().into_owned(),
306 prs: prs.clone(),
307 };
308 let Ok(json) = serde_json::to_string(&file) else {
309 return;
310 };
311
312 #[cfg(not(test))]
318 let do_sweep = SWEEP_DONE.set(()).is_ok();
319 #[cfg(test)]
320 let do_sweep = true;
321 if do_sweep {
322 if let Some(parent) = path.parent() {
323 let cutoff = SystemTime::now()
326 .checked_sub(std::time::Duration::from_secs(60))
327 .unwrap_or_else(SystemTime::now);
328 sweep_orphans(parent, cutoff);
329 }
330 }
331
332 let counter = TMP_COUNTER.fetch_add(1, Ordering::Relaxed);
342 let tmp = path.with_file_name(format!(
343 "{}.tmp.{}.{}.{}",
344 path.file_stem().unwrap_or_default().to_string_lossy(),
345 std::process::id(),
346 nanos,
347 counter,
348 ));
349 if std::fs::write(&tmp, &json).is_err() {
353 let _ = std::fs::remove_file(&tmp); return;
355 }
356 if std::fs::rename(&tmp, &path).is_err() {
357 let _ = std::fs::remove_file(&path);
362 if std::fs::rename(&tmp, &path).is_err() {
363 let _ = std::fs::remove_file(&tmp); }
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use super::*;
371 use std::path::PathBuf;
372
373 use super::super::test_env::{env_lock, EnvGuard};
381
382 #[test]
386 fn now_secs_returns_some_on_normal_system() {
387 assert!(now_secs().is_some());
388 }
389
390 #[test]
391 fn repo_hash_is_stable_and_short() {
392 let p = PathBuf::from("/tmp/some-repo-that-does-not-exist-xyz");
393 let h1 = repo_hash(&p);
394 let h2 = repo_hash(&p);
395 assert_eq!(h1, h2);
396 assert_eq!(h1.len(), 16);
397 }
398
399 #[test]
400 fn repo_hash_differs_per_path() {
401 let a = repo_hash(&PathBuf::from("/tmp/repo-a-xyz"));
402 let b = repo_hash(&PathBuf::from("/tmp/repo-b-xyz"));
403 assert_ne!(a, b);
404 }
405
406 #[test]
407 fn cache_path_contains_repo_hash() {
408 let p = PathBuf::from("/tmp/repo-xyz");
409 let cp = cache_path_for(&p).expect("cache dir available");
410 let s = cp.to_string_lossy();
411 assert!(s.contains("gw"));
412 assert!(s.contains("pr-status-"));
413 assert!(s.ends_with(".json"));
414 }
415
416 #[test]
417 fn fetch_parses_gh_json_from_env() {
418 let _g = env_lock();
419 let _env = EnvGuard::capture(&["GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
420 std::env::set_var(
421 "GW_TEST_GH_JSON",
422 r#"[{"headRefName":"feat/foo","state":"OPEN"},{"headRefName":"fix/bar","state":"MERGED"}]"#,
423 );
424 let prs = fetch_from_gh(std::path::Path::new(".")).expect("parsed");
425 assert_eq!(prs.get("feat/foo"), Some(&PrState::Open));
426 assert_eq!(prs.get("fix/bar"), Some(&PrState::Merged));
427 }
428
429 #[test]
430 fn fetch_returns_none_on_forced_failure() {
431 let _g = env_lock();
432 let _env = EnvGuard::capture(&["GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
433 std::env::set_var("GW_TEST_GH_FAIL", "1");
434 let result = fetch_from_gh(std::path::Path::new("."));
435 assert!(result.is_none());
436 }
437
438 use tempfile::tempdir;
439
440 fn with_cache_dir<F: FnOnce()>(dir: &std::path::Path, f: F) {
444 let _g = EnvGuard::capture(&["GW_TEST_CACHE_DIR"]);
445 std::env::set_var("GW_TEST_CACHE_DIR", dir);
446 f();
447 }
448
449 #[test]
450 fn load_from_disk_returns_fresh_entry() {
451 let _g = env_lock();
452 let dir = tempdir().unwrap();
453 with_cache_dir(dir.path(), || {
454 let repo = std::path::Path::new("/tmp/repo-xyz");
455 let path = cache_path_for(repo).unwrap();
456 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
457 let now = SystemTime::now()
458 .duration_since(UNIX_EPOCH)
459 .unwrap()
460 .as_secs();
461 let file = CacheFile {
462 fetched_at: now,
463 repo: repo.to_string_lossy().into_owned(),
464 prs: [("feat/a".to_string(), PrState::Open)]
465 .into_iter()
466 .collect(),
467 };
468 std::fs::write(&path, serde_json::to_string(&file).unwrap()).unwrap();
469
470 let loaded = load_from_disk(repo).expect("fresh cache");
471 assert_eq!(loaded.get("feat/a"), Some(&PrState::Open));
472 });
473 }
474
475 #[test]
476 fn load_from_disk_rejects_expired_entry() {
477 let _g = env_lock();
478 let dir = tempdir().unwrap();
479 with_cache_dir(dir.path(), || {
480 let repo = std::path::Path::new("/tmp/repo-expired-xyz");
481 let path = cache_path_for(repo).unwrap();
482 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
483 let file = CacheFile {
484 fetched_at: 0, repo: repo.to_string_lossy().into_owned(),
486 prs: HashMap::new(),
487 };
488 std::fs::write(&path, serde_json::to_string(&file).unwrap()).unwrap();
489
490 assert!(load_from_disk(repo).is_none());
491 });
492 }
493
494 #[test]
495 fn load_from_disk_rejects_future_entry() {
496 let _g = env_lock();
497 let dir = tempdir().unwrap();
498 with_cache_dir(dir.path(), || {
499 let repo = std::path::Path::new("/tmp/repo-future-xyz");
500 let path = cache_path_for(repo).unwrap();
501 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
502 let far_future = now_secs().unwrap() + 9999;
503 let file = CacheFile {
504 fetched_at: far_future,
505 repo: repo.to_string_lossy().into_owned(),
506 prs: HashMap::new(),
507 };
508 std::fs::write(&path, serde_json::to_string(&file).unwrap()).unwrap();
509
510 assert!(load_from_disk(repo).is_none());
511 });
512 }
513
514 #[test]
515 fn load_from_disk_rejects_corrupt_file() {
516 let _g = env_lock();
517 let dir = tempdir().unwrap();
518 with_cache_dir(dir.path(), || {
519 let repo = std::path::Path::new("/tmp/repo-corrupt-xyz");
520 let path = cache_path_for(repo).unwrap();
521 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
522 std::fs::write(&path, "not json").unwrap();
523
524 assert!(load_from_disk(repo).is_none());
525 });
526 }
527
528 #[test]
529 fn load_or_fetch_uses_disk_when_fresh() {
530 let _g = env_lock();
531 let _env = EnvGuard::capture(&["GW_TEST_CACHE_DIR", "GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
532 let dir = tempdir().unwrap();
533 with_cache_dir(dir.path(), || {
534 let repo = std::path::Path::new("/tmp/repo-disk-hit-xyz");
535 let path = cache_path_for(repo).unwrap();
536 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
537 let file = CacheFile {
538 fetched_at: now_secs().unwrap(),
539 repo: repo.to_string_lossy().into_owned(),
540 prs: [("feat/cached".to_string(), PrState::Merged)]
541 .into_iter()
542 .collect(),
543 };
544 std::fs::write(&path, serde_json::to_string(&file).unwrap()).unwrap();
545
546 std::env::set_var("GW_TEST_GH_FAIL", "1");
550 let cache = PrCache::load_or_fetch(repo, false);
551 assert_eq!(cache.state("feat/cached"), Some(&PrState::Merged));
552 });
553 }
554
555 #[test]
556 fn load_or_fetch_bypasses_disk_when_no_cache_true() {
557 let _g = env_lock();
558 let _env = EnvGuard::capture(&["GW_TEST_CACHE_DIR", "GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
559 let dir = tempdir().unwrap();
560 with_cache_dir(dir.path(), || {
561 let repo = std::path::Path::new("/tmp/repo-bypass-xyz");
562 let path = cache_path_for(repo).unwrap();
563 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
564 let file = CacheFile {
565 fetched_at: now_secs().unwrap(),
566 repo: repo.to_string_lossy().into_owned(),
567 prs: [("feat/old".to_string(), PrState::Open)]
568 .into_iter()
569 .collect(),
570 };
571 std::fs::write(&path, serde_json::to_string(&file).unwrap()).unwrap();
572
573 std::env::set_var(
574 "GW_TEST_GH_JSON",
575 r#"[{"headRefName":"feat/new","state":"OPEN"}]"#,
576 );
577 let cache = PrCache::load_or_fetch(repo, true);
578 assert_eq!(cache.state("feat/new"), Some(&PrState::Open));
579 assert_eq!(cache.state("feat/old"), None);
580 });
581 }
582
583 #[test]
584 fn load_or_fetch_empty_when_gh_fails_and_no_cache_file() {
585 let _g = env_lock();
586 let _env = EnvGuard::capture(&["GW_TEST_CACHE_DIR", "GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
587 let dir = tempdir().unwrap();
588 with_cache_dir(dir.path(), || {
589 let repo = std::path::Path::new("/tmp/repo-empty-xyz");
590 std::env::set_var("GW_TEST_GH_FAIL", "1");
591 let cache = PrCache::load_or_fetch(repo, false);
592 assert!(cache.state("anything").is_none());
593 });
594 }
595
596 #[test]
597 fn write_to_disk_cleans_up_tmp_file() {
598 let _g = env_lock();
599 let dir = tempdir().unwrap();
600 with_cache_dir(dir.path(), || {
601 let repo = std::path::Path::new("/tmp/repo-atomic-xyz");
602 let mut prs = HashMap::new();
603 prs.insert("feat/x".to_string(), PrState::Open);
604 write_to_disk(repo, &prs);
605
606 let final_path = cache_path_for(repo).unwrap();
607 assert!(final_path.exists(), "final cache file exists");
608
609 let parent = final_path.parent().unwrap();
611 let entries: Vec<_> = std::fs::read_dir(parent).unwrap().flatten().collect();
612 for entry in &entries {
613 let name = entry.file_name();
614 let name_str = name.to_string_lossy();
615 assert!(
616 !name_str.contains(".tmp."),
617 "no tmp file should remain: {}",
618 name_str
619 );
620 }
621 });
622 }
623
624 #[test]
625 fn from_disk_and_fetch_and_persist_split() {
626 let _g = env_lock();
627 let _env = EnvGuard::capture(&["GW_TEST_CACHE_DIR", "GW_TEST_GH_FAIL", "GW_TEST_GH_JSON"]);
628 let dir = tempdir().unwrap();
629 with_cache_dir(dir.path(), || {
630 let repo = std::path::Path::new("/tmp/repo-split-xyz");
631 assert!(PrCache::from_disk(repo).is_none());
633
634 std::env::set_var("GW_TEST_GH_FAIL", "1");
636 let empty = PrCache::fetch_and_persist(repo);
637 assert!(empty.state("anything").is_none());
638 std::env::remove_var("GW_TEST_GH_FAIL");
640
641 std::env::set_var(
643 "GW_TEST_GH_JSON",
644 r#"[{"headRefName":"main","state":"OPEN"}]"#,
645 );
646 let _ = PrCache::fetch_and_persist(repo);
647 let loaded = PrCache::from_disk(repo).expect("written to disk");
649 assert_eq!(loaded.state("main"), Some(&PrState::Open));
650 });
651 }
652
653 #[cfg(unix)]
655 #[test]
656 fn write_to_disk_sweeps_old_orphan_tmp_files() {
657 let _g = env_lock();
658 let _env = EnvGuard::capture(&["GW_TEST_CACHE_DIR"]);
659 let dir = tempdir().unwrap();
660 with_cache_dir(dir.path(), || {
661 let repo = std::path::Path::new("/tmp/repo-sweep-xyz");
662 let final_path = cache_path_for(repo).unwrap();
663 let parent = final_path.parent().unwrap();
664 std::fs::create_dir_all(parent).unwrap();
665
666 let orphan = parent.join("pr-status-orphan.tmp.99999.123456789.0");
669 std::fs::write(&orphan, "stale").unwrap();
670 {
672 use std::ffi::CString;
673 let c_path = CString::new(orphan.to_string_lossy().as_bytes()).unwrap();
674 let times = [libc::timeval {
675 tv_sec: 0,
676 tv_usec: 0,
677 }; 2];
678 unsafe { libc::utimes(c_path.as_ptr(), times.as_ptr()) };
680 }
681
682 let fresh_tmp = parent.join("pr-status-fresh.tmp.123.456.0");
684 std::fs::write(&fresh_tmp, "fresh").unwrap();
685
686 let mut prs = HashMap::new();
688 prs.insert("feat/sweep".to_string(), PrState::Open);
689 write_to_disk(repo, &prs);
690
691 assert!(
692 !orphan.exists(),
693 "old orphan tmp file should have been swept"
694 );
695 assert!(fresh_tmp.exists(), "fresh tmp file should not be swept");
696 assert!(final_path.exists(), "final cache file should exist");
697 });
698 }
699
700 #[test]
703 fn pr_state_variants_are_handled() {
704 fn _must_handle(s: &PrState) {
706 match s {
707 PrState::Open => {}
708 PrState::Merged => {}
709 PrState::Closed => {}
710 PrState::Other => {}
711 }
712 }
713 }
714}