Refactor logging to expose the default PmuxLogger

main.go

@@ -3,7 +3,6 @@ package main
 import (
 	"flag"
 	"fmt"
-	"io/ioutil"
 	"os"
 	"os/signal"
 	"syscall"
@@ -18,7 +17,7 @@ func main() {
 	cfgPath := flag.String("c", "./pmux.yml", "Path to config yaml file")
 	flag.Parse()
 
-	cfgB, err := ioutil.ReadFile(*cfgPath)
+	cfgB, err := os.ReadFile(*cfgPath)
 	if err != nil {
 		panic(fmt.Sprintf("couldn't read cfg file at %q: %v", *cfgPath, err))
 	}

pmuxlib/logger.go

@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"strings"
 	"sync"
 	"time"
@@ -13,11 +12,12 @@ import (
 // pname used by pmux itself for logging.
 const pmuxPName = "pmux"
 
-// characters used to denote different kinds of logs
+// Characters used to denote different kinds of logs in the default Pmux
+// configuration.
 const (
-	logSepStdout = '›'
-	logSepStderr = '»'
-	logSepSys    = '~'
+	LogSepStdout = '›'
+	LogSepStderr = '»'
+	LogSepSys    = '~'
 )
 
 // Logger is used by RunProcess to log process details in realtime. You can use
@@ -25,20 +25,26 @@ const (
 type Logger interface {
 	Println(string)
 	Printf(string, ...interface{})
+	Close() error
 }
 
 // NullLogger is an implementation of Logger which doesn't do anything.
 type NullLogger struct{}
 
-func (*NullLogger) Println(string)                {}
-func (*NullLogger) Printf(string, ...interface{}) {}
+var _ Logger = NullLogger{}
+
+func (NullLogger) Println(string)                {}
+func (NullLogger) Printf(string, ...interface{}) {}
+func (NullLogger) Close() error                  { return nil }
 
 // PlainLogger implements Logger by writing each line directly to the given
 // io.Writer as-is.
 type PlainLogger struct {
-	io.Writer
+	io.WriteCloser
 }
 
+var _ Logger = PlainLogger{}
+
 func (l PlainLogger) Println(line string) {
 	fmt.Fprintln(l, line)
 }
@@ -47,7 +53,15 @@ func (l PlainLogger) Printf(str string, args ...interface{}) {
 	fmt.Fprintf(l, str, args...)
 }
 
-type logger struct {
+func (l PlainLogger) Close() error {
+	return l.WriteCloser.Close()
+}
+
+// PmuxLogger is used by the pmux process itself for logging. It can prefix log
+// lines with a timestamp, the name of the process being logged, and a custom
+// separator in front of the log line to help delineate one kind of log from
+// another.
+type PmuxLogger struct {
 	timeFmt string
 
 	l      *sync.Mutex
@@ -61,16 +75,21 @@ type logger struct {
 	sep   rune
 }
 
-func newLogger(
-	out io.Writer,
-	sep rune,
-	timeFmt string,
-) *logger {
+var _ Logger = (*PmuxLogger)(nil)
 
-	pname := pmuxPName
-	maxPNameLen := uint64(len(pname))
+// NewPmuxLogger returns a PmuxLogger which will write lines to the given
+// io.Writer, separating metadata from the line itself using the given
+// separator.
+//
+// If timeFmt is not empty string then the timestamp of each line will be output
+// using that format.
+func NewPmuxLogger(out io.Writer, sep rune, timeFmt string) *PmuxLogger {
+	var (
+		pname       = pmuxPName
+		maxPNameLen = uint64(len(pname))
+	)
 
-	l := &logger{
+	return &PmuxLogger{
 		timeFmt:     timeFmt,
 		maxPNameLen: &maxPNameLen,
 		l:           new(sync.Mutex),
@@ -79,17 +98,23 @@ func newLogger(
 		pname:       pname,
 		sep:         sep,
 	}
-
-	return l
 }
 
-func (l *logger) withSep(sep rune) *logger {
+// WithSep returns a copy of the PmuxLogger which uses the given separator. The
+// original PmuxLogger will continue to function normally.
+func (l *PmuxLogger) WithSep(sep rune) *PmuxLogger {
 	l2 := *l
 	l2.sep = sep
 	return &l2
 }
 
-func (l *logger) withPName(pname string) *logger {
+// WithProcessName returns a copy of the PmuxLogger which will prefix log lines
+// with the given process name.
+//
+// All PmuxLoggers which are spawned using a With* method, including the
+// original, will left pad their process name to the length of the longest
+// process name in the group.
+func (l *PmuxLogger) WithProcessName(pname string) *PmuxLogger {
 	l2 := *l
 	l2.pname = pname
 
@@ -103,27 +128,26 @@ func (l *logger) withPName(pname string) *logger {
 	return &l2
 }
 
-func (l *logger) Close() {
-
+// Close flushes all buffered log data, and sets the logger to discard all
+// future print calls.
+//
+// PmuxLoggers created via With* methods should be each closed individually.
+func (l *PmuxLogger) Close() error {
 	l.l.Lock()
 	defer l.l.Unlock()
 
-	l.outBuf.Flush()
-
-	if syncer, ok := l.out.(interface{ Sync() error }); ok {
-		_ = syncer.Sync()
-	} else if flusher, ok := l.out.(interface{ Flush() error }); ok {
-		_ = flusher.Flush()
-	}
+	err := l.outBuf.Flush()
 
 	// 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.out = io.Discard
 	l.outBuf = bufio.NewWriter(l.out)
+
+	return err
 }
 
-func (l *logger) println(line string) {
+func (l *PmuxLogger) println(line string) {
 
 	l.l.Lock()
 	defer l.l.Unlock()
@@ -149,10 +173,10 @@ func (l *logger) println(line string) {
 	l.outBuf.Flush()
 }
 
-func (l *logger) Println(line string) {
+func (l *PmuxLogger) Println(line string) {
 	l.println(line)
 }
 
-func (l *logger) Printf(msg string, args ...interface{}) {
+func (l *PmuxLogger) Printf(msg string, args ...interface{}) {
 	l.Println(fmt.Sprintf(msg, args...))
 }

pmuxlib/pmuxlib.go

@@ -22,31 +22,30 @@ type Config struct {
 // a leftover zombie child process.
 type Pmux struct {
 	processes map[string]*Process
-	sysLogger Logger
+
+	stdoutLogger, stderrLogger, sysLogger Logger
 }
 
 // NewPmux starts a Pmux with the given configuration.
 func NewPmux(cfg Config, stdout, stderr io.Writer) *Pmux {
-	stdoutLogger := newLogger(stdout, logSepStdout, cfg.TimeFormat)
-	defer stdoutLogger.Close()
+	var (
+		stdoutLogger = NewPmuxLogger(stdout, LogSepStdout, cfg.TimeFormat)
+		stderrLogger = NewPmuxLogger(stderr, LogSepStderr, cfg.TimeFormat)
+		sysLogger    = stderrLogger.WithSep(LogSepSys)
 
-	stderrLogger := newLogger(stderr, logSepStderr, cfg.TimeFormat)
-	defer stderrLogger.Close()
-
-	sysLogger := stderrLogger.withSep(logSepSys)
-
-	p := &Pmux{
+		p = &Pmux{
 			processes:    map[string]*Process{},
+			stdoutLogger: stdoutLogger,
+			stderrLogger: stderrLogger,
 			sysLogger:    sysLogger,
 		}
+	)
 
 	for name, cfgProc := range cfg.Processes {
-		stdoutLogger := stdoutLogger.withPName(name)
-		stderrLogger := stderrLogger.withPName(name)
-		sysLogger := sysLogger.withPName(name)
-		p.processes[name] = NewProcess(
-			cfgProc, stdoutLogger, stderrLogger, sysLogger,
-		)
+		cfgProc.StdoutLogger = stdoutLogger.WithProcessName(name)
+		cfgProc.StderrLogger = stderrLogger.WithProcessName(name)
+		cfgProc.SysLogger = sysLogger.WithProcessName(name)
+		p.processes[name] = NewProcess(cfgProc)
 	}
 
 	return p
@@ -97,5 +96,14 @@ func (p *Pmux) Stop() {
 		wg.Wait()
 	}
 
-	p.sysLogger.Println("exited gracefully, ciao!")
+	if err := p.stdoutLogger.Close(); err != nil {
+		p.sysLogger.Printf("Closing stdout logger failed: %v", err)
+	}
+
+	if err := p.stderrLogger.Close(); err != nil {
+		p.sysLogger.Printf("Closing stderr logger failed: %v", err)
+	}
+
+	p.sysLogger.Println("exited gracefully, ciao!")
+	_ = p.sysLogger.Close()
 }

pmuxlib/process.go

@@ -56,7 +56,12 @@ type ProcessConfig struct {
 	// 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
+	Group int `yaml:"group"`
+
+	// Where to send stdout, stderr log lines, and lines from pmux itself.
+	StdoutLogger Logger `yaml:"-"`
+	StderrLogger Logger `yaml:"-"`
+	SysLogger    Logger `yaml:"-"`
 }
 
 func (cfg ProcessConfig) withDefaults() ProcessConfig {
@@ -103,13 +108,7 @@ func sigProcessGroup(sysLogger Logger, proc *os.Process, sig syscall.Signal) {
 //
 // 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,
-	cfg ProcessConfig,
-) (
-	int, error,
-) {
+func RunProcessOnce(ctx context.Context, cfg ProcessConfig) (int, error) {
 
 	cfg = cfg.withDefaults()
 
@@ -169,8 +168,8 @@ func RunProcessOnce(
 	}
 	defer stderr.Close()
 
-	fwdOutPipe(stdoutLogger, stdout)
-	fwdOutPipe(stderrLogger, stderr)
+	fwdOutPipe(cfg.StdoutLogger, stdout)
+	fwdOutPipe(cfg.StderrLogger, stderr)
 
 	if err := cmd.Start(); err != nil {
 		return -1, fmt.Errorf("starting process: %w", err)
@@ -182,14 +181,14 @@ func RunProcessOnce(
 
 		select {
 		case <-ctx.Done():
-			sigProcessGroup(sysLogger, proc, syscall.SIGINT)
+			sigProcessGroup(cfg.SysLogger, proc, syscall.SIGINT)
 		case <-stopCh:
 			return
 		}
 
 		select {
 		case <-time.After(cfg.SigKillWait):
-			sigProcessGroup(sysLogger, proc, syscall.SIGKILL)
+			sigProcessGroup(cfg.SysLogger, proc, syscall.SIGKILL)
 		case <-stopCh:
 		}
 
@@ -225,21 +224,14 @@ func RunProcessOnce(
 // Loggers. Various runtime events will be written to the sysLogger.
 type Process struct {
 	cfg    ProcessConfig
-	stdoutLogger, stderrLogger, sysLogger Logger
-
 	stopFn context.CancelFunc
 	doneCh chan struct{}
 }
 
 // NewProcess returns a new Process instance based on the given config.
-func NewProcess(
-	cfg ProcessConfig, stdoutLogger, stderrLogger, sysLogger Logger,
-) *Process {
+func NewProcess(cfg ProcessConfig) *Process {
 	p := &Process{
 		cfg: cfg.withDefaults(),
-		stdoutLogger: stdoutLogger,
-		stderrLogger: stderrLogger,
-		sysLogger:    sysLogger,
 	}
 	p.run()
 	return p
@@ -261,18 +253,14 @@ func (p *Process) restartLoop(ctx context.Context) {
 
 	for {
 		start := time.Now()
-		exitCode, err := RunProcessOnce(
-			ctx,
-			p.stdoutLogger, p.stderrLogger, p.sysLogger,
-			p.cfg,
-		)
+		exitCode, err := RunProcessOnce(ctx, p.cfg)
 		took := time.Since(start)
 
 		// TODO check if error was due to StartAfterFunc, change the log if so.
 		if err != nil {
-			p.sysLogger.Printf("exited: %v", err)
+			p.cfg.SysLogger.Printf("exited: %v", err)
 		} else {
-			p.sysLogger.Printf("exit code: %d", exitCode)
+			p.cfg.SysLogger.Printf("exit code: %d", exitCode)
 		}
 
 		if err := ctx.Err(); err != nil {
@@ -293,7 +281,7 @@ func (p *Process) restartLoop(ctx context.Context) {
 			wait = p.cfg.MaxWait
 		}
 
-		p.sysLogger.Printf("will restart process in %v", wait)
+		p.cfg.SysLogger.Printf("will restart process in %v", wait)
 
 		select {
 		case <-time.After(wait):

shell.nix

@@ -6,12 +6,11 @@
   },
 }: let
     pkgs = (import pkgsSrc) {};
-in {
-    shell = pkgs.mkShell {
+in pkgs.mkShell {
     name = "project-shell";
     buildInputs = [
       pkgs.go
       pkgs.golangci-lint
+      pkgs.gopls
     ];
-    };
 }