pmux/pmuxlib/pmuxlib.go

// Package pmuxlib implements the process management aspects of the pmux
// process.
package pmuxlib

import (
	"fmt"
	"io"
	"sync"
)

type Config struct {
	TimeFormat string `yaml:"timeFormat"`

	// Set of processes to run, keyed by their name.
	Processes map[string]ProcessConfig `yaml:"processes"`
}

// Pmux manages multiple child Processes. Methods on a Pmux instance are _not_
// thread-safe.
//
// Stop must be called on a Pmux before the program has exited, or there may be
// a leftover zombie child process.
type Pmux struct {
	processes map[string]*Process

	stdoutLogger, stderrLogger, sysLogger Logger
}

// NewPmux starts a Pmux with the given configuration.
func NewPmux(cfg Config, stdout, stderr io.Writer) *Pmux {
	var (
		stdoutLogger = NewPmuxLogger(stdout, LogSepStdout, cfg.TimeFormat)
		stderrLogger = NewPmuxLogger(stderr, LogSepStderr, cfg.TimeFormat)
		sysLogger    = stderrLogger.WithSep(LogSepSys)

		p = &Pmux{
			processes:    map[string]*Process{},
			stdoutLogger: stdoutLogger,
			stderrLogger: stderrLogger,
			sysLogger:    sysLogger,
		}
	)

	for name, cfgProc := range cfg.Processes {
		cfgProc.StdoutLogger = stdoutLogger.WithProcessName(name)
		cfgProc.StderrLogger = stderrLogger.WithProcessName(name)
		cfgProc.SysLogger = sysLogger.WithProcessName(name)
		p.processes[name] = NewProcess(cfgProc)
	}

	return p
}

// Restart will block until the child process of the given name has been killed
// and a new one has been spawned. If there is no child of the given name then
// Restart panics.
func (p *Pmux) Restart(name string) {
	proc, ok := p.processes[name]
	if !ok {
		panic(fmt.Sprintf("no process named %q", name))
	}
	proc.Restart()
}

// Stop will block until all child processes have been killed. The Pmux should
// not be used again after this.
func (p *Pmux) Stop() {
	var wg sync.WaitGroup

	var minGroup, maxGroup int

	for _, proc := range p.processes {
		if maxGroup < proc.cfg.Group {
			maxGroup = proc.cfg.Group
		}

		if minGroup > proc.cfg.Group {
			minGroup = proc.cfg.Group
		}
	}

	for group := maxGroup; group >= minGroup; group-- {
		p.sysLogger.Printf("killing child processes (group %d)", group)
		for _, proc := range p.processes {
			proc := proc
			if proc.cfg.Group != group {
				continue
			}

			wg.Add(1)
			go func() {
				defer wg.Done()
				proc.Stop()
			}()
		}
		wg.Wait()
	}

	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()
}