main.go

Download main.go · view raw

1// stayup: a nohup replacement backed by systemd.
2//
3// Every command becomes a real systemd service unit (Restart=always), so it
4// survives crashes, OOM kills and host reboots. stayup keeps no state of its
5// own: the unit files under systemd's unit directory are the only record.
6//
7//   - Non-root: ~/.config/systemd/user/stayup-<cmd>-<hash>.service via
8//     `systemctl --user`; `loginctl enable-linger` makes it start at boot.
9//   - Root: /etc/systemd/system/, WantedBy=multi-user.target.
10//   - The unit captures cwd, the resolved executable and the caller's
11//     environment (file mode 0600). Logs go to the journal unless -o/--stderr
12//     are given (StandardOutput=append:, systemd >= 240).
13//   - Same command + same cwd = same unit, so starting twice is idempotent.
14//   - Users see "processes"; unit names only surface as NAME in -l.
15package main
16
17import (
18	"bufio"
19	"crypto/sha256"
20	"encoding/hex"
21	"flag"
22	"fmt"
23	"os"
24	"os/exec"
25	"os/signal"
26	"path/filepath"
27	"regexp"
28	"strconv"
29	"strings"
30	"syscall"
31	"text/tabwriter"
32	"time"
33	"unsafe"
34)
35
36const prefix = "stayup-"
37
38var (
39	user    = os.Geteuid() != 0 // non-root: per-user systemd instance
40	unitDir string
41)
42
43func main() {
44	var (
45		fg      bool
46		attach  bool
47		list    bool
48		kill    bool
49		outFile string
50		errFile string
51	)
52	flag.BoolVar(&fg, "f", false, "")
53	flag.BoolVar(&fg, "foreground", false, "follow the log; ctrl-c kills, ctrl-d detaches")
54	flag.BoolVar(&attach, "a", false, "")
55	flag.BoolVar(&attach, "attach", false, "attach to a running process (by command, PID or name)")
56	flag.BoolVar(&list, "l", false, "")
57	flag.BoolVar(&list, "list", false, "list processes")
58	flag.BoolVar(&kill, "K", false, "")
59	flag.BoolVar(&kill, "kill", false, "kill a process (by command, PID or name)")
60	flag.StringVar(&outFile, "o", "", "")
61	flag.StringVar(&outFile, "stdout", "", "append stdout to this file (default: system log)")
62	flag.StringVar(&errFile, "stderr", "", "append stderr to this file (default: same as stdout)")
63	flag.Usage = usage
64	flag.Parse()
65	args := flag.Args()
66
67	if user {
68		cfg := os.Getenv("XDG_CONFIG_HOME")
69		if cfg == "" {
70			cfg = filepath.Join(os.Getenv("HOME"), ".config")
71		}
72		unitDir = filepath.Join(cfg, "systemd", "user")
73	} else {
74		unitDir = "/etc/systemd/system"
75	}
76
77	switch {
78	case list:
79		fail(doList())
80	case kill:
81		if len(args) == 0 {
82			fail(fmt.Errorf("--kill needs a command or PID"))
83		}
84		u, err := find(args)
85		fail(err)
86		c := command(u)
87		fail(remove(u))
88		fmt.Fprintln(os.Stderr, "killed", c)
89	case attach:
90		if len(args) == 0 {
91			fail(fmt.Errorf("--attach needs a command or PID"))
92		}
93		u, err := find(args)
94		fail(err)
95		fail(foreground(u))
96	case len(args) == 0:
97		usage()
98		os.Exit(2)
99	default:
100		if errFile == "" {
101			errFile = outFile
102		}
103		u, fresh, err := start(args, outFile, errFile)
104		fail(err)
105		if !fg {
106			if fresh {
107				fmt.Fprintf(os.Stderr, "started %s (pid %s)\n", shJoin(args), show(u, "MainPID"))
108			} else {
109				fmt.Fprintf(os.Stderr, "already running: %s (pid %s)\n", shJoin(args), show(u, "MainPID"))
110			}
111			fmt.Fprintf(os.Stderr, "attach: stayup -a %s\nkill:   stayup -K %s\n", shJoin(args), shJoin(args))
112			return
113		}
114		fail(foreground(u))
115	}
116}
117
118func usage() {
119	fmt.Fprint(os.Stderr, `stayup - keep a command running via systemd (survives crashes, OOM, reboots)
120
121  stayup [-f] [-o FILE] [--stderr FILE] [--] COMMAND [ARGS...]
122  stayup -a | --attach COMMAND [ARGS...] | PID | NAME
123  stayup -l | --list
124  stayup -K | --kill COMMAND [ARGS...] | PID | NAME
125
126  -f, --foreground   run attached: ctrl-c kills, ctrl-d or hangup detaches
127  -a, --attach       attach to a running process (same keys as -f)
128  -o, --stdout FILE  append stdout to FILE (default: system log)
129      --stderr FILE  append stderr to FILE (default: same as stdout)
130  -l, --list         list processes
131  -K, --kill         kill a process (by command, PID or NAME from -l)
132`)
133}
134
135func fail(err error) {
136	if err != nil {
137		fmt.Fprintln(os.Stderr, "stayup:", err)
138		os.Exit(1)
139	}
140}
141
142// ---- systemctl plumbing ---------------------------------------------------
143
144func sysctl(args ...string) *exec.Cmd {
145	if user {
146		args = append([]string{"--user"}, args...)
147	}
148	c := exec.Command("systemctl", args...)
149	c.Stderr = os.Stderr
150	return c
151}
152
153func run(args ...string) error {
154	if err := sysctl(args...).Run(); err != nil {
155		return fmt.Errorf("systemctl %s: %w", strings.Join(args, " "), err)
156	}
157	return nil
158}
159
160// command returns the user's command line for a unit.
161func command(unit string) string {
162	return strings.TrimPrefix(show(unit, "Description"), "stayup: ")
163}
164
165// show returns one property of a unit, or "" when it is unknown.
166func show(unit, prop string) string {
167	out, _ := sysctl("show", "-p", prop, "--value", unit).Output()
168	return strings.TrimSpace(string(out))
169}
170
171func props(unit string, names ...string) map[string]string {
172	m := map[string]string{}
173	out, _ := sysctl("show", "-p", strings.Join(names, ","), unit).Output()
174	for _, l := range strings.Split(string(out), "\n") {
175		if k, v, ok := strings.Cut(l, "="); ok {
176			m[k] = v
177		}
178	}
179	return m
180}
181
182// ---- unit naming -----------------------------------------------------------
183
184var badChars = regexp.MustCompile(`[^a-zA-Z0-9_.]+`)
185
186func unitName(cwd string, args []string) string {
187	h := sha256.Sum256([]byte(cwd + "\x00" + strings.Join(args, "\x00")))
188	slug := badChars.ReplaceAllString(filepath.Base(args[0]), "-")
189	if len(slug) > 24 {
190		slug = slug[:24]
191	}
192	return prefix + slug + "-" + hex.EncodeToString(h[:])[:8] + ".service"
193}
194
195func shJoin(args []string) string {
196	q := make([]string, len(args))
197	for i, a := range args {
198		if a == "" || strings.ContainsAny(a, " \t\n'\"\\$&|;<>()*?[]#~") {
199			a = "'" + strings.ReplaceAll(a, "'", `'\''`) + "'"
200		}
201		q[i] = a
202	}
203	return strings.Join(q, " ")
204}
205
206// unitQuote quotes one word for ExecStart=/Environment= in a unit file.
207func unitQuote(s string) string {
208	r := strings.NewReplacer(`\`, `\\`, `"`, `\"`, "%", "%%", "\n", `\n`)
209	return `"` + r.Replace(s) + `"`
210}
211
212func allUnits() ([]string, error) {
213	ents, err := os.ReadDir(unitDir)
214	if os.IsNotExist(err) {
215		return nil, nil
216	}
217	if err != nil {
218		return nil, err
219	}
220	var us []string
221	for _, e := range ents {
222		if n := e.Name(); strings.HasPrefix(n, prefix) && strings.HasSuffix(n, ".service") {
223			us = append(us, n)
224		}
225	}
226	return us, nil
227}
228
229// find resolves a --kill/--attach target: a unit name, a PID, the same command
230// in this cwd, or the same command in any cwd when unambiguous.
231func find(args []string) (string, error) {
232	units, err := allUnits()
233	if err != nil {
234		return "", err
235	}
236	if len(args) == 1 {
237		name := args[0]
238		if !strings.HasSuffix(name, ".service") {
239			name += ".service"
240		}
241		for _, u := range units {
242			if u == name {
243				return u, nil
244			}
245		}
246		if pid, err := strconv.Atoi(args[0]); err == nil {
247			for _, u := range units {
248				if show(u, "MainPID") == strconv.Itoa(pid) {
249					return u, nil
250				}
251			}
252			return "", fmt.Errorf("no process with pid %d (see stayup -l)", pid)
253		}
254	}
255	cwd, _ := os.Getwd()
256	want := unitName(cwd, args)
257	var matches []string
258	for _, u := range units {
259		if u == want {
260			return u, nil
261		}
262		if show(u, "Description") == "stayup: "+shJoin(args) {
263			matches = append(matches, u)
264		}
265	}
266	switch len(matches) {
267	case 0:
268		return "", fmt.Errorf("no process for %q (see stayup -l)", shJoin(args))
269	case 1:
270		return matches[0], nil
271	}
272	return "", fmt.Errorf("%d processes run %q in other directories; use the PID or NAME from stayup -l", len(matches), shJoin(args))
273}
274
275// ---- actions ---------------------------------------------------------------
276
277func start(args []string, outFile, errFile string) (unit string, fresh bool, err error) {
278	cwd, err := os.Getwd()
279	if err != nil {
280		return "", false, err
281	}
282	unit = unitName(cwd, args)
283	path := filepath.Join(unitDir, unit)
284	if _, err := os.Stat(path); err == nil {
285		if show(unit, "ActiveState") != "active" {
286			err = run("start", unit)
287		}
288		return unit, false, err
289	}
290
291	exe, err := exec.LookPath(args[0])
292	if err != nil {
293		return "", false, err
294	}
295	if exe, err = filepath.Abs(exe); err != nil {
296		return "", false, err
297	}
298	abs := func(f string) string {
299		if f == "" || filepath.IsAbs(f) {
300			return f
301		}
302		return filepath.Join(cwd, f)
303	}
304	outFile, errFile = abs(outFile), abs(errFile)
305
306	var b strings.Builder
307	fmt.Fprintf(&b, "# generated by stayup on %s\n", time.Now().Format(time.RFC3339))
308	fmt.Fprintf(&b, "[Unit]\nDescription=stayup: %s\nStartLimitIntervalSec=0\n\n", shJoin(args))
309	fmt.Fprintf(&b, "[Service]\nType=simple\nWorkingDirectory=%s\n", strings.ReplaceAll(cwd, "%", "%%"))
310	words := []string{unitQuote(exe)}
311	for _, a := range args[1:] {
312		words = append(words, unitQuote(a))
313	}
314	fmt.Fprintf(&b, "ExecStart=%s\n", strings.Join(words, " "))
315	fmt.Fprint(&b, "Restart=always\nRestartSec=1\nTimeoutStopSec=15\nOOMPolicy=continue\n")
316	if outFile != "" {
317		fmt.Fprintf(&b, "StandardOutput=append:%s\n", outFile)
318	}
319	if errFile != "" {
320		fmt.Fprintf(&b, "StandardError=append:%s\n", errFile)
321	}
322	for _, kv := range os.Environ() {
323		if k, _, _ := strings.Cut(kv, "="); k != "" && k != "_" && !strings.HasPrefix(k, "BASH_FUNC_") {
324			fmt.Fprintf(&b, "Environment=%s\n", unitQuote(kv))
325		}
326	}
327	target := "multi-user.target"
328	if user {
329		target = "default.target"
330	}
331	fmt.Fprintf(&b, "\n[Install]\nWantedBy=%s\n", target)
332
333	if err := os.MkdirAll(unitDir, 0o755); err != nil {
334		return "", false, err
335	}
336	if err := os.WriteFile(path, []byte(b.String()), 0o600); err != nil {
337		return "", false, err
338	}
339	if user {
340		// Let the user's systemd instance (and our service) outlive logins and come back after reboot.
341		exec.Command("loginctl", "enable-linger").Run()
342	}
343	if err := run("daemon-reload"); err != nil {
344		return "", false, err
345	}
346	if err := run("enable", "--now", "--quiet", unit); err != nil {
347		sysctl("disable", unit).Run()
348		os.Remove(path)
349		run("daemon-reload")
350		return "", false, fmt.Errorf("%w\nunit file was:\n%s", err, b.String())
351	}
352	return unit, true, nil
353}
354
355func remove(unit string) error {
356	err := run("disable", "--now", "--quiet", unit)
357	q := sysctl("reset-failed", unit)
358	q.Stderr = nil
359	q.Run()
360	if rerr := os.Remove(filepath.Join(unitDir, unit)); rerr != nil && !os.IsNotExist(rerr) && err == nil {
361		err = rerr
362	}
363	run("daemon-reload")
364	return err
365}
366
367func doList() error {
368	units, err := allUnits()
369	if err != nil {
370		return err
371	}
372	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', 0)
373	fmt.Fprintln(w, "NAME\tPID\tSTATE\tSTARTED\tLAST RESTART\tRESTARTS\tDIR\tCOMMAND")
374	for _, u := range units {
375		p := props(u, "MainPID", "ActiveState", "SubState", "NRestarts", "ExecMainStartTimestamp", "Description", "WorkingDirectory")
376		started := "?"
377		if fi, err := os.Stat(filepath.Join(unitDir, u)); err == nil {
378			started = fi.ModTime().Format("2006-01-02 15:04")
379		}
380		pid := p["MainPID"]
381		if pid == "0" || pid == "" {
382			pid = "-"
383		}
384		fmt.Fprintf(w, "%s\t%s\t%s/%s\t%s\t%s\t%s\t%s\t%s\n",
385			strings.TrimSuffix(u, ".service"), pid, p["ActiveState"], p["SubState"], started, fmtTS(p["ExecMainStartTimestamp"]),
386			p["NRestarts"], p["WorkingDirectory"], strings.TrimPrefix(p["Description"], "stayup: "))
387	}
388	return w.Flush()
389}
390
391// fmtTS trims systemd's "Day YYYY-MM-DD HH:MM:SS TZ" to "YYYY-MM-DD HH:MM".
392func fmtTS(s string) string {
393	if f := strings.Fields(s); len(f) >= 3 {
394		return f[1] + " " + f[2][:5]
395	}
396	return "-"
397}
398
399// ---- foreground ------------------------------------------------------------
400
401const hint = "ctrl-c to kill, ctrl-d or close connection to hang up and stay up"
402
403func foreground(unit string) error {
404	// Log follower: journalctl, or tail -F for file outputs.
405	var follow *exec.Cmd
406	p := props(unit, "StandardOutput", "StandardError")
407	var files []string
408	for _, v := range []string{p["StandardOutput"], p["StandardError"]} {
409		if f, ok := strings.CutPrefix(v, "append:"); ok && (len(files) == 0 || files[0] != f) {
410			files = append(files, f)
411		}
412	}
413	if len(files) > 0 {
414		follow = exec.Command("tail", append([]string{"-n", "50", "-F"}, files...)...)
415	} else {
416		ja := []string{"-u", unit, "-f", "-n", "50", "-o", "cat"}
417		if user {
418			ja = append([]string{"--user"}, ja...)
419		}
420		follow = exec.Command("journalctl", ja...)
421	}
422	follow.Stdout, follow.Stderr = os.Stdout, os.Stderr
423	follow.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} // keep terminal ctrl-c for us
424	if err := follow.Start(); err != nil {
425		return err
426	}
427	defer follow.Process.Kill()
428
429	tty := isTerminal(os.Stderr.Fd())
430	if tty {
431		pinHint(true)
432		defer unpinHint()
433	} else {
434		fmt.Fprintln(os.Stderr, hint)
435	}
436
437	sig := make(chan os.Signal, 1)
438	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGWINCH)
439	eof := make(chan struct{})
440	if isTerminal(os.Stdin.Fd()) {
441		go func() {
442			r := bufio.NewReader(os.Stdin)
443			for {
444				if _, err := r.ReadByte(); err != nil {
445					close(eof)
446					return
447				}
448			}
449		}()
450	}
451	done := make(chan error, 1)
452	go func() { done <- follow.Wait() }()
453
454	for {
455		select {
456		case s := <-sig:
457			switch s {
458			case syscall.SIGWINCH:
459				if tty {
460					pinHint(false)
461				}
462				continue
463			case syscall.SIGINT:
464				unpinHint()
465				fmt.Fprintln(os.Stderr, "\nkilling", command(unit))
466				return remove(unit)
467			default: // SIGHUP, SIGTERM: detach, leave it running
468				return nil
469			}
470		case <-eof:
471			unpinHint()
472			fmt.Fprintf(os.Stderr, "\ndetached; %s stays up\n", command(unit))
473			return nil
474		case err := <-done:
475			return fmt.Errorf("log output ended: %v", err)
476		}
477	}
478}
479
480// Terminal helpers: keep the hint on the last row using a scroll region.
481
482func isTerminal(fd uintptr) bool {
483	var t syscall.Termios
484	_, _, e := syscall.Syscall(syscall.SYS_IOCTL, fd, syscall.TCGETS, uintptr(unsafe.Pointer(&t)))
485	return e == 0
486}
487
488func rows() int {
489	var ws struct{ Row, Col, X, Y uint16 }
490	syscall.Syscall(syscall.SYS_IOCTL, os.Stderr.Fd(), syscall.TIOCGWINSZ, uintptr(unsafe.Pointer(&ws)))
491	if ws.Row == 0 {
492		return 24
493	}
494	return int(ws.Row)
495}
496
497func pinHint(clear bool) {
498	h := rows()
499	if clear {
500		fmt.Fprint(os.Stderr, "\x1b[2J")
501	}
502	// scroll region = rows 1..h-1, hint on row h (reverse video), cursor back into the region
503	fmt.Fprintf(os.Stderr, "\x1b[1;%dr\x1b[%d;1H\x1b[2K\x1b[7m%s\x1b[0m\x1b[%d;1H", h-1, h, hint, h-1)
504}
505
506func unpinHint() {
507	if isTerminal(os.Stderr.Fd()) {
508		fmt.Fprintf(os.Stderr, "\x1b[r\x1b[%d;1H\x1b[2K", rows())
509	}
510}