use std::cell::{Cell, RefCell};
use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use chrono::{DateTime, Utc};
use super::size_index::{DirSize, DirSizeIndex};
use super::survey::{
Classification, DiskProject, DiskRoot, DiskSurvey, DiskWorktree, GroupBy, KeepListReport,
PrRef, Reason, ReasonCode, SizeNote, TierCounts, WorktreeFacts, WorktreeTier, classify_tier,
group_by_session,
};
use crate::session_manager::worktree_keep_list::KeepList;
use crate::session_manager::worktree_ownership::{SentinelOwner, read_sentinel_owner};
use crate::session_manager::worktree_reclaim::{
AgentStateProbe, BranchPrState, LiveClaims, ReclaimGate, ReclaimVerdict, classify,
};
use crate::session_manager::worktree_registry::{ScannedWorktree, scan_registered_worktrees};
use crate::session_manager::worktree_safety::DirtyWorktree;
pub(crate) struct DiskProbes<'a> {
pub pr_state: &'a dyn Fn(&ScannedWorktree, Option<Duration>) -> BranchPrState,
pub claims: &'a LiveClaims,
pub agent_state: AgentStateProbe<'a>,
pub dirt: &'a dyn Fn(&Path) -> Option<DirtyWorktree>,
pub measure: &'a dyn Fn(&Path, Option<Duration>) -> Option<DirSize>,
pub now: &'a dyn Fn() -> Instant,
}
pub(crate) const SYSTEM_CLOCK: fn() -> Instant = Instant::now;
enum Measured {
Figure(Option<DirSize>),
DeadlineSpent,
}
enum Budget {
Left(Option<Duration>),
Spent,
}
fn budget_at(deadline: Option<Instant>, now: &dyn Fn() -> Instant) -> Budget {
match deadline {
None => Budget::Left(None),
Some(d) => match d.checked_duration_since(now()) {
Some(left) if !left.is_zero() => Budget::Left(Some(left)),
_ => Budget::Spent,
},
}
}
impl Measured {
fn figure(self) -> Option<DirSize> {
match self {
Self::Figure(size) => size,
Self::DeadlineSpent => None,
}
}
}
pub(crate) fn run(
repos_root: &Path,
keep_list: &KeepList,
probes: &DiskProbes<'_>,
deadline: Option<Instant>,
project: Option<&str>,
group_by: GroupBy,
adopted: &[PathBuf],
) -> DiskSurvey {
let partial = Cell::new(false);
let measure = |path: &Path| -> Measured {
match budget_at(deadline, probes.now) {
Budget::Spent => {
partial.set(true);
Measured::DeadlineSpent
}
Budget::Left(left) => Measured::Figure((probes.measure)(path, left)),
}
};
let mut grouped: BTreeMap<PathBuf, Vec<DiskWorktree>> = BTreeMap::new();
let cache: RefCell<HashMap<PathBuf, Option<DirtyWorktree>>> = RefCell::new(HashMap::new());
let probe_dirt = |path: &Path| -> Option<DirtyWorktree> {
if let Some(hit) = cache.borrow().get(path) {
return hit.clone();
}
let answer = (probes.dirt)(path);
cache
.borrow_mut()
.insert(path.to_path_buf(), answer.clone());
answer
};
for scanned in scan_registered_worktrees(repos_root, adopted) {
if !selected(&scanned, repos_root, project) {
continue;
}
let row = match budget_at(deadline, probes.now) {
Budget::Spent => {
partial.set(true);
not_inspected(&scanned)
}
Budget::Left(left) => inspect(&scanned, keep_list, probes, &probe_dirt, &measure, left),
};
grouped
.entry(scanned.project.clone())
.or_default()
.push(row);
}
let mut counts = TierCounts::default();
let mut stale_bytes = 0u64;
let mut stale_measured = 0usize;
let mut projects = Vec::new();
for (path, worktrees) in grouped {
for wt in &worktrees {
match wt.tier {
WorktreeTier::Stale => {
counts.stale += 1;
if let Some(bytes) = wt.bytes {
stale_bytes = stale_bytes.saturating_add(bytes);
stale_measured += 1;
}
}
WorktreeTier::Review => counts.review += 1,
WorktreeTier::Keep => counts.keep += 1,
WorktreeTier::Missing => counts.missing += 1,
}
}
let size = measure(&path).figure();
projects.push(DiskProject {
name: project_name(&path, repos_root),
bytes: size.as_ref().map(|s| s.bytes),
size: size.as_ref().map(note),
path,
worktrees,
});
}
let root_size = measure(repos_root).figure();
let by_session = match group_by {
GroupBy::None => None,
GroupBy::Session => Some(group_by_session(&projects)),
};
DiskSurvey {
generated_at: Utc::now().to_rfc3339(),
partial: partial.get(),
keep_list: KeepListReport {
patterns: keep_list.patterns().to_vec(),
invalid: keep_list.invalid().to_vec(),
error: keep_list.error().map(str::to_string),
},
root: DiskRoot {
path: repos_root.to_path_buf(),
bytes: root_size.as_ref().map(|s| s.bytes),
size: root_size.as_ref().map(note),
projects,
counts,
stale_bytes,
stale_measured,
},
by_session,
}
}
fn inspect(
scanned: &ScannedWorktree,
keep_list: &KeepList,
probes: &DiskProbes<'_>,
probe_dirt: &dyn Fn(&Path) -> Option<DirtyWorktree>,
measure: &dyn Fn(&Path) -> Measured,
left: Option<Duration>,
) -> DiskWorktree {
let pr = (probes.pr_state)(scanned, left);
let claim = probes.claims.claim_state(&scanned.path);
let verdict = classify(
&scanned.path,
scanned.admission,
&claim,
&pr,
probe_dirt,
probes.agent_state,
&probes.claims.owners,
keep_list,
);
let facts = WorktreeFacts {
path: &scanned.path,
branch: scanned.branch.as_deref(),
admission: scanned.admission,
claim: &claim,
pr: &pr,
verdict: &verdict,
};
let classification = classify_tier(&facts, keep_list, probe_dirt);
let missing = classification.tier == WorktreeTier::Missing;
let size = if missing {
None
} else {
match measure(&scanned.path) {
Measured::Figure(size) => size,
Measured::DeadlineSpent => return not_inspected(scanned),
}
};
let build_dir_bytes = if missing {
None
} else {
build_dir_bytes(&scanned.path, measure)
};
let owning = owning_session(&claim, &scanned.path);
row(
scanned,
classification,
&verdict,
&pr,
&claim,
size,
build_dir_bytes,
owning,
)
}
fn owning_session(
claim: &crate::session_manager::worktree_reclaim_claim::ClaimState,
path: &Path,
) -> Option<String> {
if let Some(session) = claiming_session(claim) {
return Some(session);
}
match read_sentinel_owner(path) {
SentinelOwner::Agent(owner, _) => Some(owner.parent_session_id.0.to_string()),
SentinelOwner::Known(session, _) => Some(session.to_string()),
SentinelOwner::Unknown => None,
}
}
fn build_dir_bytes(worktree: &Path, measure: &dyn Fn(&Path) -> Measured) -> Option<u64> {
let Ok(entries) = std::fs::read_dir(worktree) else {
return None;
};
let mut total = 0u64;
for entry in entries {
let Ok(entry) = entry else {
return None;
};
if !entry.file_name().to_string_lossy().starts_with("target") {
continue;
}
if !entry.file_type().is_ok_and(|t| t.is_dir()) {
continue;
}
match measure(&entry.path()) {
Measured::Figure(Some(size)) => total = total.saturating_add(size.bytes),
Measured::Figure(None) | Measured::DeadlineSpent => return None,
}
}
Some(total)
}
#[allow(clippy::too_many_arguments)]
fn row(
scanned: &ScannedWorktree,
classification: Classification,
verdict: &ReclaimVerdict,
pr: &BranchPrState,
claim: &crate::session_manager::worktree_reclaim_claim::ClaimState,
size: Option<DirSize>,
build_dir_bytes: Option<u64>,
owning_session: Option<String>,
) -> DiskWorktree {
let (gate, reason) = match verdict {
ReclaimVerdict::Reclaimable { .. } | ReclaimVerdict::ReclaimableLandedContent { .. } => {
(None, None)
}
ReclaimVerdict::Blocked { gate, reason }
| ReclaimVerdict::BlockedByAgent { gate, reason } => (Some(*gate), Some(reason.clone())),
};
DiskWorktree {
id: scanned.path.to_string_lossy().into_owned(),
path: scanned.path.clone(),
branch: scanned.branch.clone(),
tier: classification.tier,
reasons: classification.reasons,
gate,
reason,
reclaimable: verdict.is_reclaimable(),
bytes: size.as_ref().map(|s| s.bytes),
size: size.as_ref().map(note),
pr: PrRef::from_state(pr),
session: claiming_session(claim),
owning_session,
build_dir_bytes,
}
}
fn not_inspected(scanned: &ScannedWorktree) -> DiskWorktree {
DiskWorktree {
id: scanned.path.to_string_lossy().into_owned(),
path: scanned.path.clone(),
branch: scanned.branch.clone(),
tier: WorktreeTier::Review,
reasons: vec![Reason {
code: ReasonCode::UnknownBranchState,
detail: "the survey deadline was reached before this worktree was inspected"
.to_string(),
}],
gate: Some(ReclaimGate::Deadline),
reason: Some("survey deadline reached before inspection".to_string()),
reclaimable: false,
bytes: None,
size: None,
pr: None,
session: None,
owning_session: None,
build_dir_bytes: None,
}
}
fn claiming_session(
claim: &crate::session_manager::worktree_reclaim_claim::ClaimState,
) -> Option<String> {
use crate::session_manager::worktree_reclaim_claim::ClaimState;
match claim {
ClaimState::Unclaimed | ClaimState::DeadClaimsDiscarded { .. } => None,
ClaimState::CallerNested { session }
| ClaimState::CallerWorkspace { session }
| ClaimState::ForeignNested { session, .. }
| ClaimState::Foreign { session, .. } => Some(session.clone()),
}
}
pub(crate) fn measure(
index: &mut DirSizeIndex,
path: &Path,
budget: Option<Duration>,
) -> Option<DirSize> {
match index.measure_within(path, budget) {
Ok(size) => Some(size),
Err(e) => {
tracing::debug!(path = %path.display(), error = %e, "disk-survey: no byte figure (#6927)");
None
}
}
}
fn note(size: &DirSize) -> SizeNote {
SizeNote {
from_cache: size.from_cache,
truncated: size.truncated,
measured_at: DateTime::<Utc>::from(size.measured_at).to_rfc3339(),
unreadable: size.unreadable.len(),
}
}
fn project_name(project: &Path, repos_root: &Path) -> String {
project
.strip_prefix(repos_root)
.unwrap_or(project)
.to_string_lossy()
.into_owned()
}
fn selected(scanned: &ScannedWorktree, repos_root: &Path, project: Option<&str>) -> bool {
let Some(filter) = project else {
return true;
};
let name = project_name(&scanned.project, repos_root);
name == filter
|| scanned
.project
.file_name()
.is_some_and(|n| n.to_string_lossy() == filter)
|| scanned.project.starts_with(filter)
}