sqlite_graphrag/commands/
slots.rs1use clap::{Args, Subcommand};
10use serde::Serialize;
11
12use crate::errors::AppError;
13use crate::llm_slots::{slot_path, slots_dir};
14use crate::output::emit_json_compact;
15use crate::output::OutputFormat;
16
17#[derive(Debug, Args)]
20pub struct SlotsArgs {
21 #[command(subcommand)]
22 pub cmd: SlotsCmd,
23}
24
25#[derive(Debug, Subcommand)]
26pub enum SlotsCmd {
27 Status(SlotsStatusArgs),
29 Release {
31 #[arg(long)]
33 slot_id: u32,
34 #[arg(long)]
36 yes: bool,
37 },
38 Cleanup {
40 #[arg(long, default_value_t = 3600)]
42 stale_after: u64,
43 #[arg(long)]
45 yes: bool,
46 #[arg(long)]
48 dry_run: bool,
49 },
50}
51
52#[derive(Debug, clap::Args)]
53pub struct SlotsStatusArgs {
54 #[arg(long, value_enum, default_value_t = OutputFormat::Json)]
56 pub format: OutputFormat,
57}
58
59#[derive(Serialize)]
60struct SlotEntry {
61 slot_id: u32,
62 path: String,
63 age_secs: u64,
64 pid_hint: Option<u32>,
65}
66
67#[derive(Serialize)]
68struct SlotsStatusOutput {
69 action: &'static str,
70 max_concurrency: u32,
71 active: usize,
72 free: usize,
73 slots: Vec<SlotEntry>,
74 elapsed_ms: u64,
75}
76
77pub fn run(args: SlotsArgs) -> Result<(), AppError> {
78 run_cmd(args.cmd)
79}
80
81fn run_cmd(cmd: SlotsCmd) -> Result<(), AppError> {
82 match cmd {
83 SlotsCmd::Status(args) => run_status(args),
84 SlotsCmd::Release { slot_id, yes } => run_release(slot_id, yes),
85 SlotsCmd::Cleanup {
86 stale_after,
87 yes,
88 dry_run,
89 } => run_cleanup(stale_after, yes, dry_run),
90 }
91}
92
93fn run_status(args: SlotsStatusArgs) -> Result<(), AppError> {
94 let start = std::time::Instant::now();
95 let max = crate::llm_slots::default_max_concurrency();
96 let dir = slots_dir();
97 let mut entries: Vec<SlotEntry> = Vec::new();
98
99 if dir.is_dir() {
100 for slot_id in 0..max {
101 let path = slot_path(slot_id);
102 if path.is_file() {
103 let age_secs = path
104 .metadata()
105 .and_then(|m| m.modified())
106 .ok()
107 .and_then(|t| t.elapsed().ok())
108 .map(|d| d.as_secs())
109 .unwrap_or(0);
110 let pid_hint = std::fs::read_to_string(&path)
111 .ok()
112 .and_then(|s| s.trim().parse::<u32>().ok());
113 entries.push(SlotEntry {
114 slot_id,
115 path: path.to_string_lossy().into_owned(),
116 age_secs,
117 pid_hint,
118 });
119 }
120 }
121 }
122
123 let output = SlotsStatusOutput {
124 action: "slots_status",
125 max_concurrency: max,
126 active: entries.len(),
127 free: (max as usize).saturating_sub(entries.len()),
128 slots: entries,
129 elapsed_ms: start.elapsed().as_millis() as u64,
130 };
131
132 if matches!(args.format, OutputFormat::Json) {
133 let json = serde_json::to_string_pretty(&output).map_err(AppError::Json)?;
134 println!("{json}");
135 } else {
136 println!("max_concurrency: {}", output.max_concurrency);
137 println!("active: {} / free: {}", output.active, output.free);
138 for s in &output.slots {
139 let pid = s.pid_hint.map(|p| p.to_string()).unwrap_or_default();
140 println!(
141 " slot {} — age={}s pid={} {}",
142 s.slot_id, s.age_secs, pid, s.path
143 );
144 }
145 }
146 Ok(())
147}
148
149fn run_release(slot_id: u32, yes: bool) -> Result<(), AppError> {
150 let path = slot_path(slot_id);
151 if !path.is_file() {
152 return Err(AppError::NotFound(format!(
153 "slot {slot_id} is not held (no file at {})",
154 path.display()
155 )));
156 }
157 if !yes {
158 eprintln!(
159 "About to release slot {slot_id} at {}. Pass --yes to skip confirmation.",
160 path.display()
161 );
162 }
163 std::fs::remove_file(&path).map_err(AppError::Io)?;
164 let out = serde_json::json!({
165 "action": "slot_released",
166 "slot_id": slot_id,
167 "path": path.to_string_lossy(),
168 });
169 let _ = emit_json_compact(&out);
170 Ok(())
171}
172
173fn run_cleanup(stale_after: u64, yes: bool, dry_run: bool) -> Result<(), AppError> {
174 let start = std::time::Instant::now();
175 let max = crate::llm_slots::default_max_concurrency();
176 let mut removed: Vec<u32> = Vec::new();
177 for slot_id in 0..max {
178 let path = slot_path(slot_id);
179 if !path.is_file() {
180 continue;
181 }
182 let age = path
183 .metadata()
184 .and_then(|m| m.modified())
185 .ok()
186 .and_then(|t| t.elapsed().ok())
187 .map(|d| d.as_secs())
188 .unwrap_or(0);
189 if age >= stale_after {
190 if !dry_run {
191 if let Err(e) = std::fs::remove_file(&path) {
192 tracing::warn!(target: "slots", slot_id, error = %e, "stale slot removal failed");
193 continue;
194 }
195 }
196 removed.push(slot_id);
197 }
198 }
199 let out = serde_json::json!({
200 "action": if dry_run { "slots_cleanup_dry_run" } else { "slots_cleanup" },
201 "stale_after_secs": stale_after,
202 "removed": removed,
203 "removed_count": removed.len(),
204 "elapsed_ms": start.elapsed().as_millis() as u64,
205 "yes": yes,
206 });
207 let _ = emit_json_compact(&out);
208 Ok(())
209}
210
211#[cfg(test)]
215mod tests {
216 use super::*;
217 use crate::llm_slots::acquire_llm_slot;
218
219 #[test]
220 fn acquire_then_drop_releases_slot() {
221 let _ = std::fs::remove_dir_all(crate::llm_slots::slots_dir());
222 let guard = acquire_llm_slot(2, 5).expect("acquire");
223 let path = slot_path(guard.slot_id());
224 assert!(path.is_file(), "slot file must exist after acquire");
225 drop(guard);
226 assert!(
227 !path.is_file(),
228 "slot file must be removed after Drop (RAII guarantee)"
229 );
230 }
231}