diff --git a/shutdown_test.go b/shutdown_test.go new file mode 100644 index 00000000..4c7ee9d4 --- /dev/null +++ b/shutdown_test.go @@ -0,0 +1,67 @@ +package main + +import ( + "os/exec" + "runtime" + "testing" + "time" +) + +func TestShutdownEmptyVHS(t *testing.T) { + // shutdown should be a no-op (and never panic) on a fresh VHS that + // hasn't been Start()ed: both tty and browser are nil. + var vhs VHS + if err := vhs.shutdown(); err != nil { + t.Fatalf("shutdown on empty VHS: %v", err) + } +} + +func TestShutdownKillsTTY(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sleep binary not available on Windows runners") + } + + // Use a long-running sleep as a stand-in for ttyd so we can check that + // shutdown actually killed it. Regression for the ttyd leak when + // Evaluate returns before Record begins. + cmd := exec.Command("sleep", "30") + if err := cmd.Start(); err != nil { + t.Fatalf("start sleep: %v", err) + } + + vhs := VHS{tty: cmd} + if err := vhs.shutdown(); err != nil { + t.Fatalf("shutdown: %v", err) + } + + done := make(chan error, 1) + go func() { done <- cmd.Wait() }() + select { + case <-done: + // The Wait return value is intentionally ignored: the process was + // killed, so an error is expected. + case <-time.After(2 * time.Second): + t.Fatal("ttyd stand-in still alive 2s after shutdown") + } +} + +func TestShutdownIsIdempotent(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sleep binary not available on Windows runners") + } + + // terminate() already runs at the end of the happy-path recording, then + // the deferred close() runs from Evaluate. Calling shutdown twice in a + // row should not surface an error from the second call. + cmd := exec.Command("sleep", "30") + if err := cmd.Start(); err != nil { + t.Fatalf("start sleep: %v", err) + } + vhs := VHS{tty: cmd} + if err := vhs.shutdown(); err != nil { + t.Fatalf("first shutdown: %v", err) + } + if err := vhs.shutdown(); err != nil { + t.Fatalf("second shutdown: %v", err) + } +} diff --git a/vhs.go b/vhs.go index 235e7b42..5c9f3003 100644 --- a/vhs.go +++ b/vhs.go @@ -136,6 +136,15 @@ func (vhs *VHS) Start() error { return fmt.Errorf("could not start tty: %w", err) } + // If anything below this point fails, the ttyd we just spawned would + // outlive the program. Kill it on the way out unless Start completes. + cleanupTTY := true + defer func() { + if cleanupTTY { + _ = vhs.tty.Process.Kill() + } + }() + path, _ := launcher.LookPath() enableNoSandbox := os.Getenv("VHS_NO_SANDBOX") != "" u, err := launcher.New().Leakless(false).Bin(path).NoSandbox(enableNoSandbox).Launch() @@ -145,16 +154,36 @@ func (vhs *VHS) Start() error { browser := rod.New().ControlURL(u).MustConnect() page, err := browser.Page(proto.TargetCreateTarget{URL: fmt.Sprintf("http://localhost:%d", port)}) if err != nil { + _ = browser.Close() return fmt.Errorf("could not open ttyd: %w", err) } vhs.browser = browser vhs.Page = page - vhs.close = vhs.browser.Close + vhs.close = vhs.shutdown vhs.started = true + cleanupTTY = false return nil } +// shutdown stops the browser and ttyd that this VHS instance owns. It's +// safe to call multiple times and after terminate() has already run; +// already-closed browsers and already-exited processes are not treated as +// errors so callers can use this from deferred cleanups without worrying +// about ordering. +func (vhs *VHS) shutdown() error { + var ttyErr error + if vhs.tty != nil && vhs.tty.Process != nil { + if err := vhs.tty.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) { + ttyErr = err + } + } + if vhs.browser != nil { + _ = vhs.browser.Close() + } + return ttyErr +} + // Setup sets up the VHS instance and performs the necessary actions to reflect // the options that are default and set by the user. func (vhs *VHS) Setup() {