| 1 | package proc |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "log/slog" |
| 6 | "os/exec" |
| 7 | "sync" |
| 8 | "time" |
| 9 | ) |
| 10 | |
| 11 | const ( |
| 12 | defaultCancelRetryInterval = 500 * time.Millisecond |
| 13 | defaultCancelRetryFor = 5 * time.Second |
| 14 | ) |
| 15 | |
| 16 | // RunOptions controls process-tree tracking for a foreground command. |
| 17 | type RunOptions struct { |
| 18 | Track bool |
| 19 | CancelWaitGrace time.Duration |
| 20 | CancelRetryInterval time.Duration |
| 21 | CancelRetryFor time.Duration |
| 22 | Source string |
| 23 | ShellKind string |
| 24 | ShellPath string |
| 25 | CommandPreview string |
| 26 | } |
| 27 | |
| 28 | // RunDiagnostics is local-only state for debugging stuck cancellation. |
| 29 | type RunDiagnostics struct { |
| 30 | Source string |
| 31 | ShellKind string |
| 32 | ShellPath string |
| 33 | CommandPreview string |
| 34 | RootPID int |
| 35 | Tracked bool |
| 36 | JobObjectCreated bool |
| 37 | TreeTrackerStarted bool |
| 38 | KillCalls int |
| 39 | TreeKillAttempts int |
| 40 | RetryKillCalls int |
| 41 | CancelWaitGraceExpired bool |
| 42 | CancelRetryWindowMillis int64 |
| 43 | } |
| 44 | |
| 45 | // TrackedCommand owns the cancellation state for one started command. |
| 46 | type TrackedCommand struct { |
| 47 | cmd *exec.Cmd |
| 48 | |
| 49 | mu sync.Mutex |
| 50 | job uintptr |
| 51 | tree *TreeTracker |
| 52 | killed bool |
| 53 | diag RunDiagnostics |
| 54 | } |
| 55 | |
| 56 | // RunCommand starts cmd and waits for it. When tracking is enabled, cancellation |
| 57 | // kills the tracked process tree and returns after a bounded wait even if the |
| 58 | // platform wait path remains wedged. |
| 59 | func RunCommand(ctx context.Context, cmd *exec.Cmd, opts RunOptions) (*TrackedCommand, error) { |
| 60 | opts = normalizeRunOptions(opts) |
| 61 | if !opts.Track { |
| 62 | SetCancelKillsTree(cmd) |
| 63 | return nil, cmd.Run() |
| 64 | } |
| 65 | |
| 66 | tracked := &TrackedCommand{cmd: cmd} |
| 67 | tracked.setMetadata(opts) |
| 68 | HideWindow(cmd) |
| 69 | cmd.Cancel = func() error { |
| 70 | tracked.Kill() |
| 71 | return context.Canceled |
| 72 | } |
| 73 | job, err := StartTracked(cmd) |
| 74 | if err != nil { |
| 75 | return tracked, err |
| 76 | } |
| 77 | tracked.setStarted(job) |
| 78 | tracked.setTree(TrackTree(cmd)) |
| 79 | return tracked, waitForTrackedCommand(ctx, tracked, cmd.Wait, opts.CancelWaitGrace, opts.CancelRetryInterval, opts.CancelRetryFor) |
| 80 | } |
| 81 | |
| 82 | // SetCancelKillsTree configures cmd so context cancellation kills the entire |
| 83 | // process tree instead of only the direct child. |
| 84 | func SetCancelKillsTree(cmd *exec.Cmd) { |
| 85 | if cmd == nil { |
| 86 | return |
| 87 | } |
| 88 | HideWindow(cmd) |
| 89 | SetProcessGroupKill(cmd) |
| 90 | cmd.Cancel = func() error { |
| 91 | KillTree(cmd) |
| 92 | return context.Canceled |
| 93 | } |
| 94 | } |
| 95 | |
| 96 | // CanceledWaitError preserves the underlying wait error while presenting the |
| 97 | // context cancellation as the primary command result. |
| 98 | type CanceledWaitError struct { |
| 99 | Cause error |
| 100 | WaitErr error |
| 101 | } |
| 102 | |
| 103 | func (e CanceledWaitError) Error() string { |
| 104 | if e.Cause != nil { |
| 105 | return e.Cause.Error() |
| 106 | } |
| 107 | if e.WaitErr != nil { |
| 108 | return e.WaitErr.Error() |
| 109 | } |
| 110 | return "command wait canceled" |
| 111 | } |
| 112 | |
| 113 | func (e CanceledWaitError) Unwrap() []error { |
| 114 | if e.Cause != nil && e.WaitErr != nil { |
| 115 | return []error{e.Cause, e.WaitErr} |
| 116 | } |
| 117 | if e.Cause != nil { |
| 118 | return []error{e.Cause} |
| 119 | } |
| 120 | if e.WaitErr != nil { |
| 121 | return []error{e.WaitErr} |
| 122 | } |
| 123 | return nil |
| 124 | } |
| 125 | |
| 126 | func normalizeRunOptions(opts RunOptions) RunOptions { |
| 127 | if opts.CancelWaitGrace <= 0 { |
| 128 | opts.CancelWaitGrace = time.Second |
| 129 | } |
| 130 | if opts.CancelRetryInterval <= 0 { |
| 131 | opts.CancelRetryInterval = defaultCancelRetryInterval |
| 132 | } |
| 133 | if opts.CancelRetryFor <= 0 { |
| 134 | opts.CancelRetryFor = defaultCancelRetryFor |
| 135 | } |
| 136 | return opts |
| 137 | } |
| 138 | |
| 139 | func waitForTrackedCommand(ctx context.Context, tracked *TrackedCommand, wait func() error, grace, retryEvery, retryFor time.Duration) error { |
| 140 | waitCh := make(chan error, 1) |
| 141 | go func() { waitCh <- wait() }() |
| 142 | |
| 143 | select { |
| 144 | case err := <-waitCh: |
| 145 | tracked.StopTracking() |
| 146 | return err |
| 147 | case <-ctx.Done(): |
| 148 | } |
| 149 | |
| 150 | tracked.Kill() |
| 151 | select { |
| 152 | case err := <-waitCh: |
| 153 | tracked.StopTracking() |
| 154 | return CanceledWaitError{Cause: context.Cause(ctx), WaitErr: err} |
| 155 | case <-time.After(grace): |
| 156 | tracked.markGraceExpired(retryFor) |
| 157 | diag := tracked.Diagnostics() |
| 158 | slog.Warn("proc: command wait still blocked after cancellation", |
| 159 | "source", diag.Source, |
| 160 | "shell_kind", diag.ShellKind, |
| 161 | "shell_path", diag.ShellPath, |
| 162 | "command_preview", diag.CommandPreview, |
| 163 | "root_pid", diag.RootPID, |
| 164 | "tracked", diag.Tracked, |
| 165 | "job_object_created", diag.JobObjectCreated, |
| 166 | "tree_tracker_started", diag.TreeTrackerStarted, |
| 167 | "retry_window_ms", diag.CancelRetryWindowMillis) |
| 168 | go tracked.retryKillUntilWait(waitCh, retryEvery, retryFor) |
| 169 | return context.Cause(ctx) |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | // Kill terminates the tracked command tree. |
| 174 | func (p *TrackedCommand) Kill() { |
| 175 | if p == nil { |
| 176 | return |
| 177 | } |
| 178 | p.mu.Lock() |
| 179 | firstKill := !p.killed |
| 180 | p.killed = true |
| 181 | p.diag.KillCalls++ |
| 182 | job := p.job |
| 183 | p.job = 0 |
| 184 | tree := p.tree |
| 185 | p.mu.Unlock() |
| 186 | if !firstKill { |
| 187 | job = 0 |
| 188 | } |
| 189 | KillTracked(p.cmd, job) |
| 190 | if tree != nil { |
| 191 | treeKills := tree.Kill() |
| 192 | tree.Stop() |
| 193 | p.mu.Lock() |
| 194 | p.diag.TreeKillAttempts += treeKills |
| 195 | p.mu.Unlock() |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | // StopTracking stops background tree observation without killing the command. |
| 200 | func (p *TrackedCommand) StopTracking() { |
| 201 | if p == nil { |
| 202 | return |
| 203 | } |
| 204 | p.mu.Lock() |
| 205 | tree := p.tree |
| 206 | p.tree = nil |
| 207 | p.mu.Unlock() |
| 208 | if tree != nil { |
| 209 | tree.Stop() |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | // Diagnostics returns a snapshot of local cancellation state. |
| 214 | func (p *TrackedCommand) Diagnostics() RunDiagnostics { |
| 215 | if p == nil { |
| 216 | return RunDiagnostics{} |
| 217 | } |
| 218 | p.mu.Lock() |
| 219 | defer p.mu.Unlock() |
| 220 | return p.diag |
| 221 | } |
| 222 | |
| 223 | func (p *TrackedCommand) setStarted(job uintptr) { |
| 224 | if p == nil { |
| 225 | return |
| 226 | } |
| 227 | rootPID := 0 |
| 228 | if p.cmd != nil && p.cmd.Process != nil { |
| 229 | rootPID = p.cmd.Process.Pid |
| 230 | } |
| 231 | p.mu.Lock() |
| 232 | killed := p.killed |
| 233 | if !killed { |
| 234 | p.job = job |
| 235 | } |
| 236 | p.diag.RootPID = rootPID |
| 237 | p.diag.Tracked = true |
| 238 | p.diag.JobObjectCreated = job != 0 |
| 239 | p.mu.Unlock() |
| 240 | if killed && job != 0 { |
| 241 | KillTracked(p.cmd, job) |
| 242 | } |
| 243 | } |
| 244 | |
| 245 | func (p *TrackedCommand) setTree(tree *TreeTracker) { |
| 246 | if p == nil || tree == nil { |
| 247 | return |
| 248 | } |
| 249 | p.mu.Lock() |
| 250 | killed := p.killed |
| 251 | if !killed { |
| 252 | p.tree = tree |
| 253 | p.diag.TreeTrackerStarted = true |
| 254 | } |
| 255 | p.mu.Unlock() |
| 256 | if killed { |
| 257 | tree.Kill() |
| 258 | tree.Stop() |
| 259 | } |
| 260 | } |
| 261 | |
| 262 | func (p *TrackedCommand) markGraceExpired(retryFor time.Duration) { |
| 263 | if p == nil { |
| 264 | return |
| 265 | } |
| 266 | p.mu.Lock() |
| 267 | p.diag.CancelWaitGraceExpired = true |
| 268 | p.diag.CancelRetryWindowMillis = retryFor.Milliseconds() |
| 269 | p.mu.Unlock() |
| 270 | } |
| 271 | |
| 272 | func (p *TrackedCommand) retryKillUntilWait(waitCh <-chan error, interval, max time.Duration) { |
| 273 | defer p.StopTracking() |
| 274 | deadline := time.NewTimer(max) |
| 275 | defer deadline.Stop() |
| 276 | ticker := time.NewTicker(interval) |
| 277 | defer ticker.Stop() |
| 278 | for { |
| 279 | select { |
| 280 | case <-waitCh: |
| 281 | return |
| 282 | case <-ticker.C: |
| 283 | p.mu.Lock() |
| 284 | p.diag.RetryKillCalls++ |
| 285 | p.mu.Unlock() |
| 286 | p.Kill() |
| 287 | case <-deadline.C: |
| 288 | diag := p.Diagnostics() |
| 289 | if diag.CancelWaitGraceExpired { |
| 290 | slog.Warn("proc: command cleanup retry window ended", |
| 291 | "source", diag.Source, |
| 292 | "shell_kind", diag.ShellKind, |
| 293 | "shell_path", diag.ShellPath, |
| 294 | "command_preview", diag.CommandPreview, |
| 295 | "root_pid", diag.RootPID, |
| 296 | "kill_calls", diag.KillCalls, |
| 297 | "retry_kill_calls", diag.RetryKillCalls) |
| 298 | } |
| 299 | return |
| 300 | } |
| 301 | } |
| 302 | } |
| 303 | |
| 304 | func (p *TrackedCommand) setMetadata(opts RunOptions) { |
| 305 | if p == nil { |
| 306 | return |
| 307 | } |
| 308 | p.mu.Lock() |
| 309 | p.diag.Source = opts.Source |
| 310 | p.diag.ShellKind = opts.ShellKind |
| 311 | p.diag.ShellPath = opts.ShellPath |
| 312 | p.diag.CommandPreview = opts.CommandPreview |
| 313 | p.mu.Unlock() |
| 314 | } |
| 315 |