Skip to content

Commit f2c9c46

Browse files
arcanisarcanis
committed
fixes
1 parent 3bb8959 commit f2c9c46

File tree

16 files changed

+780
-19
lines changed

16 files changed

+780
-19
lines changed

‎packages/zpm-switch/src/commands/switch/daemon_kill.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ impl DaemonKillCommand {
4444
return Ok(());
4545
}
4646

47-
if daemons::kill_process(daemon.pid) {
47+
if daemons::kill_daemon_gracefully(daemon.pid) {
4848
daemons::unregister_daemon(&detected_root)?;
4949
println!(
5050
"{} Stopped daemon for {} (PID: {})",

‎packages/zpm-switch/src/commands/switch/daemon_kill_all.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ impl DaemonKillAllCommand {
3535
continue;
3636
}
3737

38-
if daemons::kill_process(daemon.pid) {
38+
if daemons::kill_daemon_gracefully(daemon.pid) {
3939
daemons::unregister_daemon(&daemon.project_cwd)?;
4040
println!(
4141
"{} Stopped daemon for {} (PID: {})",

‎packages/zpm-switch/src/commands/switch/daemon_open.rs‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@ impl DaemonOpenCommand {
4545
println!("ws://127.0.0.1:{}", existing.port);
4646
return Ok(());
4747
}
48-
// Process alive but not answering — terminate it before replacing
49-
daemons::kill_process(existing.pid);
48+
// Process alive but not answering — terminate it and its children before replacing
49+
daemons::kill_daemon_gracefully(existing.pid);
5050
}
5151

5252
daemons::unregister_daemon(&detected_root)?;
@@ -187,7 +187,7 @@ impl DaemonOpenCommand {
187187
= format!("ws://127.0.0.1:{}", port);
188188

189189
// Just attempt to establish a WebSocket connection - if it succeeds, daemon is ready
190-
tokio_tungstenite::connect_async(&url)
190+
let (mut ws, _) = tokio_tungstenite::connect_async(&url)
191191
.await
192192
.map_err(|e| {
193193
Error::DaemonConnectionFailed(Arc::new(std::io::Error::new(
@@ -196,6 +196,9 @@ impl DaemonOpenCommand {
196196
)))
197197
})?;
198198

199+
// Close the connection properly to avoid connection resets
200+
ws.close(None).await.ok();
201+
199202
Ok(())
200203
}
201204
}

‎packages/zpm-switch/src/daemons.rs‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,62 @@ pub fn kill_process(pid: u32) -> bool {
159159
}
160160
}
161161

162+
/// Kill a daemon process and all its children (process group).
163+
/// Sends SIGTERM first, waits for the process to exit, then sends SIGKILL if needed.
164+
/// Returns true if the process was successfully killed.
165+
pub fn kill_daemon_gracefully(pid: u32) -> bool {
166+
#[cfg(unix)]
167+
{
168+
// First, send SIGTERM to the daemon process itself
169+
// The daemon's signal handler should propagate to children
170+
let term_result = unsafe { libc::kill(pid as i32, libc::SIGTERM) };
171+
if term_result != 0 {
172+
// Process doesn't exist or we don't have permission
173+
return false;
174+
}
175+
176+
// Wait up to 6 seconds for the daemon to shut down gracefully
177+
// (daemon waits 5s internally for children to exit, plus 1s buffer)
178+
for _ in 0..60 {
179+
std::thread::sleep(std::time::Duration::from_millis(100));
180+
if !is_process_alive(pid) {
181+
return true;
182+
}
183+
}
184+
185+
// If still alive after 6 seconds, send SIGKILL to the process group
186+
// Use negative PID to target the entire process group
187+
let pgid = unsafe { libc::getpgid(pid as i32) };
188+
if pgid > 0 {
189+
let _ = unsafe { libc::killpg(pgid, libc::SIGKILL) };
190+
}
191+
// Also send SIGKILL directly to the daemon in case it's not the process group leader
192+
let _ = unsafe { libc::kill(pid as i32, libc::SIGKILL) };
193+
194+
// Wait a bit for the process to actually die
195+
for _ in 0..10 {
196+
std::thread::sleep(std::time::Duration::from_millis(100));
197+
if !is_process_alive(pid) {
198+
return true;
199+
}
200+
}
201+
202+
// Return true even if we couldn't verify death - we did our best
203+
true
204+
}
205+
206+
#[cfg(windows)]
207+
{
208+
// On Windows, just use TerminateProcess (no graceful shutdown)
209+
kill_process(pid)
210+
}
211+
212+
#[cfg(not(any(unix, windows)))]
213+
{
214+
false
215+
}
216+
}
217+
162218
pub fn cleanup_stale_daemons() -> Result<(), Error> {
163219
let daemons
164220
= list_daemons()?;

‎packages/zpm/src/commands/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ pub enum YarnCli {
112112
TaskRunInterlaced(tasks::run_interlaced::TaskRunInterlaced),
113113
TaskRunBuffered(tasks::run_buffered::TaskRunBuffered),
114114
TaskRunSilentDependencies(tasks::run_silent_dependencies::TaskRunSilentDependencies),
115+
TaskStats(tasks::stats::TaskStats),
115116
TaskStop(tasks::stop::TaskStop),
116117
Unlink(unlink::Unlink),
117118
Unplug(unplug::Unplug),

‎packages/zpm/src/commands/tasks/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ pub mod push;
66
pub mod run_buffered;
77
pub mod run_interlaced;
88
pub mod run_silent_dependencies;
9+
pub mod stats;
910
pub mod stop;
1011

1112
use zpm_tasks::{parse, TaskName};

‎packages/zpm/src/commands/tasks/push.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,6 @@ impl TaskPush {
6363

6464
client.push_tasks(task_subscriptions, parent_task_id, None, None).await?;
6565

66-
Ok(ExitStatus::from_raw(0 << 8))
66+
Ok(ExitStatus::from_raw(0))
6767
}
6868
}

‎packages/zpm/src/commands/tasks/stats.rs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
use std::{os::unix::process::ExitStatusExt, process::ExitStatus};
2+
3+
use clipanion::cli;
4+
5+
use crate::daemon::DaemonClient;
6+
use crate::error::Error;
7+
use crate::project::Project;
8+
9+
/// Get internal state statistics from the daemon
10+
///
11+
/// This command returns statistics about the daemon's internal state,
12+
/// useful for debugging and testing memory management.
13+
#[cli::command]
14+
#[cli::path("tasks", "stats")]
15+
#[cli::category("Task management commands")]
16+
pub struct TaskStats {
17+
/// Output as JSON
18+
#[cli::option("--json", default = false)]
19+
json: bool,
20+
}
21+
22+
impl TaskStats {
23+
pub async fn execute(&self) -> Result<ExitStatus, Error> {
24+
let project = Project::new(None).await?;
25+
26+
let mut client = DaemonClient::connect(&project.project_cwd).await?;
27+
28+
let stats = client.get_stats().await?;
29+
30+
client.close();
31+
32+
if self.json {
33+
println!(
34+
"{}",
35+
serde_json::json!({
36+
"tasksCount": stats.tasks_count,
37+
"preparedCount": stats.prepared_count,
38+
"subtasksCount": stats.subtasks_count,
39+
"outputBufferCount": stats.output_buffer_count,
40+
"closedTasksCount": stats.closed_tasks_count,
41+
})
42+
);
43+
} else {
44+
println!("Daemon State Statistics:");
45+
println!(" tasks: {}", stats.tasks_count);
46+
println!(" prepared: {}", stats.prepared_count);
47+
println!(" subtasks: {}", stats.subtasks_count);
48+
println!(" output_buffer: {}", stats.output_buffer_count);
49+
println!(" closed_tasks: {}", stats.closed_tasks_count);
50+
}
51+
52+
Ok(ExitStatus::from_raw(0))
53+
}
54+
}

‎packages/zpm/src/daemon/client.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,38 @@ impl DaemonClient {
389389
_ => Err(Error::IpcError("Unexpected response".to_string())),
390390
}
391391
}
392+
393+
/// Get internal state statistics from the daemon (for debugging/testing)
394+
pub async fn get_stats(&mut self) -> Result<DaemonStatsResult, Error> {
395+
let request = DaemonRequest::GetStats;
396+
397+
match self.send_request(request).await? {
398+
DaemonResponse::Stats {
399+
tasks_count,
400+
prepared_count,
401+
subtasks_count,
402+
output_buffer_count,
403+
closed_tasks_count,
404+
} => Ok(DaemonStatsResult {
405+
tasks_count,
406+
prepared_count,
407+
subtasks_count,
408+
output_buffer_count,
409+
closed_tasks_count,
410+
}),
411+
DaemonResponse::Error { message } => Err(Error::IpcError(message)),
412+
_ => Err(Error::IpcError("Unexpected response".to_string())),
413+
}
414+
}
415+
}
416+
417+
/// Result of getting daemon statistics
418+
pub struct DaemonStatsResult {
419+
pub tasks_count: usize,
420+
pub prepared_count: usize,
421+
pub subtasks_count: usize,
422+
pub output_buffer_count: usize,
423+
pub closed_tasks_count: usize,
392424
}
393425

394426
async fn start_daemon(project_root: &Path) -> Result<String, Error> {

‎packages/zpm/src/daemon/coordinator.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use zpm_utils::{Path, ToFileString};
1919

2020
use super::coordinator_commands::{
2121
CancelContextResult, CommandSender, CoordinatorCommand, LongLivedTaskInfo,
22-
PushTasksResult, StopTaskResult, TaskCompletionResult,
22+
PushTasksResult, StatsResult, StopTaskResult, TaskCompletionResult,
2323
};
2424
use super::coordinator_state::{
2525
format_contextual_task_id, parse_base_task_id, CoordinatorState, ContextualTaskId, PreparedTask,
@@ -354,6 +354,17 @@ async fn handle_command(
354354
let _ = response_tx.send(infos);
355355
}
356356

357+
CoordinatorCommand::GetStats { response_tx } => {
358+
let (tasks_count, prepared_count, subtasks_count, output_buffer_count, closed_tasks_count) = state.get_stats();
359+
let _ = response_tx.send(StatsResult {
360+
tasks_count,
361+
prepared_count,
362+
subtasks_count,
363+
output_buffer_count,
364+
closed_tasks_count,
365+
});
366+
}
367+
357368
// ====================================================================
358369
// Subscription Commands
359370
// ====================================================================

0 commit comments

Comments
 (0)