Skip to content

Commit

Permalink
Tool to implement robust cleanups even on process crash
Browse files Browse the repository at this point in the history
  • Loading branch information
kovidgoyal committed Jan 5, 2025
1 parent d23adce commit 0d5bcff
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 0 deletions.
72 changes: 72 additions & 0 deletions kitty_tests/atexit.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
#!/usr/bin/env python
# License: GPLv3 Copyright: 2025, Kovid Goyal <kovid at kovidgoyal.net>


import os
import select
import shutil
import signal
import subprocess
import tempfile

from kitty.constants import kitten_exe, kitty_exe

from . import BaseTest


class Atexit(BaseTest):

def setUp(self):
self.tdir = tempfile.mkdtemp()

def tearDown(self):
shutil.rmtree(self.tdir)

def test_atexit(self):

def r(action='close'):
p = subprocess.Popen([kitty_exe(), '+runpy', f'''\
import subprocess
p = subprocess.Popen(['{kitten_exe()}', '__atexit__'])
print(p.pid, flush=True)
raise SystemExit(p.wait())
'''], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
readers = [p.stdout.fileno()]
def read():
r, _, _ = select.select(readers, [], [], 10)
if not r:
raise TimeoutError('Timed out waiting for read from child')
return p.stdout.readline().rstrip().decode()
atexit_pid = int(read())
for i in range(2):
with open(os.path.join(self.tdir, str(i)), 'w') as f:
p.stdin.write(f'unlink {f.name}\n'.encode())
p.stdin.flush()
select.select(readers, [], [], 10)
self.ae(read(), str(i+1))
self.assertTrue(os.listdir(self.tdir))

# Ensure child is ignoring signals
os.kill(atexit_pid, signal.SIGINT)
os.kill(atexit_pid, signal.SIGTERM)
if action == 'close':
p.stdin.close()
elif action == 'terminate':
p.terminate()
else:
p.kill()
p.wait(10)
if action != 'close':
p.stdin.close()
select.select(readers, [], [], 10)
self.assertFalse(read())
p.stdout.close()
self.assertFalse(os.listdir(self.tdir))
try:
os.waitpid(atexit_pid, 0)
except ChildProcessError:
pass

r('close')
r('terminate')
r('kill')
65 changes: 65 additions & 0 deletions tools/cmd/atexit/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
package atexit

import (
"bufio"
"fmt"
"os"
"os/signal"
"strings"

"kitty/tools/cli"
)

var _ = fmt.Print

func main() (rc int, err error) {
signal.Ignore()
done_channel := make(chan bool)
lines := []string{}

defer os.Stdout.Close()

go func() {
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
lines = append(lines, scanner.Text())
fmt.Println(len(lines))
}
done_channel <- true
}()

keep_going := true
for keep_going {
select {
case <-done_channel:
keep_going = false
}
}
rc = 0
for _, line := range lines {
if action, rest, found := strings.Cut(line, " "); found {
switch action {
case "unlink":
if err := os.Remove(rest); err != nil {
fmt.Fprintln(os.Stderr, "Failed to remove:", rest, "with error:", err)
rc = 1
}
}
}
}
return
}

func EntryPoint(root *cli.Command) {
root.AddSubCommand(&cli.Command{
Name: "__atexit__",
Hidden: true,
OnlyArgsAllowed: true,
Run: func(cmd *cli.Command, args []string) (rc int, err error) {
if len(args) != 0 {
return 1, fmt.Errorf("Usage: __atexit__")
}
return main()
},
})
}
3 changes: 3 additions & 0 deletions tools/cmd/tool/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"kitty/kittens/unicode_input"
"kitty/tools/cli"
"kitty/tools/cmd/at"
"kitty/tools/cmd/atexit"
"kitty/tools/cmd/benchmark"
"kitty/tools/cmd/edit_in_kitty"
"kitty/tools/cmd/mouse_demo"
Expand Down Expand Up @@ -109,6 +110,8 @@ func KittyToolEntryPoints(root *cli.Command) {
})
// __convert_image__
images.ConvertEntryPoint(root)
// __atexit__
atexit.EntryPoint(root)
// __generate_man_pages__
root.AddSubCommand(&cli.Command{
Name: "__generate_man_pages__",
Expand Down

0 comments on commit 0d5bcff

Please sign in to comment.