2022-06-19 01:52:04 +00:00
|
|
|
package pmuxlib
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
|
|
import (
|
|
|
|
"bufio"
|
|
|
|
"context"
|
|
|
|
"errors"
|
|
|
|
"fmt"
|
|
|
|
"io"
|
|
|
|
"os"
|
|
|
|
"os/exec"
|
|
|
|
"strings"
|
|
|
|
"sync"
|
2022-06-30 19:31:32 +00:00
|
|
|
"syscall"
|
2022-01-23 03:03:13 +00:00
|
|
|
"time"
|
|
|
|
)
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
// ProcessConfig is used to configure a process.
|
2022-06-19 01:52:04 +00:00
|
|
|
type ProcessConfig struct {
|
2022-01-23 03:03:13 +00:00
|
|
|
// Cmd and Args describe the actual process to run.
|
|
|
|
Cmd string `yaml:"cmd"`
|
|
|
|
Args []string `yaml:"args"`
|
|
|
|
|
|
|
|
// Env describes the environment variables to set on the process.
|
|
|
|
Env map[string]string `yaml:"env"`
|
|
|
|
|
2022-01-23 03:39:34 +00:00
|
|
|
// Dir is the directory the process will be run in. If not set then the
|
|
|
|
// process is run in the same directory as this parent process.
|
2022-02-25 03:48:37 +00:00
|
|
|
Dir string `yaml:"dir"`
|
2022-01-23 03:39:34 +00:00
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
// MinWait and MaxWait are the minimum and maximum amount of time between
|
2024-07-19 13:35:38 +00:00
|
|
|
// restarts that Process will wait.
|
2022-01-23 03:03:13 +00:00
|
|
|
//
|
|
|
|
// MinWait defaults to 1 second.
|
|
|
|
// MaxWait defaults to 64 seconds.
|
|
|
|
MinWait time.Duration `yaml:"minWait"`
|
|
|
|
MaxWait time.Duration `yaml:"maxWait"`
|
|
|
|
|
|
|
|
// SigKillWait is the amount of time after the process is sent a SIGINT
|
2024-07-19 13:35:38 +00:00
|
|
|
// before a SIGKILL is sent.
|
2022-01-23 03:03:13 +00:00
|
|
|
//
|
|
|
|
// Defalts to 10 seconds.
|
|
|
|
SigKillWait time.Duration `yaml:"sigKillWait"`
|
2022-02-25 03:48:37 +00:00
|
|
|
|
|
|
|
// NoRestartOn indicates which exit codes should result in the process not
|
|
|
|
// being restarted any further.
|
2022-02-27 20:26:18 +00:00
|
|
|
NoRestartOn []int `yaml:"noRestartOn"`
|
2022-10-20 18:52:52 +00:00
|
|
|
|
|
|
|
// StartAfterFunc will cause the starting of the process to be blocked until
|
|
|
|
// this function returns. If used with RunProcessOnce and an error is
|
|
|
|
// returned, RunProcessOnce will return that error.
|
|
|
|
StartAfterFunc func(context.Context) error
|
2024-07-19 13:49:13 +00:00
|
|
|
|
|
|
|
// Group can be used to control the order and grouping of processes as they
|
|
|
|
// shut down.
|
|
|
|
//
|
|
|
|
// Processes will not be shut down until all processes with a higher group
|
|
|
|
// number are already shut down. Processes with the same group number will
|
|
|
|
// be shut down simultaneously.
|
|
|
|
Group int
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
2022-06-19 01:52:04 +00:00
|
|
|
func (cfg ProcessConfig) withDefaults() ProcessConfig {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
|
|
if cfg.MinWait == 0 {
|
|
|
|
cfg.MinWait = 1 * time.Second
|
|
|
|
}
|
|
|
|
|
|
|
|
if cfg.MaxWait == 0 {
|
|
|
|
cfg.MaxWait = 64 * time.Second
|
|
|
|
}
|
|
|
|
|
|
|
|
if cfg.SigKillWait == 0 {
|
|
|
|
cfg.SigKillWait = 10 * time.Second
|
|
|
|
}
|
|
|
|
|
|
|
|
return cfg
|
|
|
|
}
|
|
|
|
|
2022-06-30 19:31:32 +00:00
|
|
|
func sigProcessGroup(sysLogger Logger, proc *os.Process, sig syscall.Signal) {
|
|
|
|
sysLogger.Printf("sending %v signal", sig)
|
|
|
|
|
|
|
|
// Because we use Setpgid when starting child processes, child processes
|
|
|
|
// will have the same PGID as their PID. To send a signal to all processes
|
|
|
|
// in a group, you send the signal to the negation of the PGID, which in
|
|
|
|
// this case is equivalent to -PID.
|
|
|
|
//
|
|
|
|
// POSIX is a fucking joke.
|
|
|
|
if err := syscall.Kill(-proc.Pid, sig); err != nil {
|
|
|
|
|
|
|
|
panic(fmt.Errorf(
|
|
|
|
"failed to send %v signal to %d: %w",
|
|
|
|
sig, -proc.Pid, err,
|
|
|
|
))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-06-19 01:52:04 +00:00
|
|
|
// RunProcessOnce runs the process described by the ProcessConfig (though it
|
2022-06-30 19:42:57 +00:00
|
|
|
// doesn't use all fields from the ProcessConfig).
|
2022-01-23 03:36:07 +00:00
|
|
|
//
|
2022-06-30 19:42:57 +00:00
|
|
|
// The process is killed if-and-only-if the context is canceled, returning -1
|
|
|
|
// and context.Canceled. Otherwise the exit status of the process is returned,
|
|
|
|
// or -1 and an error.
|
2022-01-23 03:36:07 +00:00
|
|
|
//
|
2022-02-27 18:12:26 +00:00
|
|
|
// The stdout and stderr of the process will be written to the corresponding
|
|
|
|
// Loggers. Various runtime events will be written to the sysLogger.
|
|
|
|
func RunProcessOnce(
|
|
|
|
ctx context.Context,
|
|
|
|
stdoutLogger, stderrLogger, sysLogger Logger,
|
2022-06-19 01:52:04 +00:00
|
|
|
cfg ProcessConfig,
|
2022-02-27 18:12:26 +00:00
|
|
|
) (
|
|
|
|
int, error,
|
|
|
|
) {
|
2022-01-23 03:36:07 +00:00
|
|
|
|
|
|
|
cfg = cfg.withDefaults()
|
2022-01-23 03:03:13 +00:00
|
|
|
|
2022-10-20 18:52:52 +00:00
|
|
|
if cfg.StartAfterFunc != nil {
|
|
|
|
if err := cfg.StartAfterFunc(ctx); err != nil {
|
|
|
|
return -1, err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
var wg sync.WaitGroup
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
fwdOutPipe := func(logger Logger, r io.Reader) {
|
2022-01-23 03:03:13 +00:00
|
|
|
wg.Add(1)
|
|
|
|
go func() {
|
|
|
|
defer wg.Done()
|
|
|
|
bufR := bufio.NewReader(r)
|
|
|
|
for {
|
|
|
|
line, err := bufR.ReadString('\n')
|
|
|
|
if errors.Is(err, io.EOF) {
|
|
|
|
return
|
|
|
|
} else if err != nil {
|
|
|
|
logger.Printf("reading output: %v", err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
logger.Println(strings.TrimSuffix(line, "\n"))
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
}
|
|
|
|
|
|
|
|
cmd := exec.Command(cfg.Cmd, cfg.Args...)
|
2022-01-23 03:57:59 +00:00
|
|
|
|
2022-01-23 03:39:34 +00:00
|
|
|
cmd.Dir = cfg.Dir
|
2022-01-23 03:03:13 +00:00
|
|
|
|
2022-06-30 19:31:32 +00:00
|
|
|
cmd.SysProcAttr = &syscall.SysProcAttr{
|
|
|
|
// Indicates that the child process should be a part of a separate
|
|
|
|
// process group than the parent, so that it does not receive signals
|
|
|
|
// that the parent receives. This is what ensures that context
|
|
|
|
// cancellation is the only way to interrupt the child processes.
|
|
|
|
Setpgid: true,
|
|
|
|
}
|
|
|
|
|
|
|
|
cmd.Env = os.Environ()
|
2022-01-23 03:03:13 +00:00
|
|
|
for k, v := range cfg.Env {
|
|
|
|
cmd.Env = append(cmd.Env, k+"="+v)
|
|
|
|
}
|
|
|
|
|
|
|
|
stdout, err := cmd.StdoutPipe()
|
|
|
|
if err != nil {
|
2022-02-25 03:48:37 +00:00
|
|
|
return -1, fmt.Errorf("getting stdout pipe: %w", err)
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
defer stdout.Close()
|
|
|
|
|
|
|
|
stderr, err := cmd.StderrPipe()
|
|
|
|
if err != nil {
|
2022-02-25 03:48:37 +00:00
|
|
|
return -1, fmt.Errorf("getting stderr pipe: %w", err)
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
defer stderr.Close()
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
fwdOutPipe(stdoutLogger, stdout)
|
|
|
|
fwdOutPipe(stderrLogger, stderr)
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
|
|
if err := cmd.Start(); err != nil {
|
2022-02-25 03:48:37 +00:00
|
|
|
return -1, fmt.Errorf("starting process: %w", err)
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
stopCh := make(chan struct{})
|
2022-01-23 03:36:07 +00:00
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
go func(proc *os.Process) {
|
2022-06-30 19:31:32 +00:00
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
select {
|
|
|
|
case <-ctx.Done():
|
2022-06-30 19:31:32 +00:00
|
|
|
sigProcessGroup(sysLogger, proc, syscall.SIGINT)
|
2022-01-23 03:03:13 +00:00
|
|
|
case <-stopCh:
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
select {
|
|
|
|
case <-time.After(cfg.SigKillWait):
|
2022-06-30 19:31:32 +00:00
|
|
|
sigProcessGroup(sysLogger, proc, syscall.SIGKILL)
|
2022-01-23 03:03:13 +00:00
|
|
|
case <-stopCh:
|
|
|
|
}
|
2022-06-30 19:31:32 +00:00
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
}(cmd.Process)
|
|
|
|
|
|
|
|
wg.Wait()
|
|
|
|
|
2022-02-25 03:48:37 +00:00
|
|
|
err = cmd.Wait()
|
2022-02-27 18:45:33 +00:00
|
|
|
close(stopCh)
|
|
|
|
|
2022-06-30 19:42:57 +00:00
|
|
|
if err := ctx.Err(); err != nil {
|
|
|
|
return -1, err
|
|
|
|
}
|
2022-02-25 03:48:37 +00:00
|
|
|
|
|
|
|
if err != nil {
|
2022-06-30 19:42:57 +00:00
|
|
|
return -1, fmt.Errorf("process exited: %w", err)
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
2022-06-30 19:42:57 +00:00
|
|
|
return cmd.ProcessState.ExitCode(), nil
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
// Process is used to manage a running process. Methods on a Process are _not_
|
|
|
|
// thread-safe.
|
|
|
|
//
|
|
|
|
// Stop must be called on a Process before the program has exited, or there may
|
|
|
|
// be a leftover zombie child process.
|
2022-01-23 03:03:13 +00:00
|
|
|
//
|
|
|
|
// The process will be restarted if it exits of its own accord. There will be a
|
2022-06-30 19:31:32 +00:00
|
|
|
// brief wait time between each restart, with an exponential backoff mechanism
|
|
|
|
// so that the wait time increases upon repeated restarts.
|
2022-01-23 03:03:13 +00:00
|
|
|
//
|
2022-02-27 18:12:26 +00:00
|
|
|
// The stdout and stderr of the process will be written to the corresponding
|
|
|
|
// Loggers. Various runtime events will be written to the sysLogger.
|
2024-07-19 13:35:38 +00:00
|
|
|
type Process struct {
|
|
|
|
cfg ProcessConfig
|
|
|
|
stdoutLogger, stderrLogger, sysLogger Logger
|
2022-01-23 03:03:13 +00:00
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
stopFn context.CancelFunc
|
|
|
|
doneCh chan struct{}
|
|
|
|
}
|
2022-01-23 03:03:13 +00:00
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
// NewProcess returns a new Process instance based on the given config.
|
|
|
|
func NewProcess(
|
|
|
|
cfg ProcessConfig, stdoutLogger, stderrLogger, sysLogger Logger,
|
|
|
|
) *Process {
|
|
|
|
p := &Process{
|
|
|
|
cfg: cfg.withDefaults(),
|
|
|
|
stdoutLogger: stdoutLogger,
|
|
|
|
stderrLogger: stderrLogger,
|
|
|
|
sysLogger: sysLogger,
|
|
|
|
}
|
|
|
|
p.run()
|
|
|
|
return p
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *Process) run() {
|
|
|
|
var ctx context.Context
|
|
|
|
ctx, p.stopFn = context.WithCancel(context.Background())
|
|
|
|
p.doneCh = make(chan struct{})
|
|
|
|
|
|
|
|
go func() {
|
|
|
|
defer close(p.doneCh)
|
|
|
|
p.restartLoop(ctx)
|
|
|
|
}()
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *Process) restartLoop(ctx context.Context) {
|
2022-01-23 03:03:13 +00:00
|
|
|
var wait time.Duration
|
|
|
|
|
|
|
|
for {
|
|
|
|
start := time.Now()
|
2022-02-27 18:12:26 +00:00
|
|
|
exitCode, err := RunProcessOnce(
|
|
|
|
ctx,
|
2024-07-19 13:35:38 +00:00
|
|
|
p.stdoutLogger, p.stderrLogger, p.sysLogger,
|
|
|
|
p.cfg,
|
2022-02-27 18:12:26 +00:00
|
|
|
)
|
2022-01-23 03:03:13 +00:00
|
|
|
took := time.Since(start)
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
// TODO check if error was due to StartAfterFunc, change the log if so.
|
2022-01-23 03:03:13 +00:00
|
|
|
if err != nil {
|
2024-07-19 13:35:38 +00:00
|
|
|
p.sysLogger.Printf("exited: %v", err)
|
2022-01-23 03:03:13 +00:00
|
|
|
} else {
|
2024-07-19 13:35:38 +00:00
|
|
|
p.sysLogger.Printf("exit code: %d", exitCode)
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
if err := ctx.Err(); err != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
for i := range p.cfg.NoRestartOn {
|
|
|
|
if p.cfg.NoRestartOn[i] == exitCode {
|
2022-02-25 03:48:37 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
wait = ((wait * 2) - took).Truncate(time.Millisecond)
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
if wait < p.cfg.MinWait {
|
|
|
|
wait = p.cfg.MinWait
|
|
|
|
} else if wait > p.cfg.MaxWait {
|
|
|
|
wait = p.cfg.MaxWait
|
2022-01-23 03:03:13 +00:00
|
|
|
}
|
|
|
|
|
2024-07-19 13:35:38 +00:00
|
|
|
p.sysLogger.Printf("will restart process in %v", wait)
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
|
|
select {
|
|
|
|
case <-time.After(wait):
|
|
|
|
case <-ctx.Done():
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2024-07-19 13:35:38 +00:00
|
|
|
|
|
|
|
// Restart will block until the currently running child process has been killed
|
|
|
|
// and a new one has been spawned.
|
|
|
|
func (p *Process) Restart() {
|
|
|
|
p.Stop()
|
|
|
|
p.run()
|
|
|
|
}
|
|
|
|
|
|
|
|
// Stop will block until the child process has been killed. The Process should
|
|
|
|
// not be used again after this.
|
|
|
|
func (p *Process) Stop() {
|
|
|
|
p.stopFn()
|
|
|
|
<-p.doneCh
|
|
|
|
}
|