pmux/pmuxlib/logger.go

package pmuxlib

import (
	"bufio"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"
)

// pname used by pmux itself for logging.
const pmuxPName = "pmux"

// Characters used to denote different kinds of logs in the default Pmux
// configuration.
const (
	LogSepStdout = '›'
	LogSepStderr = '»'
	LogSepSys    = '~'
)

// Logger is used by RunProcess to log process details in realtime. You can use
// a new(NullLogger) if you don't care.
type Logger interface {
	Println(string)
	Printf(string, ...interface{})
}

// NullLogger is an implementation of Logger which doesn't do anything.
type NullLogger struct{}

var _ Logger = NullLogger{}

func (NullLogger) Println(string)                {}
func (NullLogger) Printf(string, ...interface{}) {}

// PlainLogger implements Logger by writing each line directly to the given
// io.Writer as-is.
type PlainLogger struct {
	io.WriteCloser
}

var _ Logger = PlainLogger{}

func (l PlainLogger) Println(line string) {
	fmt.Fprintln(l, line)
}

func (l PlainLogger) Printf(str string, args ...interface{}) {
	fmt.Fprintf(l, str, args...)
}

// 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
	out    io.Writer
	outBuf *bufio.Writer

	// maxPNameLen is a pointer because it changes when WithPrefix is called.
	maxPNameLen *uint64

	pname string
	sep   rune
}

var _ Logger = (*PmuxLogger)(nil)

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

	return &PmuxLogger{
		timeFmt:     timeFmt,
		maxPNameLen: &maxPNameLen,
		l:           new(sync.Mutex),
		out:         out,
		outBuf:      bufio.NewWriter(out),
		pname:       pname,
		sep:         sep,
	}
}

// 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
}

// 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

	l2.l.Lock()
	defer l2.l.Unlock()

	if pnameLen := uint64(len(pname)); pnameLen > *l2.maxPNameLen {
		*l2.maxPNameLen = pnameLen
	}

	return &l2
}

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

	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 = io.Discard
	l.outBuf = bufio.NewWriter(l.out)

	return err
}

func (l *PmuxLogger) println(line string) {

	l.l.Lock()
	defer l.l.Unlock()

	if l.timeFmt != "" {
		fmt.Fprintf(
			l.outBuf,
			"%s %c ",
			time.Now().Format(l.timeFmt),
			l.sep,
		)
	}

	fmt.Fprintf(
		l.outBuf,
		"%s%s%c %s\n",
		l.pname,
		strings.Repeat(" ", int(*l.maxPNameLen+1)-len(l.pname)),
		l.sep,
		line,
	)

	l.outBuf.Flush()
}

func (l *PmuxLogger) Println(line string) {
	l.println(line)
}

func (l *PmuxLogger) Printf(msg string, args ...interface{}) {
	l.Println(fmt.Sprintf(msg, args...))
}