Skip to content

Commit cb8bf7f

Browse files
committed
feat: add Ctrl+C cancellation support for script execution
- Implement real-time Ctrl+C detection using crossterm event polling - Add graceful process termination for both parallel and sequential execution - Introduce cancellation flag shared across all execution threads - Kill child processes when cancellation is requested
1 parent e8aa449 commit cb8bf7f

File tree

2 files changed

+88
-10
lines changed

2 files changed

+88
-10
lines changed

Cargo.lock

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

src/scripts.rs

Lines changed: 87 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,20 @@ impl ScriptManager {
2424
cursor, execute,
2525
style::{Color, Print, SetForegroundColor},
2626
terminal::{self, Clear, ClearType},
27+
event::{self, Event, KeyCode, KeyEvent},
2728
};
2829
use parking_lot::Mutex;
2930
use std::io::stdout;
30-
use std::sync::atomic::{AtomicUsize, Ordering};
31+
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
3132
use std::sync::Arc;
3233
use std::thread::{self, JoinHandle};
33-
use std::time::Instant;
34+
use std::time::{Duration, Instant};
3435

3536
// Create a mutex-wrapped stdout for thread-safe access
3637
let stdout = Arc::new(Mutex::new(stdout()));
38+
39+
// Create cancellation flag for Ctrl+C handling
40+
let cancelled = Arc::new(AtomicBool::new(false));
3741

3842
// Initialize terminal
3943
{
@@ -60,7 +64,7 @@ impl ScriptManager {
6064
}
6165

6266
let commands = Arc::new(script_config.commands.clone());
63-
let mut handles: Vec<JoinHandle<Result<()>>> = vec![]; // Add type annotation here
67+
let mut handles: Vec<JoinHandle<Result<()>>> = vec![];
6468

6569
// Calculate screen layout
6670
let (term_width, term_height) = terminal::size()?;
@@ -95,15 +99,46 @@ impl ScriptManager {
9599
// For parallel execution with max threads control
96100
let active_threads = Arc::new(AtomicUsize::new(0));
97101

98-
// Execute commands
102+
// Start Ctrl+C detection thread
103+
let cancelled_clone = Arc::clone(&cancelled);
104+
let stdout_clone = Arc::clone(&stdout);
105+
let ctrl_c_handle = thread::spawn(move || {
106+
loop {
107+
if event::poll(Duration::from_millis(100)).unwrap_or(false) {
108+
if let Ok(Event::Key(KeyEvent { code: KeyCode::Char('c'), modifiers, .. })) = event::read() {
109+
if modifiers.contains(crossterm::event::KeyModifiers::CONTROL) {
110+
cancelled_clone.store(true, Ordering::SeqCst);
111+
let mut stdout = stdout_clone.lock();
112+
let _ = execute!(
113+
stdout,
114+
cursor::MoveTo(0, 0),
115+
SetForegroundColor(Color::Red),
116+
Print("⚠️ Cancelling execution... (Ctrl+C detected)")
117+
);
118+
break;
119+
}
120+
}
121+
}
122+
if cancelled_clone.load(Ordering::SeqCst) {
123+
break;
124+
}
125+
}
126+
});
127+
99128
// Execute commands
100129
for (idx, cmd) in commands.iter().enumerate() {
130+
// Check for cancellation before starting new commands
131+
if cancelled.load(Ordering::SeqCst) {
132+
break;
133+
}
134+
101135
let cmd = cmd.clone();
102136
let name = name.to_string();
103137
let stdout = Arc::clone(&stdout);
104138
let y_pos = idx as u16 * section_height;
105139
// let term_width = term_width;
106140
let active_threads = Arc::clone(&active_threads);
141+
let cancelled_flag = Arc::clone(&cancelled);
107142

108143
// If parallel execution is disabled, wait for previous command to complete
109144
if !script_config.parallel {
@@ -146,6 +181,15 @@ impl ScriptManager {
146181
let handle = thread::spawn(move || -> Result<()> {
147182
let command_start = Instant::now();
148183

184+
// Check for cancellation before starting
185+
if cancelled_flag.load(Ordering::SeqCst) {
186+
active_threads.fetch_sub(1, Ordering::SeqCst);
187+
return Err(HookError::ScriptExecutionError {
188+
script_name: name,
189+
reason: "Execution cancelled".to_string(),
190+
});
191+
}
192+
149193
// Build command with working directory and environment variables
150194
let mut child = std::process::Command::new("sh");
151195
child.arg("-c").arg(&cmd.command);
@@ -173,10 +217,19 @@ impl ScriptManager {
173217
let mut current_line = y_pos + 1;
174218

175219
// Process output in real-time
176-
177220
if let Some(child_stdout) = child.stdout.take() {
178221
let reader = std::io::BufReader::new(child_stdout);
179222
for line in std::io::BufRead::lines(reader).flatten() {
223+
// Check for cancellation during output processing
224+
if cancelled_flag.load(Ordering::SeqCst) {
225+
let _ = child.kill();
226+
active_threads.fetch_sub(1, Ordering::SeqCst);
227+
return Err(HookError::ScriptExecutionError {
228+
script_name: name,
229+
reason: "Execution cancelled".to_string(),
230+
});
231+
}
232+
180233
let mut stdout = stdout.lock();
181234
execute!(
182235
stdout,
@@ -188,12 +241,23 @@ impl ScriptManager {
188241
}
189242
}
190243

244+
// Check for cancellation before waiting for process completion
245+
if cancelled_flag.load(Ordering::SeqCst) {
246+
let _ = child.kill();
247+
active_threads.fetch_sub(1, Ordering::SeqCst);
248+
return Err(HookError::ScriptExecutionError {
249+
script_name: name,
250+
reason: "Execution cancelled".to_string(),
251+
});
252+
}
253+
191254
let status = child.wait().map_err(|e| HookError::ScriptExecutionError {
192255
script_name: name.clone(),
193256
reason: e.to_string(),
194257
})?;
195258

196259
if !status.success() {
260+
active_threads.fetch_sub(1, Ordering::SeqCst);
197261
return Err(HookError::ScriptExecutionError {
198262
script_name: name,
199263
reason: format!("Command '{}' failed with status {}", cmd.command, status),
@@ -218,6 +282,8 @@ impl ScriptManager {
218282

219283
// Wait for all remaining threads to complete
220284
let mut all_successful = true;
285+
let was_cancelled = cancelled.load(Ordering::SeqCst);
286+
221287
while let Some(handle) = handles.pop() {
222288
match handle.join().map_err(|_| HookError::ScriptExecutionError {
223289
script_name: name.to_string(),
@@ -240,21 +306,29 @@ impl ScriptManager {
240306
}
241307
}
242308

309+
// Signal the Ctrl+C thread to stop and wait for it
310+
cancelled.store(true, Ordering::SeqCst);
311+
let _ = ctrl_c_handle.join();
312+
243313
// Show final status
244314
let total_duration = start_time.elapsed();
245315
{
246316
let mut stdout = stdout.lock();
247317
execute!(
248318
stdout,
249319
cursor::MoveTo(0, term_height - 1),
250-
SetForegroundColor(if all_successful {
320+
SetForegroundColor(if was_cancelled {
321+
Color::Yellow
322+
} else if all_successful {
251323
Color::Green
252324
} else {
253325
Color::Red
254326
}),
255327
Print(format!(
256328
"Execution {} in {:.2?} ({})",
257-
if all_successful {
329+
if was_cancelled {
330+
"cancelled"
331+
} else if all_successful {
258332
"completed"
259333
} else {
260334
"failed"
@@ -278,10 +352,14 @@ impl ScriptManager {
278352
execute!(stdout, terminal::LeaveAlternateScreen)?;
279353
}
280354

281-
if !all_successful {
355+
if was_cancelled || !all_successful {
282356
return Err(HookError::ScriptExecutionError {
283357
script_name: name.to_string(),
284-
reason: "One or more commands failed".to_string(),
358+
reason: if was_cancelled {
359+
"Execution was cancelled by user".to_string()
360+
} else {
361+
"One or more commands failed".to_string()
362+
},
285363
});
286364
}
287365

0 commit comments

Comments
 (0)