2021-09-21 22:36:50 +00:00
|
|
|
|
package main
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
"bufio"
|
|
|
|
|
"context"
|
|
|
|
|
"errors"
|
|
|
|
|
"flag"
|
|
|
|
|
"fmt"
|
|
|
|
|
"io"
|
|
|
|
|
"io/ioutil"
|
|
|
|
|
"os"
|
|
|
|
|
"os/signal"
|
|
|
|
|
"strings"
|
|
|
|
|
"sync"
|
|
|
|
|
"syscall"
|
|
|
|
|
"time"
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
"github.com/cryptic-io/pmux/pmuxproc"
|
|
|
|
|
|
2021-09-21 22:36:50 +00:00
|
|
|
|
"gopkg.in/yaml.v2"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
// pname used by pmux itself for logging.
|
|
|
|
|
const pmuxPName = "pmux"
|
|
|
|
|
|
|
|
|
|
// characters used to denote different kinds of logs
|
|
|
|
|
const (
|
2022-02-27 18:12:26 +00:00
|
|
|
|
logSepStdout = '›'
|
|
|
|
|
logSepStderr = '»'
|
|
|
|
|
logSepSys = '~'
|
2021-09-21 22:36:50 +00:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
type logger struct {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
timeFmt string
|
2021-09-21 22:36:50 +00:00
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
l *sync.Mutex
|
2021-09-21 22:36:50 +00:00
|
|
|
|
out io.Writer
|
|
|
|
|
outBuf *bufio.Writer
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
// maxPNameLen is a pointer because it changes when WithPrefix is called.
|
|
|
|
|
maxPNameLen *uint64
|
|
|
|
|
|
|
|
|
|
pname string
|
|
|
|
|
sep rune
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func newLogger(
|
|
|
|
|
out io.Writer,
|
2022-02-27 18:12:26 +00:00
|
|
|
|
sep rune,
|
2021-09-21 22:36:50 +00:00
|
|
|
|
timeFmt string,
|
|
|
|
|
) *logger {
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
pname := pmuxPName
|
|
|
|
|
maxPNameLen := uint64(len(pname))
|
2021-09-21 22:36:50 +00:00
|
|
|
|
|
|
|
|
|
l := &logger{
|
|
|
|
|
timeFmt: timeFmt,
|
2022-01-23 03:03:13 +00:00
|
|
|
|
maxPNameLen: &maxPNameLen,
|
|
|
|
|
l: new(sync.Mutex),
|
2021-09-21 22:36:50 +00:00
|
|
|
|
out: out,
|
|
|
|
|
outBuf: bufio.NewWriter(out),
|
2022-01-23 03:03:13 +00:00
|
|
|
|
pname: pname,
|
2022-02-27 18:12:26 +00:00
|
|
|
|
sep: sep,
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return l
|
|
|
|
|
}
|
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
|
func (l *logger) withSep(sep rune) *logger {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
l2 := *l
|
|
|
|
|
l2.sep = sep
|
2022-02-27 18:12:26 +00:00
|
|
|
|
return &l2
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (l *logger) withPName(pname string) *logger {
|
|
|
|
|
l2 := *l
|
|
|
|
|
l2.pname = pname
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
|
|
|
|
l2.l.Lock()
|
|
|
|
|
defer l2.l.Unlock()
|
|
|
|
|
|
|
|
|
|
if pnameLen := uint64(len(pname)); pnameLen > *l2.maxPNameLen {
|
|
|
|
|
*l2.maxPNameLen = pnameLen
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return &l2
|
|
|
|
|
}
|
|
|
|
|
|
2021-09-21 22:36:50 +00:00
|
|
|
|
func (l *logger) Close() {
|
|
|
|
|
|
|
|
|
|
l.l.Lock()
|
|
|
|
|
defer l.l.Unlock()
|
|
|
|
|
|
|
|
|
|
l.outBuf.Flush()
|
|
|
|
|
|
|
|
|
|
if syncer, ok := l.out.(interface{ Sync() error }); ok {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
_ = syncer.Sync()
|
2021-09-21 22:36:50 +00:00
|
|
|
|
} else if flusher, ok := l.out.(interface{ Flush() error }); ok {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
_ = flusher.Flush()
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// this generally shouldn't be necessary, but we could run into cases (e.g.
|
|
|
|
|
// during a force-kill) where further Prints are called after a Close. These
|
|
|
|
|
// should just do nothing.
|
|
|
|
|
l.out = ioutil.Discard
|
|
|
|
|
l.outBuf = bufio.NewWriter(l.out)
|
|
|
|
|
}
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
func (l *logger) println(line string) {
|
2021-09-21 22:36:50 +00:00
|
|
|
|
|
|
|
|
|
l.l.Lock()
|
|
|
|
|
defer l.l.Unlock()
|
|
|
|
|
|
2022-02-27 17:37:22 +00:00
|
|
|
|
if l.timeFmt != "" {
|
|
|
|
|
fmt.Fprintf(
|
|
|
|
|
l.outBuf,
|
|
|
|
|
"%s %c ",
|
|
|
|
|
time.Now().Format(l.timeFmt),
|
|
|
|
|
l.sep,
|
|
|
|
|
)
|
|
|
|
|
}
|
|
|
|
|
|
2021-09-21 22:36:50 +00:00
|
|
|
|
fmt.Fprintf(
|
|
|
|
|
l.outBuf,
|
2022-02-27 17:37:22 +00:00
|
|
|
|
"%s%s%c %s\n",
|
2022-01-23 03:03:13 +00:00
|
|
|
|
l.pname,
|
|
|
|
|
strings.Repeat(" ", int(*l.maxPNameLen+1)-len(l.pname)),
|
|
|
|
|
l.sep,
|
2021-09-21 22:36:50 +00:00
|
|
|
|
line,
|
|
|
|
|
)
|
2022-02-27 18:12:26 +00:00
|
|
|
|
|
|
|
|
|
l.outBuf.Flush()
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
func (l *logger) Println(line string) {
|
|
|
|
|
l.println(line)
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
func (l *logger) Printf(msg string, args ...interface{}) {
|
|
|
|
|
l.Println(fmt.Sprintf(msg, args...))
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
2022-01-23 03:03:13 +00:00
|
|
|
|
type processConfig struct {
|
|
|
|
|
pmuxproc.Config `yaml:",inline"`
|
|
|
|
|
Name string `yaml:"name"`
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type config struct {
|
|
|
|
|
TimeFormat string `yaml:"timeFormat"`
|
2022-01-23 03:03:13 +00:00
|
|
|
|
Processes []processConfig `yaml:"processes"`
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (cfg config) init() (config, error) {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
2021-09-21 22:36:50 +00:00
|
|
|
|
if len(cfg.Processes) == 0 {
|
|
|
|
|
return config{}, errors.New("no processes defined")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return cfg, nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func main() {
|
|
|
|
|
|
|
|
|
|
cfgPath := flag.String("c", "./pmux.yml", "Path to config yaml file")
|
|
|
|
|
flag.Parse()
|
|
|
|
|
|
|
|
|
|
cfgB, err := ioutil.ReadFile(*cfgPath)
|
|
|
|
|
if err != nil {
|
|
|
|
|
panic(fmt.Sprintf("couldn't read cfg file at %q: %v", *cfgPath, err))
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
var cfg config
|
|
|
|
|
if err := yaml.Unmarshal(cfgB, &cfg); err != nil {
|
|
|
|
|
panic(fmt.Sprintf("couldn't parse cfg file: %v", err))
|
|
|
|
|
|
|
|
|
|
} else if cfg, err = cfg.init(); err != nil {
|
|
|
|
|
panic(fmt.Sprintf("initializing config: %v", err))
|
|
|
|
|
}
|
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
|
stdoutLogger := newLogger(os.Stdout, logSepStdout, cfg.TimeFormat)
|
|
|
|
|
defer stdoutLogger.Close()
|
|
|
|
|
|
|
|
|
|
stderrLogger := newLogger(os.Stderr, logSepStderr, cfg.TimeFormat)
|
|
|
|
|
defer stderrLogger.Close()
|
|
|
|
|
|
|
|
|
|
sysLogger := stderrLogger.withSep(logSepSys)
|
|
|
|
|
|
|
|
|
|
defer sysLogger.Println("exited gracefully, ciao!")
|
2021-09-21 22:36:50 +00:00
|
|
|
|
|
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
|
go func() {
|
2022-01-23 03:03:13 +00:00
|
|
|
|
sigCh := make(chan os.Signal, 2)
|
2021-09-21 22:36:50 +00:00
|
|
|
|
signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
|
|
|
|
|
|
|
|
|
|
sig := <-sigCh
|
2022-02-27 18:12:26 +00:00
|
|
|
|
sysLogger.Printf("%v signal received, killing all sub-processes", sig)
|
2021-09-21 22:36:50 +00:00
|
|
|
|
cancel()
|
|
|
|
|
|
|
|
|
|
<-sigCh
|
2022-02-27 18:12:26 +00:00
|
|
|
|
sysLogger.Printf("forcefully exiting pmux process, there may be zombie child processes being left behind, good luck!")
|
|
|
|
|
sysLogger.Close()
|
2021-09-21 22:36:50 +00:00
|
|
|
|
os.Exit(1)
|
|
|
|
|
}()
|
|
|
|
|
|
|
|
|
|
var wg sync.WaitGroup
|
|
|
|
|
defer wg.Wait()
|
|
|
|
|
|
|
|
|
|
for _, cfgProc := range cfg.Processes {
|
|
|
|
|
wg.Add(1)
|
2022-01-23 03:03:13 +00:00
|
|
|
|
go func(procCfg processConfig) {
|
2021-09-21 22:36:50 +00:00
|
|
|
|
defer wg.Done()
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
|
stdoutLogger := stdoutLogger.withPName(procCfg.Name)
|
|
|
|
|
stderrLogger := stderrLogger.withPName(procCfg.Name)
|
|
|
|
|
sysLogger := sysLogger.withPName(procCfg.Name)
|
|
|
|
|
|
|
|
|
|
defer sysLogger.Printf("stopped process handler")
|
|
|
|
|
|
|
|
|
|
sysLogger.Println("starting process")
|
2022-01-23 03:03:13 +00:00
|
|
|
|
|
2022-02-27 18:12:26 +00:00
|
|
|
|
pmuxproc.RunProcess(
|
|
|
|
|
ctx, stdoutLogger, stderrLogger, sysLogger, procCfg.Config,
|
|
|
|
|
)
|
2021-09-21 22:36:50 +00:00
|
|
|
|
}(cfgProc)
|
|
|
|
|
}
|
|
|
|
|
}
|