use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::{Mutex, MutexGuard, PoisonError};
use std::time::Duration;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
const DEFAULT_TTL: Duration = Duration::from_secs(30);
const MAX_WINDOWS: usize = 256;
#[derive(Debug, Clone, Deserialize)]
pub struct RegisterRequest {
pub key: String,
#[serde(default)]
pub folders: Vec<PathBuf>,
#[serde(default)]
pub repo: Option<String>,
#[serde(default)]
pub title: Option<String>,
#[serde(default)]
pub pid: Option<u32>,
}
#[derive(Debug, Clone, Serialize)]
pub struct WindowEntry {
pub key: String,
pub folders: Vec<PathBuf>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pid: Option<u32>,
pub last_seen: DateTime<Utc>,
}
pub struct WorktreesRegistry {
windows: Mutex<HashMap<String, WindowEntry>>,
ttl: Duration,
changes: watch::Sender<u64>,
close_pending: Mutex<HashSet<String>>,
}
impl WorktreesRegistry {
#[must_use]
pub fn new() -> Self {
Self {
windows: Mutex::new(HashMap::new()),
ttl: DEFAULT_TTL,
changes: watch::channel(0).0,
close_pending: Mutex::new(HashSet::new()),
}
}
#[must_use]
pub fn subscribe_changes(&self) -> watch::Receiver<u64> {
self.changes.subscribe()
}
fn bump(&self) {
self.changes.send_modify(|v| *v = v.wrapping_add(1));
}
fn lock(&self) -> MutexGuard<'_, HashMap<String, WindowEntry>> {
self.windows.lock().unwrap_or_else(PoisonError::into_inner)
}
pub fn register(&self, req: RegisterRequest) {
let now = Utc::now();
{
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
if !windows.contains_key(&req.key) && windows.len() >= MAX_WINDOWS {
evict_oldest(&mut windows);
}
windows.insert(
req.key.clone(),
WindowEntry {
key: req.key,
folders: req.folders,
repo: req.repo,
title: req.title,
pid: req.pid,
last_seen: now,
},
);
}
self.bump();
}
pub fn heartbeat(&self, key: &str) -> bool {
let now = Utc::now();
let (known, reaped) = {
let mut windows = self.lock();
let reaped = reap(&mut windows, self.ttl, now);
let known = match windows.get_mut(key) {
Some(entry) => {
entry.last_seen = now;
true
}
None => false,
};
(known, reaped)
};
if reaped > 0 {
self.bump();
}
known
}
pub fn unregister(&self, key: &str) -> bool {
let now = Utc::now();
let (removed, reaped) = {
let mut windows = self.lock();
let removed = windows.remove(key).is_some();
let reaped = reap(&mut windows, self.ttl, now);
(removed, reaped)
};
self.take_close_pending(key);
if removed || reaped > 0 {
self.bump();
}
removed
}
pub fn mark_close_pending(&self, key: &str) {
self.close_pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(key.to_string());
}
pub fn take_close_pending(&self, key: &str) -> bool {
self.close_pending
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(key)
}
pub fn list(&self) -> Vec<WindowEntry> {
let now = Utc::now();
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
sorted_entries(&windows)
}
pub fn first_folder(&self, key: &str) -> Option<PathBuf> {
let windows = self.lock();
windows.get(key).and_then(|e| e.folders.first().cloned())
}
pub fn open_folders(&self) -> Vec<PathBuf> {
let now = Utc::now();
let mut windows = self.lock();
reap(&mut windows, self.ttl, now);
let mut folders: Vec<PathBuf> = windows
.values()
.flat_map(|e| e.folders.iter().cloned())
.collect();
folders.sort();
folders.dedup();
folders
}
}
impl Default for WorktreesRegistry {
fn default() -> Self {
Self::new()
}
}
fn reap(windows: &mut HashMap<String, WindowEntry>, ttl: Duration, now: DateTime<Utc>) -> usize {
let max_age = ttl.as_secs() as i64;
let before = windows.len();
windows.retain(|_, e| (now - e.last_seen).num_seconds() <= max_age);
before - windows.len()
}
fn evict_oldest(windows: &mut HashMap<String, WindowEntry>) {
let oldest = windows
.values()
.min_by(|a, b| {
a.last_seen
.cmp(&b.last_seen)
.then_with(|| a.key.cmp(&b.key))
})
.map(|e| e.key.clone());
if let Some(key) = oldest {
windows.remove(&key);
}
}
fn sorted_entries(windows: &HashMap<String, WindowEntry>) -> Vec<WindowEntry> {
let mut entries: Vec<WindowEntry> = windows.values().cloned().collect();
entries.sort_by(|a, b| a.repo.cmp(&b.repo).then_with(|| a.key.cmp(&b.key)));
entries
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
fn register_request(key: &str, repo: Option<&str>, folder: &str) -> RegisterRequest {
RegisterRequest {
key: key.to_string(),
folders: vec![PathBuf::from(folder)],
repo: repo.map(str::to_string),
title: Some(format!("{key}-title")),
pid: Some(1234),
}
}
#[test]
fn list_is_empty_initially() {
let reg = WorktreesRegistry::new();
assert!(reg.list().is_empty());
}
#[test]
fn register_then_list_round_trips() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
let windows = reg.list();
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].key, "w1");
assert_eq!(windows[0].repo.as_deref(), Some("repo-a"));
}
#[test]
fn register_is_idempotent_upsert() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", Some("repo-a"), "/tmp/a"));
reg.register(register_request("w1", Some("repo-b"), "/tmp/b"));
let windows = reg.list();
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].repo.as_deref(), Some("repo-b"));
}
#[test]
fn heartbeat_reports_known_and_unknown() {
let reg = WorktreesRegistry::new();
assert!(!reg.heartbeat("w1"));
reg.register(register_request("w1", None, "/tmp/a"));
assert!(reg.heartbeat("w1"));
}
#[test]
fn unregister_removes() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
assert!(reg.unregister("w1"));
assert!(!reg.unregister("w1"));
}
#[test]
fn first_folder_returns_first_folder_or_none() {
let reg = WorktreesRegistry::new();
assert!(reg.first_folder("missing").is_none());
reg.register(register_request("w1", None, "/tmp/a"));
assert_eq!(reg.first_folder("w1"), Some(PathBuf::from("/tmp/a")));
reg.register(RegisterRequest {
key: "w2".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
});
assert!(reg.first_folder("w2").is_none());
}
#[test]
fn open_folders_dedups_and_sorts_across_windows() {
let reg = WorktreesRegistry::new();
assert!(reg.open_folders().is_empty());
reg.register(register_request("w1", Some("repo-a"), "/tmp/shared"));
reg.register(RegisterRequest {
key: "w2".to_string(),
folders: vec![PathBuf::from("/tmp/shared"), PathBuf::from("/tmp/b")],
repo: Some("repo-a".to_string()),
title: None,
pid: None,
});
reg.register(register_request("w3", Some("repo-b"), "/tmp/a"));
assert_eq!(
reg.open_folders(),
vec![
PathBuf::from("/tmp/a"),
PathBuf::from("/tmp/b"),
PathBuf::from("/tmp/shared"),
]
);
}
#[test]
fn open_folders_reaps_stale_windows() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
windows.insert(
"fresh".to_string(),
WindowEntry {
key: "fresh".to_string(),
folders: vec![PathBuf::from("/tmp/fresh")],
repo: None,
title: None,
pid: None,
last_seen: Utc::now(),
},
);
windows.insert(
"stale".to_string(),
WindowEntry {
key: "stale".to_string(),
folders: vec![PathBuf::from("/tmp/stale")],
repo: None,
title: None,
pid: None,
last_seen: Utc::now() - chrono::Duration::seconds(120),
},
);
}
assert_eq!(reg.open_folders(), vec![PathBuf::from("/tmp/fresh")]);
}
#[test]
fn reap_evicts_only_stale_entries() {
let now = Utc::now();
let mut windows = HashMap::new();
windows.insert(
"fresh".to_string(),
WindowEntry {
key: "fresh".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(5),
},
);
windows.insert(
"stale".to_string(),
WindowEntry {
key: "stale".to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen: now - chrono::Duration::seconds(120),
},
);
reap(&mut windows, DEFAULT_TTL, now);
assert!(windows.contains_key("fresh"));
assert!(!windows.contains_key("stale"));
}
fn entry_at(key: &str, last_seen: DateTime<Utc>) -> WindowEntry {
WindowEntry {
key: key.to_string(),
folders: vec![],
repo: None,
title: None,
pid: None,
last_seen,
}
}
#[test]
fn evict_oldest_removes_oldest_with_key_tiebreak() {
let now = Utc::now();
let mut windows = HashMap::new();
windows.insert("young".to_string(), entry_at("young", now));
windows.insert(
"old-b".to_string(),
entry_at("old-b", now - chrono::Duration::seconds(10)),
);
windows.insert(
"old-a".to_string(),
entry_at("old-a", now - chrono::Duration::seconds(10)),
);
evict_oldest(&mut windows);
assert!(!windows.contains_key("old-a"));
assert!(windows.contains_key("old-b"));
assert!(windows.contains_key("young"));
let mut empty: HashMap<String, WindowEntry> = HashMap::new();
evict_oldest(&mut empty);
assert!(empty.is_empty());
}
#[test]
fn register_at_cap_evicts_only_the_oldest() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
let base = Utc::now();
for i in 0..MAX_WINDOWS {
let key = format!("w{i:03}");
windows.insert(
key.clone(),
entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
);
}
}
reg.register(register_request("fresh", None, "/tmp/f"));
let windows = reg.lock();
assert_eq!(windows.len(), MAX_WINDOWS);
assert!(windows.contains_key("fresh"));
assert!(!windows.contains_key(&format!("w{:03}", MAX_WINDOWS - 1)));
assert!(windows.contains_key("w000"));
}
#[test]
fn register_upsert_at_cap_does_not_evict() {
let reg = WorktreesRegistry::new();
{
let mut windows = reg.lock();
let base = Utc::now();
for i in 0..MAX_WINDOWS {
let key = format!("w{i:03}");
windows.insert(
key.clone(),
entry_at(&key, base - chrono::Duration::milliseconds(i as i64)),
);
}
}
let oldest = format!("w{:03}", MAX_WINDOWS - 1);
reg.register(register_request(&oldest, Some("r"), "/tmp/a"));
let windows = reg.lock();
assert_eq!(windows.len(), MAX_WINDOWS);
assert!(windows.contains_key(&oldest));
assert!(windows.contains_key("w000"));
}
#[test]
fn sorted_entries_orders_by_repo_then_key() {
let now = Utc::now();
let mut windows = HashMap::new();
for (key, repo) in [("z", "repo-a"), ("a", "repo-b"), ("m", "repo-a")] {
windows.insert(
key.to_string(),
WindowEntry {
key: key.to_string(),
folders: vec![],
repo: Some(repo.to_string()),
title: None,
pid: None,
last_seen: now,
},
);
}
let entries = sorted_entries(&windows);
let ordered: Vec<(&str, &str)> = entries
.iter()
.map(|e| (e.key.as_str(), e.repo.as_deref().unwrap()))
.collect();
assert_eq!(
ordered,
vec![("m", "repo-a"), ("z", "repo-a"), ("a", "repo-b")]
);
}
#[test]
fn default_constructs_an_empty_registry() {
let reg = WorktreesRegistry::default();
assert!(reg.lock().is_empty());
}
#[test]
fn subscribe_changes_starts_seen_and_register_bumps() {
let reg = WorktreesRegistry::new();
let mut rx = reg.subscribe_changes();
assert!(!rx.has_changed().unwrap());
reg.register(register_request("w1", None, "/tmp/a"));
assert!(rx.has_changed().unwrap(), "register should bump");
rx.borrow_and_update();
assert!(!rx.has_changed().unwrap());
}
#[test]
fn unregister_bumps_only_when_it_removes() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
let rx = reg.subscribe_changes();
assert!(!reg.unregister("ghost"));
assert!(
!rx.has_changed().unwrap(),
"a no-op unregister must not bump"
);
assert!(reg.unregister("w1"));
assert!(
rx.has_changed().unwrap(),
"a removing unregister should bump"
);
}
#[test]
fn heartbeat_bumps_only_when_it_reaps() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
let rx = reg.subscribe_changes();
assert!(reg.heartbeat("w1"));
assert!(!rx.has_changed().unwrap(), "a pure heartbeat must not bump");
{
let mut windows = reg.lock();
windows.insert(
"stale".to_string(),
entry_at("stale", Utc::now() - chrono::Duration::seconds(120)),
);
}
assert!(reg.heartbeat("w1"));
assert!(
rx.has_changed().unwrap(),
"a heartbeat that reaps a stale sibling should bump"
);
}
#[test]
fn close_pending_is_taken_once_then_cleared() {
let reg = WorktreesRegistry::new();
assert!(!reg.take_close_pending("w1"));
reg.mark_close_pending("w1");
assert!(reg.take_close_pending("w1"));
assert!(!reg.take_close_pending("w1"));
}
#[test]
fn unregister_clears_a_pending_close_directive() {
let reg = WorktreesRegistry::new();
reg.register(register_request("w1", None, "/tmp/a"));
reg.mark_close_pending("w1");
assert!(reg.unregister("w1"));
assert!(!reg.take_close_pending("w1"));
}
}