diff --git a/worker/internal/worker/worker.go b/worker/internal/worker/worker.go index cd17c90..ce0ad13 100644 --- a/worker/internal/worker/worker.go +++ b/worker/internal/worker/worker.go @@ -177,17 +177,12 @@ func Run(ctx context.Context, config Config, payload Payload, publisher Publishe defer close(watcherDone) select { case <-workContext.Done(): - _ = syscall.Kill(-command.Process.Pid, syscall.SIGTERM) - timer := time.NewTimer(config.GracePeriod) - defer timer.Stop() - select { - case <-timer.C: - _ = syscall.Kill(-command.Process.Pid, syscall.SIGKILL) - case <-processDone: - // Descendants may have survived the shell's termination. - _ = syscall.Kill(-command.Process.Pid, syscall.SIGKILL) - } + terminateProcessGroup(command.Process.Pid, config.GracePeriod) case <-processDone: + if workContext.Err() != nil { + terminateProcessGroup(command.Process.Pid, config.GracePeriod) + return + } _ = syscall.Kill(-command.Process.Pid, syscall.SIGKILL) } }() @@ -234,6 +229,27 @@ func Run(ctx context.Context, config Config, payload Payload, publisher Publishe return code, errors.Join(runError, outputError) } +// terminateProcessGroup gives surviving descendants their grace period even when +// the leader has exited. Return early only when the entire process group is gone. +func terminateProcessGroup(pid int, grace time.Duration) { + _ = syscall.Kill(-pid, syscall.SIGTERM) + deadline := time.NewTimer(grace) + defer deadline.Stop() + poll := time.NewTicker(10 * time.Millisecond) + defer poll.Stop() + for { + if errors.Is(syscall.Kill(-pid, 0), syscall.ESRCH) { + return + } + select { + case <-deadline.C: + _ = syscall.Kill(-pid, syscall.SIGKILL) + return + case <-poll.C: + } + } +} + // resolveHome uses the node's NSS lookup for site accounts (LDAP/SSSD). // A static Go binary cannot rely on user.Current or libc's user lookup. func resolveHome() (string, error) { diff --git a/worker/internal/worker/worker_test.go b/worker/internal/worker/worker_test.go index 46b61d4..2aef980 100644 --- a/worker/internal/worker/worker_test.go +++ b/worker/internal/worker/worker_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" "os" "path/filepath" @@ -355,3 +356,45 @@ func TestUnicodeAcrossChunkBoundariesSurvivesJSON(t *testing.T) { } } } + +func TestCancellationLetsRedirectedChildFinishCleanup(t *testing.T) { + directory := t.TempDir() + marker := filepath.Join(directory, "flushed") + ready := filepath.Join(directory, "ready") + script := fmt.Sprintf(`trap 'exit 0' TERM +/bin/bash -c 'trap "sleep 0.2; touch %s; exit" TERM; touch %s; while :; do sleep 0.01; done' >/dev/null 2>&1 & +wait +`, marker, ready) + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + watcherDone := make(chan struct{}) + go func() { + defer close(watcherDone) + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + if _, err := os.Stat(ready); err == nil { + cancel() + return + } + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } + }() + var out, errout bytes.Buffer + cfg := config(&out, &errout) + cfg.GracePeriod = time.Second + publisher := &fakePublisher{} + code, err := Run(ctx, cfg, Payload{Script: script}, publisher) + cancel() + <-watcherDone + if code != 143 || lastStatus(t, publisher) != "cancelled" { + t.Fatalf("code=%d err=%v", code, err) + } + if _, err := os.Stat(marker); err != nil { + t.Fatalf("child cleanup was interrupted before the grace period elapsed: %v", err) + } +}