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

	stderrLogger := newLogger(stderr, logSepStderr, cfg.TimeFormat)
	defer stderrLogger.Close()

	sysLogger := stderrLogger.withSep(logSepSys)

	p := &Pmux{
		processes: map[string]*Process{},
		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,
		)
	}

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

	p.sysLogger.Println("exited gracefully, ciao!")
}