<commit_before>// Package log connects to a local or remote syslog server with fallback to
// stderr output.
package log

import (
	"fmt"
	"log"
	"log/syslog"
	"os"
	"runtime"
	"strings"

	"github.com/registrobr/gostk/path"
)

// pathDeep defines the number of directories that are visible when logging a
// message with the logging location.
const pathDeep = 3

// Syslog level message, defined in RFC 5424, section 6.2.1
const (
	// LevelEmergency sets a high priority level of problem advising that system
	// is unusable.
	LevelEmergency Level = 0

	// LevelAlert sets a high priority level of problem advising to correct
	// immediately.
	LevelAlert Level = 1

	// LevelCritical sets a medium priority level of problem indicating a failure
	// in a primary system.
	LevelCritical Level = 2

	// LevelError sets a medium priority level of problem indicating a non-urgent
	// failure.
	LevelError Level = 3

	// LevelWarning sets a low priority level indicating that an error will occur
	// if action is not taken.
	LevelWarning Level = 4

	// LevelNotice sets a low priority level indicating events that are unusual,
	// but not error conditions.
	LevelNotice Level = 5

	// LevelInfo sets a very low priority level indicating normal operational
	// messages that require no action.
	LevelInfo Level = 6

	// LevelDebug sets a very low priority level indicating information useful to
	// developers for debugging the application.
	LevelDebug Level = 7
)

// Level defines the severity of an error. For example, if a custom error is
// created as bellow:
//
//    import "github.com/registrobr/gostk/log"
//
//    type ErrDatabaseFailure struct {
//    }
//
//    func (e ErrDatabaseFailure) Error() string {
//      return "database failure!"
//    }
//
//    func (e ErrDatabaseFailure) Level() log.Level {
//      return log.LevelEmergency
//    }
//
//  When used with the Logger type will be written in the syslog in the
//  corresponding log level.
type Level int

type leveler interface {
	Level() Level
}

// syslogWriter is useful to mock a low level syslog writer for unit tests.
type syslogWriter interface {
	Close() error
	Emerg(m string) (err error)
	Alert(m string) (err error)
	Crit(m string) (err error)
	Err(m string) (err error)
	Warning(m string) (err error)
	Notice(m string) (err error)
	Info(m string) (err error)
	Debug(m string) (err error)
}

var (
	remoteLogger syslogWriter
	localLogger  *log.Logger
)

func init() {
	localLogger = log.New(os.Stderr, "", log.LstdFlags)
}

// Dial establishes a connection to a log daemon by connecting to
// address raddr on the specified network.  Each write to the returned
// writer sends a log message with the given facility, severity and
// tag. If network is empty, Dial will connect to the local syslog server.
func Dial(network, raddr, tag string) (err error) {
	remoteLogger, err = syslog.Dial(network, raddr, syslog.LOG_INFO|syslog.LOG_LOCAL0, tag)
	return
}

// Close closes a connection to the syslog daemon.
func Close() error {
	if remoteLogger == nil {
		return nil
	}

	err := remoteLogger.Close()
	if err == nil {
		remoteLogger = nil
	}
	return err
}

// Logger allows logging messages in all different level types. As it is an
// interface it can be replaced by mocks for test purposes.
type Logger interface {
	Emerg(m ...interface{})
	Emergf(m string, a ...interface{})
	Alert(m ...interface{})
	Alertf(m string, a ...interface{})
	Crit(m ...interface{})
	Critf(m string, a ...interface{})
	Error(e error)
	Errorf(m string, a ...interface{})
	Warning(m ...interface{})
	Warningf(m string, a ...interface{})
	Notice(m ...interface{})
	Noticef(m string, a ...interface{})
	Info(m ...interface{})
	Infof(m string, a ...interface{})
	Debug(m ...interface{})
	Debugf(m string, a ...interface{})

	// setCaller defines the number of invocations to follow-up to retrieve the
	// actual caller of the log entry. For now is only used by the package easy
	// functions.
	setCaller(n int)
}

type logger struct {
	identifier string
	caller     int
}

// NewLogger returns a internal instance of the Logger type tagging an
// identifier to every message logged. This identifier is useful to group many
// messages to one related transaction id.
var NewLogger = func(id string) Logger {
	return &logger{
		identifier: "[" + id + "] ",
		caller:     3,
	}
}

func (l logger) Emerg(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Emerg
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Emergf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Emerg
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Alert(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Alert
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Alertf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Alert
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Crit(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Crit
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Critf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Crit
	}
	l.logWithSourceInfof(f, m, a...)
}

// Error converts an Go error into an error message. The responsibility of
// knowing the file and line where the error occurred is from the Error()
// function of the specific error.
func (l logger) Error(e error) {
	if e == nil {
		return
	}

	msg := l.identifier + e.Error()
	if remoteLogger == nil {
		localLogger.Println(msg)
		return
	}

	var err error

	if levelError, ok := e.(leveler); ok {
		switch levelError.Level() {
		case LevelEmergency:
			err = remoteLogger.Emerg(msg)
		case LevelAlert:
			err = remoteLogger.Alert(msg)
		case LevelCritical:
			err = remoteLogger.Crit(msg)
		case LevelError:
			err = remoteLogger.Err(msg)
		case LevelWarning:
			err = remoteLogger.Warning(msg)
		case LevelNotice:
			err = remoteLogger.Notice(msg)
		case LevelInfo:
			err = remoteLogger.Info(msg)
		case LevelDebug:
			err = remoteLogger.Debug(msg)
		default:
			l.Warningf("Wrong error level: %d", levelError.Level())
			err = remoteLogger.Err(msg)
		}
	} else {
		err = remoteLogger.Err(msg)
	}

	if err != nil {
		localLogger.Println("Error writing to syslog. Details:", err)
		localLogger.Println(msg)
	}
}

func (l logger) Errorf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Err
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Warning(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Warning
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Warningf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Warning
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Notice(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Notice
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Noticef(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Notice
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Info(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Info
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Infof(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Info
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Debug(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Debug
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Debugf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Debug
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l *logger) setCaller(n int) {
	l.caller = n
}

// Emerg log an emergency message
func Emerg(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Emerg(a...)
}

// Emergf log an emergency message with arguments
func Emergf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Emergf(m, a...)
}

// Alert log an emergency message
func Alert(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Alert(a...)
}

// Alertf log an emergency message with arguments
func Alertf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Alertf(m, a...)
}

// Crit log an emergency message
func Crit(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Crit(a...)
}

// Critf log an emergency message with arguments
func Critf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Critf(m, a...)
}

// Error log an emergency message
func Error(err error) {
	l := NewLogger("")
	l.setCaller(4)
	l.Error(err)
}

// Errorf log an emergency message with arguments
func Errorf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Errorf(m, a...)
}

// Warning log an emergency message
func Warning(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Warning(a...)
}

// Warningf log an emergency message with arguments
func Warningf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Warningf(m, a...)
}

// Notice log an emergency message
func Notice(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Notice(a...)
}

// Noticef log an emergency message with arguments
func Noticef(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Noticef(m, a...)
}

// Info log an emergency message
func Info(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Info(a...)
}

// Infof log an emergency message with arguments
func Infof(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Infof(m, a...)
}

// Debug log an emergency message
func Debug(a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Debug(a...)
}

// Debugf log an emergency message with arguments
func Debugf(m string, a ...interface{}) {
	l := NewLogger("")
	l.setCaller(4)
	l.Debugf(m, a...)
}

type logFunc func(string) error

func (l logger) logWithSourceInfo(f logFunc, a ...interface{}) {
	// identify the caller from 3 levels above, as this function is never called
	// directly from the place that logged the message
	_, file, line, _ := runtime.Caller(l.caller)
	file = path.RelevantPath(file, pathDeep)
	doLog(f, l.identifier, fmt.Sprint(a...), file, line)
}

func (l logger) logWithSourceInfof(f logFunc, message string, a ...interface{}) {
	// identify the caller from 3 levels above, as this function is never called
	// directly from the place that logged the message
	_, file, line, _ := runtime.Caller(l.caller)
	file = path.RelevantPath(file, pathDeep)
	doLog(f, l.identifier, fmt.Sprintf(message, a...), file, line)
}

func doLog(f logFunc, prefix, message, file string, line int) {
	// support multiline log message, breaking it in many log entries
	for _, item := range strings.Split(message, "\n") {
		if item == "" {
			continue
		}

		msg := fmt.Sprintf("%s%s:%d: %s", prefix, file, line, item)

		if f == nil {
			localLogger.Println(msg)

		} else if err := f(msg); err != nil {
			localLogger.Println("Error writing to syslog. Details:", err)
			localLogger.Println(msg)
		}
	}
}
<commit_msg>Export Logger private methods<commit_after>// Package log connects to a local or remote syslog server with fallback to
// stderr output.
package log

import (
	"fmt"
	"log"
	"log/syslog"
	"os"
	"runtime"
	"strings"

	"github.com/registrobr/gostk/path"
)

// pathDeep defines the number of directories that are visible when logging a
// message with the logging location.
const pathDeep = 3

// Syslog level message, defined in RFC 5424, section 6.2.1
const (
	// LevelEmergency sets a high priority level of problem advising that system
	// is unusable.
	LevelEmergency Level = 0

	// LevelAlert sets a high priority level of problem advising to correct
	// immediately.
	LevelAlert Level = 1

	// LevelCritical sets a medium priority level of problem indicating a failure
	// in a primary system.
	LevelCritical Level = 2

	// LevelError sets a medium priority level of problem indicating a non-urgent
	// failure.
	LevelError Level = 3

	// LevelWarning sets a low priority level indicating that an error will occur
	// if action is not taken.
	LevelWarning Level = 4

	// LevelNotice sets a low priority level indicating events that are unusual,
	// but not error conditions.
	LevelNotice Level = 5

	// LevelInfo sets a very low priority level indicating normal operational
	// messages that require no action.
	LevelInfo Level = 6

	// LevelDebug sets a very low priority level indicating information useful to
	// developers for debugging the application.
	LevelDebug Level = 7
)

// Level defines the severity of an error. For example, if a custom error is
// created as bellow:
//
//    import "github.com/registrobr/gostk/log"
//
//    type ErrDatabaseFailure struct {
//    }
//
//    func (e ErrDatabaseFailure) Error() string {
//      return "database failure!"
//    }
//
//    func (e ErrDatabaseFailure) Level() log.Level {
//      return log.LevelEmergency
//    }
//
//  When used with the Logger type will be written in the syslog in the
//  corresponding log level.
type Level int

type leveler interface {
	Level() Level
}

// syslogWriter is useful to mock a low level syslog writer for unit tests.
type syslogWriter interface {
	Close() error
	Emerg(m string) (err error)
	Alert(m string) (err error)
	Crit(m string) (err error)
	Err(m string) (err error)
	Warning(m string) (err error)
	Notice(m string) (err error)
	Info(m string) (err error)
	Debug(m string) (err error)
}

var (
	remoteLogger syslogWriter
	localLogger  *log.Logger
)

func init() {
	localLogger = log.New(os.Stderr, "", log.LstdFlags)
}

// Dial establishes a connection to a log daemon by connecting to
// address raddr on the specified network.  Each write to the returned
// writer sends a log message with the given facility, severity and
// tag. If network is empty, Dial will connect to the local syslog server.
func Dial(network, raddr, tag string) (err error) {
	remoteLogger, err = syslog.Dial(network, raddr, syslog.LOG_INFO|syslog.LOG_LOCAL0, tag)
	return
}

// Close closes a connection to the syslog daemon.
func Close() error {
	if remoteLogger == nil {
		return nil
	}

	err := remoteLogger.Close()
	if err == nil {
		remoteLogger = nil
	}
	return err
}

// Logger allows logging messages in all different level types. As it is an
// interface it can be replaced by mocks for test purposes.
type Logger interface {
	Emerg(m ...interface{})
	Emergf(m string, a ...interface{})
	Alert(m ...interface{})
	Alertf(m string, a ...interface{})
	Crit(m ...interface{})
	Critf(m string, a ...interface{})
	Error(e error)
	Errorf(m string, a ...interface{})
	Warning(m ...interface{})
	Warningf(m string, a ...interface{})
	Notice(m ...interface{})
	Noticef(m string, a ...interface{})
	Info(m ...interface{})
	Infof(m string, a ...interface{})
	Debug(m ...interface{})
	Debugf(m string, a ...interface{})

	// SetCaller defines the number of invocations to follow-up to retrieve the
	// actual caller of the log entry. For now is only used by the package easy
	// functions.
	SetCaller(n int)
}

type logger struct {
	identifier string
	caller     int
}

// NewLogger returns a internal instance of the Logger type tagging an
// identifier to every message logged. This identifier is useful to group many
// messages to one related transaction id.
var NewLogger = func(id string) Logger {
	return &logger{
		identifier: "[" + id + "] ",
		caller:     3,
	}
}

func (l logger) Emerg(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Emerg
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Emergf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Emerg
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Alert(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Alert
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Alertf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Alert
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Crit(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Crit
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Critf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Crit
	}
	l.logWithSourceInfof(f, m, a...)
}

// Error converts an Go error into an error message. The responsibility of
// knowing the file and line where the error occurred is from the Error()
// function of the specific error.
func (l logger) Error(e error) {
	if e == nil {
		return
	}

	msg := l.identifier + e.Error()
	if remoteLogger == nil {
		localLogger.Println(msg)
		return
	}

	var err error

	if levelError, ok := e.(leveler); ok {
		switch levelError.Level() {
		case LevelEmergency:
			err = remoteLogger.Emerg(msg)
		case LevelAlert:
			err = remoteLogger.Alert(msg)
		case LevelCritical:
			err = remoteLogger.Crit(msg)
		case LevelError:
			err = remoteLogger.Err(msg)
		case LevelWarning:
			err = remoteLogger.Warning(msg)
		case LevelNotice:
			err = remoteLogger.Notice(msg)
		case LevelInfo:
			err = remoteLogger.Info(msg)
		case LevelDebug:
			err = remoteLogger.Debug(msg)
		default:
			l.Warningf("Wrong error level: %d", levelError.Level())
			err = remoteLogger.Err(msg)
		}
	} else {
		err = remoteLogger.Err(msg)
	}

	if err != nil {
		localLogger.Println("Error writing to syslog. Details:", err)
		localLogger.Println(msg)
	}
}

func (l logger) Errorf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Err
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Warning(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Warning
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Warningf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Warning
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Notice(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Notice
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Noticef(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Notice
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Info(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Info
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Infof(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Info
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l logger) Debug(a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Debug
	}
	l.logWithSourceInfo(f, a...)
}

func (l logger) Debugf(m string, a ...interface{}) {
	var f logFunc
	if remoteLogger != nil {
		f = remoteLogger.Debug
	}
	l.logWithSourceInfof(f, m, a...)
}

func (l *logger) SetCaller(n int) {
	l.caller = n
}

// Emerg log an emergency message
func Emerg(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Emerg(a...)
}

// Emergf log an emergency message with arguments
func Emergf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Emergf(m, a...)
}

// Alert log an emergency message
func Alert(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Alert(a...)
}

// Alertf log an emergency message with arguments
func Alertf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Alertf(m, a...)
}

// Crit log an emergency message
func Crit(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Crit(a...)
}

// Critf log an emergency message with arguments
func Critf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Critf(m, a...)
}

// Error log an emergency message
func Error(err error) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Error(err)
}

// Errorf log an emergency message with arguments
func Errorf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Errorf(m, a...)
}

// Warning log an emergency message
func Warning(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Warning(a...)
}

// Warningf log an emergency message with arguments
func Warningf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Warningf(m, a...)
}

// Notice log an emergency message
func Notice(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Notice(a...)
}

// Noticef log an emergency message with arguments
func Noticef(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Noticef(m, a...)
}

// Info log an emergency message
func Info(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Info(a...)
}

// Infof log an emergency message with arguments
func Infof(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Infof(m, a...)
}

// Debug log an emergency message
func Debug(a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Debug(a...)
}

// Debugf log an emergency message with arguments
func Debugf(m string, a ...interface{}) {
	l := NewLogger("")
	l.SetCaller(4)
	l.Debugf(m, a...)
}

type logFunc func(string) error

func (l logger) logWithSourceInfo(f logFunc, a ...interface{}) {
	// identify the caller from 3 levels above, as this function is never called
	// directly from the place that logged the message
	_, file, line, _ := runtime.Caller(l.caller)
	file = path.RelevantPath(file, pathDeep)
	doLog(f, l.identifier, fmt.Sprint(a...), file, line)
}

func (l logger) logWithSourceInfof(f logFunc, message string, a ...interface{}) {
	// identify the caller from 3 levels above, as this function is never called
	// directly from the place that logged the message
	_, file, line, _ := runtime.Caller(l.caller)
	file = path.RelevantPath(file, pathDeep)
	doLog(f, l.identifier, fmt.Sprintf(message, a...), file, line)
}

func doLog(f logFunc, prefix, message, file string, line int) {
	// support multiline log message, breaking it in many log entries
	for _, item := range strings.Split(message, "\n") {
		if item == "" {
			continue
		}

		msg := fmt.Sprintf("%s%s:%d: %s", prefix, file, line, item)

		if f == nil {
			localLogger.Println(msg)

		} else if err := f(msg); err != nil {
			localLogger.Println("Error writing to syslog. Details:", err)
			localLogger.Println(msg)
		}
	}
}
<|endoftext|>


<commit_before>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

func getRealValue(value reflect.Value, columns []string) (results []interface{}) {
	// If value is a nil pointer, Indirect returns a zero Value!
	// Therefor we need to check for a zero value,
	// as FieldByName could panic
	if pointedValue := reflect.Indirect(value); pointedValue.IsValid() {
		for _, column := range columns {
			if pointedValue.FieldByName(column).IsValid() {
				result := pointedValue.FieldByName(column).Interface()
				if r, ok := result.(driver.Valuer); ok {
					result, _ = r.Value()
				}
				results = append(results, result)
			}
		}
	}
	return
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	if scope.Search.preload == nil {
		return
	}

	preloadMap := map[string]bool{}
	fields := scope.Fields()
	for _, preload := range scope.Search.preload {
		schema, conditions := preload.schema, preload.conditions
		keys := strings.Split(schema, ".")
		currentScope := scope
		currentFields := fields
		originalConditions := conditions
		conditions = []interface{}{}
		for i, key := range keys {
			var found bool
			if preloadMap[strings.Join(keys[:i+1], ".")] {
				goto nextLoop
			}

			if i == len(keys)-1 {
				conditions = originalConditions
			}

			for _, field := range currentFields {
				if field.Name != key || field.Relationship == nil {
					continue
				}

				found = true
				switch field.Relationship.Kind {
				case "has_one":
					currentScope.handleHasOnePreload(field, conditions)
				case "has_many":
					currentScope.handleHasManyPreload(field, conditions)
				case "belongs_to":
					currentScope.handleBelongsToPreload(field, conditions)
				case "many_to_many":
					currentScope.handleManyToManyPreload(field, conditions)
				default:
					currentScope.Err(errors.New("not supported relation"))
				}
				break
			}

			if !found {
				value := reflect.ValueOf(currentScope.Value)
				if value.Kind() == reflect.Slice && value.Type().Elem().Kind() == reflect.Interface {
					value = value.Index(0).Elem()
				}
				scope.Err(fmt.Errorf("can't find field %s in %s", key, value.Type()))
				return
			}

			preloadMap[strings.Join(keys[:i+1], ".")] = true

		nextLoop:
			if i < len(keys)-1 {
				currentScope = currentScope.getColumnsAsScope(key)
				currentFields = currentScope.Fields()
			}
		}
	}

}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) handleHasOnePreload(field *Field, conditions []interface{}) {
	relation := field.Relationship

	primaryKeys := scope.getColumnAsArray(relation.AssociationForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.ForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	for i := 0; i < resultValues.Len(); i++ {
		result := resultValues.Index(i)
		if scope.IndirectValue().Kind() == reflect.Slice {
			value := getRealValue(result, relation.ForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				if equalAsString(getRealValue(objects.Index(j), relation.AssociationForeignFieldNames), value) {
					reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
					break
				}
			}
		} else {
			if err := scope.SetColumn(field, result); err != nil {
				scope.Err(err)
				return
			}
		}
	}
}

func (scope *Scope) handleHasManyPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	primaryKeys := scope.getColumnAsArray(relation.AssociationForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.ForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	if scope.IndirectValue().Kind() == reflect.Slice {
		for i := 0; i < resultValues.Len(); i++ {
			result := resultValues.Index(i)
			value := getRealValue(result, relation.ForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				object := reflect.Indirect(objects.Index(j))
				if equalAsString(getRealValue(object, relation.AssociationForeignFieldNames), value) {
					f := object.FieldByName(field.Name)
					f.Set(reflect.Append(f, result))
					break
				}
			}
		}
	} else {
		scope.SetColumn(field, resultValues)
	}
}

func (scope *Scope) handleBelongsToPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	primaryKeys := scope.getColumnAsArray(relation.ForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.AssociationForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	for i := 0; i < resultValues.Len(); i++ {
		result := resultValues.Index(i)
		if scope.IndirectValue().Kind() == reflect.Slice {
			value := getRealValue(result, relation.AssociationForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				object := reflect.Indirect(objects.Index(j))
				if equalAsString(getRealValue(object, relation.ForeignFieldNames), value) {
					object.FieldByName(field.Name).Set(result)
				}
			}
		} else {
			scope.SetColumn(field, result)
		}
	}
}

func (scope *Scope) handleManyToManyPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	joinTableHandler := relation.JoinTableHandler
	destType := field.StructField.Struct.Type.Elem()
	var isPtr bool
	if destType.Kind() == reflect.Ptr {
		isPtr = true
		destType = destType.Elem()
	}

	var sourceKeys []string
	var linkHash = make(map[string][]reflect.Value)

	for _, key := range joinTableHandler.SourceForeignKeys() {
		sourceKeys = append(sourceKeys, key.DBName)
	}

	db := scope.NewDB().Table(scope.New(reflect.New(destType).Interface()).TableName()).Select("*")

	preloadJoinDB := joinTableHandler.JoinWith(joinTableHandler, db, scope.Value)

	if len(conditions) > 0 {
		preloadJoinDB = preloadJoinDB.Where(conditions[0], conditions[1:]...)
	}
	rows, err := preloadJoinDB.Rows()

	if scope.Err(err) != nil {
		return
	}
	defer rows.Close()

	columns, _ := rows.Columns()
	for rows.Next() {
		elem := reflect.New(destType).Elem()
		var values = make([]interface{}, len(columns))

		fields := scope.New(elem.Addr().Interface()).Fields()

		var foundFields = map[string]bool{}
		for index, column := range columns {
			if field, ok := fields[column]; ok && !foundFields[column] {
				if field.Field.Kind() == reflect.Ptr {
					values[index] = field.Field.Addr().Interface()
				} else {
					values[index] = reflect.New(reflect.PtrTo(field.Field.Type())).Interface()
				}
				foundFields[column] = true
			} else {
				var i interface{}
				values[index] = &i
			}
		}

		scope.Err(rows.Scan(values...))

		var sourceKey []interface{}

		var scannedFields = map[string]bool{}
		for index, column := range columns {
			value := values[index]
			if field, ok := fields[column]; ok && !scannedFields[column] {
				if field.Field.Kind() == reflect.Ptr {
					field.Field.Set(reflect.ValueOf(value).Elem())
				} else if v := reflect.ValueOf(value).Elem().Elem(); v.IsValid() {
					field.Field.Set(v)
				}
				scannedFields[column] = true
			} else if strInSlice(column, sourceKeys) {
				sourceKey = append(sourceKey, *(value.(*interface{})))
			}
		}

		if len(sourceKey) != 0 {
			if isPtr {
				linkHash[toString(sourceKey)] = append(linkHash[toString(sourceKey)], elem.Addr())
			} else {
				linkHash[toString(sourceKey)] = append(linkHash[toString(sourceKey)], elem)
			}
		}
	}

	var associationForeignStructFieldNames []string
	for _, dbName := range relation.AssociationForeignFieldNames {
		if field, ok := scope.FieldByName(dbName); ok {
			associationForeignStructFieldNames = append(associationForeignStructFieldNames, field.Name)
		}
	}

	if scope.IndirectValue().Kind() == reflect.Slice {
		objects := scope.IndirectValue()
		for j := 0; j < objects.Len(); j++ {
			object := reflect.Indirect(objects.Index(j))
			source := getRealValue(object, associationForeignStructFieldNames)
			field := object.FieldByName(field.Name)
			for _, link := range linkHash[toString(source)] {
				field.Set(reflect.Append(field, link))
			}
		}
	} else {
		if object := scope.IndirectValue(); object.IsValid() {
			source := getRealValue(object, associationForeignStructFieldNames)
			field := object.FieldByName(field.Name)
			for _, link := range linkHash[toString(source)] {
				field.Set(reflect.Append(field, link))
			}
		}
	}
}

func (scope *Scope) getColumnAsArray(columns []string) (results [][]interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		for i := 0; i < values.Len(); i++ {
			var result []interface{}
			for _, column := range columns {
				result = append(result, reflect.Indirect(values.Index(i)).FieldByName(column).Interface())
			}
			results = append(results, result)
		}
	case reflect.Struct:
		var result []interface{}
		for _, column := range columns {
			result = append(result, values.FieldByName(column).Interface())
		}
		return [][]interface{}{result}
	}
	return
}

func (scope *Scope) getColumnsAsScope(column string) *Scope {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		modelType := values.Type().Elem()
		if modelType.Kind() == reflect.Ptr {
			modelType = modelType.Elem()
		}
		fieldStruct, _ := modelType.FieldByName(column)
		var columns reflect.Value
		if fieldStruct.Type.Kind() == reflect.Slice || fieldStruct.Type.Kind() == reflect.Ptr {
			columns = reflect.New(reflect.SliceOf(reflect.PtrTo(fieldStruct.Type.Elem()))).Elem()
		} else {
			columns = reflect.New(reflect.SliceOf(reflect.PtrTo(fieldStruct.Type))).Elem()
		}
		for i := 0; i < values.Len(); i++ {
			column := reflect.Indirect(values.Index(i)).FieldByName(column)
			if column.Kind() == reflect.Ptr {
				column = column.Elem()
			}
			if column.Kind() == reflect.Slice {
				for i := 0; i < column.Len(); i++ {
					elem := column.Index(i)
					if elem.CanAddr() {
						columns = reflect.Append(columns, elem.Addr())
					}
				}
			} else {
				if column.CanAddr() {
					columns = reflect.Append(columns, column.Addr())
				}
			}
		}
		return scope.New(columns.Interface())
	case reflect.Struct:
		field := values.FieldByName(column)
		if !field.CanAddr() {
			return nil
		}
		return scope.New(field.Addr().Interface())
	}
	return nil
}
<commit_msg>Don't preload if has any error<commit_after>package gorm

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

func getRealValue(value reflect.Value, columns []string) (results []interface{}) {
	// If value is a nil pointer, Indirect returns a zero Value!
	// Therefor we need to check for a zero value,
	// as FieldByName could panic
	if pointedValue := reflect.Indirect(value); pointedValue.IsValid() {
		for _, column := range columns {
			if pointedValue.FieldByName(column).IsValid() {
				result := pointedValue.FieldByName(column).Interface()
				if r, ok := result.(driver.Valuer); ok {
					result, _ = r.Value()
				}
				results = append(results, result)
			}
		}
	}
	return
}

func equalAsString(a interface{}, b interface{}) bool {
	return fmt.Sprintf("%v", a) == fmt.Sprintf("%v", b)
}

func Preload(scope *Scope) {
	if scope.Search.preload == nil || scope.HasError() {
		return
	}

	preloadMap := map[string]bool{}
	fields := scope.Fields()
	for _, preload := range scope.Search.preload {
		schema, conditions := preload.schema, preload.conditions
		keys := strings.Split(schema, ".")
		currentScope := scope
		currentFields := fields
		originalConditions := conditions
		conditions = []interface{}{}
		for i, key := range keys {
			var found bool
			if preloadMap[strings.Join(keys[:i+1], ".")] {
				goto nextLoop
			}

			if i == len(keys)-1 {
				conditions = originalConditions
			}

			for _, field := range currentFields {
				if field.Name != key || field.Relationship == nil {
					continue
				}

				found = true
				switch field.Relationship.Kind {
				case "has_one":
					currentScope.handleHasOnePreload(field, conditions)
				case "has_many":
					currentScope.handleHasManyPreload(field, conditions)
				case "belongs_to":
					currentScope.handleBelongsToPreload(field, conditions)
				case "many_to_many":
					currentScope.handleManyToManyPreload(field, conditions)
				default:
					currentScope.Err(errors.New("not supported relation"))
				}
				break
			}

			if !found {
				value := reflect.ValueOf(currentScope.Value)
				if value.Kind() == reflect.Slice && value.Type().Elem().Kind() == reflect.Interface {
					value = value.Index(0).Elem()
				}
				scope.Err(fmt.Errorf("can't find field %s in %s", key, value.Type()))
				return
			}

			preloadMap[strings.Join(keys[:i+1], ".")] = true

		nextLoop:
			if i < len(keys)-1 {
				currentScope = currentScope.getColumnsAsScope(key)
				currentFields = currentScope.Fields()
			}
		}
	}

}

func makeSlice(typ reflect.Type) interface{} {
	if typ.Kind() == reflect.Slice {
		typ = typ.Elem()
	}
	sliceType := reflect.SliceOf(typ)
	slice := reflect.New(sliceType)
	slice.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
	return slice.Interface()
}

func (scope *Scope) handleHasOnePreload(field *Field, conditions []interface{}) {
	relation := field.Relationship

	primaryKeys := scope.getColumnAsArray(relation.AssociationForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.ForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	for i := 0; i < resultValues.Len(); i++ {
		result := resultValues.Index(i)
		if scope.IndirectValue().Kind() == reflect.Slice {
			value := getRealValue(result, relation.ForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				if equalAsString(getRealValue(objects.Index(j), relation.AssociationForeignFieldNames), value) {
					reflect.Indirect(objects.Index(j)).FieldByName(field.Name).Set(result)
					break
				}
			}
		} else {
			if err := scope.SetColumn(field, result); err != nil {
				scope.Err(err)
				return
			}
		}
	}
}

func (scope *Scope) handleHasManyPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	primaryKeys := scope.getColumnAsArray(relation.AssociationForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.ForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	if scope.IndirectValue().Kind() == reflect.Slice {
		for i := 0; i < resultValues.Len(); i++ {
			result := resultValues.Index(i)
			value := getRealValue(result, relation.ForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				object := reflect.Indirect(objects.Index(j))
				if equalAsString(getRealValue(object, relation.AssociationForeignFieldNames), value) {
					f := object.FieldByName(field.Name)
					f.Set(reflect.Append(f, result))
					break
				}
			}
		}
	} else {
		scope.SetColumn(field, resultValues)
	}
}

func (scope *Scope) handleBelongsToPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	primaryKeys := scope.getColumnAsArray(relation.ForeignFieldNames)
	if len(primaryKeys) == 0 {
		return
	}

	results := makeSlice(field.Struct.Type)
	scope.Err(scope.NewDB().Where(fmt.Sprintf("%v IN (%v)", toQueryCondition(scope, relation.AssociationForeignDBNames), toQueryMarks(primaryKeys)), toQueryValues(primaryKeys)...).Find(results, conditions...).Error)
	resultValues := reflect.Indirect(reflect.ValueOf(results))

	for i := 0; i < resultValues.Len(); i++ {
		result := resultValues.Index(i)
		if scope.IndirectValue().Kind() == reflect.Slice {
			value := getRealValue(result, relation.AssociationForeignFieldNames)
			objects := scope.IndirectValue()
			for j := 0; j < objects.Len(); j++ {
				object := reflect.Indirect(objects.Index(j))
				if equalAsString(getRealValue(object, relation.ForeignFieldNames), value) {
					object.FieldByName(field.Name).Set(result)
				}
			}
		} else {
			scope.SetColumn(field, result)
		}
	}
}

func (scope *Scope) handleManyToManyPreload(field *Field, conditions []interface{}) {
	relation := field.Relationship
	joinTableHandler := relation.JoinTableHandler
	destType := field.StructField.Struct.Type.Elem()
	var isPtr bool
	if destType.Kind() == reflect.Ptr {
		isPtr = true
		destType = destType.Elem()
	}

	var sourceKeys []string
	var linkHash = make(map[string][]reflect.Value)

	for _, key := range joinTableHandler.SourceForeignKeys() {
		sourceKeys = append(sourceKeys, key.DBName)
	}

	db := scope.NewDB().Table(scope.New(reflect.New(destType).Interface()).TableName()).Select("*")

	preloadJoinDB := joinTableHandler.JoinWith(joinTableHandler, db, scope.Value)

	if len(conditions) > 0 {
		preloadJoinDB = preloadJoinDB.Where(conditions[0], conditions[1:]...)
	}
	rows, err := preloadJoinDB.Rows()

	if scope.Err(err) != nil {
		return
	}
	defer rows.Close()

	columns, _ := rows.Columns()
	for rows.Next() {
		elem := reflect.New(destType).Elem()
		var values = make([]interface{}, len(columns))

		fields := scope.New(elem.Addr().Interface()).Fields()

		var foundFields = map[string]bool{}
		for index, column := range columns {
			if field, ok := fields[column]; ok && !foundFields[column] {
				if field.Field.Kind() == reflect.Ptr {
					values[index] = field.Field.Addr().Interface()
				} else {
					values[index] = reflect.New(reflect.PtrTo(field.Field.Type())).Interface()
				}
				foundFields[column] = true
			} else {
				var i interface{}
				values[index] = &i
			}
		}

		scope.Err(rows.Scan(values...))

		var sourceKey []interface{}

		var scannedFields = map[string]bool{}
		for index, column := range columns {
			value := values[index]
			if field, ok := fields[column]; ok && !scannedFields[column] {
				if field.Field.Kind() == reflect.Ptr {
					field.Field.Set(reflect.ValueOf(value).Elem())
				} else if v := reflect.ValueOf(value).Elem().Elem(); v.IsValid() {
					field.Field.Set(v)
				}
				scannedFields[column] = true
			} else if strInSlice(column, sourceKeys) {
				sourceKey = append(sourceKey, *(value.(*interface{})))
			}
		}

		if len(sourceKey) != 0 {
			if isPtr {
				linkHash[toString(sourceKey)] = append(linkHash[toString(sourceKey)], elem.Addr())
			} else {
				linkHash[toString(sourceKey)] = append(linkHash[toString(sourceKey)], elem)
			}
		}
	}

	var associationForeignStructFieldNames []string
	for _, dbName := range relation.AssociationForeignFieldNames {
		if field, ok := scope.FieldByName(dbName); ok {
			associationForeignStructFieldNames = append(associationForeignStructFieldNames, field.Name)
		}
	}

	if scope.IndirectValue().Kind() == reflect.Slice {
		objects := scope.IndirectValue()
		for j := 0; j < objects.Len(); j++ {
			object := reflect.Indirect(objects.Index(j))
			source := getRealValue(object, associationForeignStructFieldNames)
			field := object.FieldByName(field.Name)
			for _, link := range linkHash[toString(source)] {
				field.Set(reflect.Append(field, link))
			}
		}
	} else {
		if object := scope.IndirectValue(); object.IsValid() {
			source := getRealValue(object, associationForeignStructFieldNames)
			field := object.FieldByName(field.Name)
			for _, link := range linkHash[toString(source)] {
				field.Set(reflect.Append(field, link))
			}
		}
	}
}

func (scope *Scope) getColumnAsArray(columns []string) (results [][]interface{}) {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		for i := 0; i < values.Len(); i++ {
			var result []interface{}
			for _, column := range columns {
				result = append(result, reflect.Indirect(values.Index(i)).FieldByName(column).Interface())
			}
			results = append(results, result)
		}
	case reflect.Struct:
		var result []interface{}
		for _, column := range columns {
			result = append(result, values.FieldByName(column).Interface())
		}
		return [][]interface{}{result}
	}
	return
}

func (scope *Scope) getColumnsAsScope(column string) *Scope {
	values := scope.IndirectValue()
	switch values.Kind() {
	case reflect.Slice:
		modelType := values.Type().Elem()
		if modelType.Kind() == reflect.Ptr {
			modelType = modelType.Elem()
		}
		fieldStruct, _ := modelType.FieldByName(column)
		var columns reflect.Value
		if fieldStruct.Type.Kind() == reflect.Slice || fieldStruct.Type.Kind() == reflect.Ptr {
			columns = reflect.New(reflect.SliceOf(reflect.PtrTo(fieldStruct.Type.Elem()))).Elem()
		} else {
			columns = reflect.New(reflect.SliceOf(reflect.PtrTo(fieldStruct.Type))).Elem()
		}
		for i := 0; i < values.Len(); i++ {
			column := reflect.Indirect(values.Index(i)).FieldByName(column)
			if column.Kind() == reflect.Ptr {
				column = column.Elem()
			}
			if column.Kind() == reflect.Slice {
				for i := 0; i < column.Len(); i++ {
					elem := column.Index(i)
					if elem.CanAddr() {
						columns = reflect.Append(columns, elem.Addr())
					}
				}
			} else {
				if column.CanAddr() {
					columns = reflect.Append(columns, column.Addr())
				}
			}
		}
		return scope.New(columns.Interface())
	case reflect.Struct:
		field := values.FieldByName(column)
		if !field.CanAddr() {
			return nil
		}
		return scope.New(field.Addr().Interface())
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
)

var cmdPrepare = &Command{
	Run:       runPrepare,
	UsageLine: "prepare",
	Short:     "prepares everything for building and running",
	Long:      ``,
}

func runPrepare(cmd *Command, args []string) {
	config, err := parseConfiguration()

	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		return
	}

	for i := range config.Prepare {
		c := exec.Command("bash", "-c", config.Prepare[i])

		c.Stdin = os.Stdin
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr

		if err = c.Start(); err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
			return
		}

		if err = c.Wait(); err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
		}
	}
}
<commit_msg>Make prepare exit with non zero on actual error<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
)

var cmdPrepare = &Command{
	Run:       runPrepare,
	UsageLine: "prepare",
	Short:     "prepares everything for building and running",
	Long:      ``,
}

func runPrepare(cmd *Command, args []string) {
	config, err := parseConfiguration()

	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		return
	}

	for i := range config.Prepare {
		c := exec.Command("bash", "-c", config.Prepare[i])

		c.Stdin = os.Stdin
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr

		if err = c.Start(); err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
			return
		}

		if err = c.Wait(); err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		}
	}
}
<|endoftext|>


<commit_before>package gocb

import (
	"fmt"
	"gopkg.in/couchbase/gocbcore.v7"
	"log"
	"strings"
)

// LogLevel specifies the severity of a log message.
type LogLevel gocbcore.LogLevel

// Various logging levels (or subsystems) which can categorize the message.
// Currently these are ordered in decreasing severity.
const (
	LogError        = LogLevel(gocbcore.LogError)
	LogWarn         = LogLevel(gocbcore.LogWarn)
	LogInfo         = LogLevel(gocbcore.LogInfo)
	LogDebug        = LogLevel(gocbcore.LogDebug)
	LogTrace        = LogLevel(gocbcore.LogTrace)
	LogSched        = LogLevel(gocbcore.LogSched)
	LogMaxVerbosity = LogLevel(gocbcore.LogMaxVerbosity)
)

// Logger defines a logging interface. You can either use one of the default loggers
// (DefaultStdioLogger(), VerboseStdioLogger()) or implement your own.
type Logger interface {
	// Outputs logging information:
	// level is the verbosity level
	// offset is the position within the calling stack from which the message
	// originated. This is useful for contextual loggers which retrieve file/line
	// information.
	Log(level LogLevel, offset int, format string, v ...interface{}) error
}

var (
	globalLogger Logger
)

type coreLogWrapper struct {
	wrapped gocbcore.Logger
}

func (wrapper coreLogWrapper) Log(level LogLevel, offset int, format string, v ...interface{}) error {
	return wrapper.wrapped.Log(gocbcore.LogLevel(level), offset+2, format, v...)
}

// DefaultStdioLogger gets the default standard I/O logger.
//  gocb.SetLogger(gocb.DefaultStdioLogger())
func DefaultStdioLogger() Logger {
	return &coreLogWrapper{
		wrapped: gocbcore.DefaultStdioLogger(),
	}
}

// VerboseStdioLogger is a more verbose level of DefaultStdioLogger(). Messages
// pertaining to the scheduling of ordinary commands (and their responses) will
// also be emitted.
//  gocb.SetLogger(gocb.VerboseStdioLogger())
func VerboseStdioLogger() Logger {
	return coreLogWrapper{
		wrapped: gocbcore.VerboseStdioLogger(),
	}
}

type coreLogger struct {
	wrapped Logger
}

func (wrapper coreLogger) Log(level gocbcore.LogLevel, offset int, format string, v ...interface{}) error {
	return wrapper.wrapped.Log(LogLevel(level), offset+2, format, v...)
}

func getCoreLogger(logger Logger) gocbcore.Logger {
	typedLogger, isCoreLogger := logger.(*coreLogWrapper)
	if isCoreLogger {
		return typedLogger.wrapped
	}

	return &coreLogger{
		wrapped: logger,
	}
}

// SetLogger sets a logger to be used by the library. A logger can be obtained via
// the DefaultStdioLogger() or VerboseStdioLogger() functions. You can also implement
// your own logger using the Logger interface.
func SetLogger(logger Logger) {
	globalLogger = logger
	gocbcore.SetLogger(getCoreLogger(logger))
}

func logExf(level LogLevel, offset int, format string, v ...interface{}) {
	if globalLogger != nil {
		err := globalLogger.Log(level, offset+1, format, v...)
		if err != nil {
			log.Printf("Logger error occurred (%s)\n", err)
		}
	}
}

func logDebugf(format string, v ...interface{}) {
	logExf(LogDebug, 1, format, v...)
}

func logSchedf(format string, v ...interface{}) {
	logExf(LogSched, 1, format, v...)
}

func logWarnf(format string, v ...interface{}) {
	logExf(LogWarn, 1, format, v...)
}

func logErrorf(format string, v ...interface{}) {
	logExf(LogError, 1, format, v...)
}

func reindentLog(indent, message string) string {
	reindentedMessage := strings.Replace(message, "\n", "\n"+indent, -1)
	return fmt.Sprintf("%s%s", indent, reindentedMessage)
}
<commit_msg>GOCBC-257: Implement log redaction.<commit_after>package gocb

import (
	"fmt"
	"gopkg.in/couchbase/gocbcore.v7"
	"log"
	"strings"
)

// LogLevel specifies the severity of a log message.
type LogLevel gocbcore.LogLevel

// Various logging levels (or subsystems) which can categorize the message.
// Currently these are ordered in decreasing severity.
const (
	LogError        = LogLevel(gocbcore.LogError)
	LogWarn         = LogLevel(gocbcore.LogWarn)
	LogInfo         = LogLevel(gocbcore.LogInfo)
	LogDebug        = LogLevel(gocbcore.LogDebug)
	LogTrace        = LogLevel(gocbcore.LogTrace)
	LogSched        = LogLevel(gocbcore.LogSched)
	LogMaxVerbosity = LogLevel(gocbcore.LogMaxVerbosity)
)

// LogRedactLevel specifies the degree with which to redact the logs.
type LogRedactLevel int

const (
	// RedactNone indicates to perform no redactions
	RedactNone = LogRedactLevel(0)

	// RedactPartial indicates to redact all possible user-identifying information from logs.
	RedactPartial = LogRedactLevel(1)

	// RedactFull indicates to fully redact all possible identifying information from logs.
	RedactFull = LogRedactLevel(1)
)

// SetLogRedactionLevel specifies the level with which logs should be redacted.
func SetLogRedactionLevel(level LogRedactLevel) {
	// We don't current log any data that falls under our current redaction rules.
	// This function is included as a stub for future implementations of log redaction
	// that act at a higher level and may need to perform actual redaction's.
}

// Logger defines a logging interface. You can either use one of the default loggers
// (DefaultStdioLogger(), VerboseStdioLogger()) or implement your own.
type Logger interface {
	// Outputs logging information:
	// level is the verbosity level
	// offset is the position within the calling stack from which the message
	// originated. This is useful for contextual loggers which retrieve file/line
	// information.
	Log(level LogLevel, offset int, format string, v ...interface{}) error
}

var (
	globalLogger Logger
)

type coreLogWrapper struct {
	wrapped gocbcore.Logger
}

func (wrapper coreLogWrapper) Log(level LogLevel, offset int, format string, v ...interface{}) error {
	return wrapper.wrapped.Log(gocbcore.LogLevel(level), offset+2, format, v...)
}

// DefaultStdioLogger gets the default standard I/O logger.
//  gocb.SetLogger(gocb.DefaultStdioLogger())
func DefaultStdioLogger() Logger {
	return &coreLogWrapper{
		wrapped: gocbcore.DefaultStdioLogger(),
	}
}

// VerboseStdioLogger is a more verbose level of DefaultStdioLogger(). Messages
// pertaining to the scheduling of ordinary commands (and their responses) will
// also be emitted.
//  gocb.SetLogger(gocb.VerboseStdioLogger())
func VerboseStdioLogger() Logger {
	return coreLogWrapper{
		wrapped: gocbcore.VerboseStdioLogger(),
	}
}

type coreLogger struct {
	wrapped Logger
}

func (wrapper coreLogger) Log(level gocbcore.LogLevel, offset int, format string, v ...interface{}) error {
	return wrapper.wrapped.Log(LogLevel(level), offset+2, format, v...)
}

func getCoreLogger(logger Logger) gocbcore.Logger {
	typedLogger, isCoreLogger := logger.(*coreLogWrapper)
	if isCoreLogger {
		return typedLogger.wrapped
	}

	return &coreLogger{
		wrapped: logger,
	}
}

// SetLogger sets a logger to be used by the library. A logger can be obtained via
// the DefaultStdioLogger() or VerboseStdioLogger() functions. You can also implement
// your own logger using the Logger interface.
func SetLogger(logger Logger) {
	globalLogger = logger
	gocbcore.SetLogger(getCoreLogger(logger))
}

func logExf(level LogLevel, offset int, format string, v ...interface{}) {
	if globalLogger != nil {
		err := globalLogger.Log(level, offset+1, format, v...)
		if err != nil {
			log.Printf("Logger error occurred (%s)\n", err)
		}
	}
}

func logInfof(format string, v ...interface{}) {
	logExf(LogInfo, 1, format, v...)
}

func logDebugf(format string, v ...interface{}) {
	logExf(LogDebug, 1, format, v...)
}

func logSchedf(format string, v ...interface{}) {
	logExf(LogSched, 1, format, v...)
}

func logWarnf(format string, v ...interface{}) {
	logExf(LogWarn, 1, format, v...)
}

func logErrorf(format string, v ...interface{}) {
	logExf(LogError, 1, format, v...)
}

func reindentLog(indent, message string) string {
	reindentedMessage := strings.Replace(message, "\n", "\n"+indent, -1)
	return fmt.Sprintf("%s%s", indent, reindentedMessage)
}
<|endoftext|>


<commit_before>/* vim: set sw=4 sts=4 et foldmethod=syntax : */

/*
 * Copyright (c) 2011 Alexander Færøy <ahf@0x90.dk>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
    "crypto/rand"
    "crypto/tls"
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "strconv"
    "strings"
    "time"
)
type Ircd struct {
    *log.Logger
    listeners []Listener
    config *ConfigurationFile

    motd_file string
    motd_content []string

    clientRegistry *ClientRegistry
}

func NewIrcd() *Ircd {
    ircd := new(Ircd)
    ircd.Logger = log.New(os.Stderr, "", log.Ldate | log.Ltime)
    ircd.listeners = make([]Listener, 0)
    ircd.clientRegistry = NewClientRegistry()

    return ircd
}

func (this *Ircd) SetConfigurationFile(config *ConfigurationFile) {
    this.config = config

    for i := range this.config.Ircd.Listeners {
        listener := this.config.Ircd.Listeners[i]
        hostport := listener.Host + ":" + strconv.Itoa(listener.Port)
        protocol := ProtocolFromString(listener.Type)

        if protocol == nil {
            this.Printf("Unknown protocol type: %s\n", listener.Type)
            continue
        }

        if listener.Tls {
            this.addSecureListener(*protocol, hostport)
        } else {
            this.addListener(*protocol, hostport)
        }
    }
}

func (this *Ircd) addCommonListener(p Protocol, address string, config *tls.Config) {
    var listener Listener

    switch p {
        case TCP: listener = NewTCPListener(this, address, config)
        case WebSocket: listener = NewWebSocketListener(this, address, config)
        default: panic("Unhandled Protocol.")
    }

    if listener != nil {
        this.listeners = append(this.listeners, listener)
    }
}

func (this *Ircd) addListener(protocol Protocol, address string) {
    this.addCommonListener(protocol, address, nil)
}

func (this *Ircd) addSecureListener(protocol Protocol, address string) {
    cert := this.config.Ircd.ServerInfo.Tls.Certificate
    key := this.config.Ircd.ServerInfo.Tls.Key
    errorMessage := fmt.Sprintf("Unable to add secure listener for %s", address)

    if cert == "" {
        this.Printf("%s: %s", errorMessage, "Empty TLS certificate in configuration file.")
        return
    }

    if key == "" {
        this.Printf("%s: %s", errorMessage, "Empty TLS key in configuration file.")
        return
    }

    certificate, error := tls.LoadX509KeyPair(cert, key)

    if error != nil {
        this.Printf("Error Loading Certificate: %s", error)
        return
    }

    config := &tls.Config{
        Rand: rand.Reader,
        Time: time.Seconds,
    }

    config.Certificates = make([]tls.Certificate, 1)
    config.Certificates[0] = certificate

    if protocol == WebSocket {
        config.NextProtos = []string{"http/1.1"}
    }

    this.addCommonListener(protocol, address, config)
}

func (this *Ircd) Run() {
    if len(this.listeners) == 0 {
        fmt.Printf("Error: No Listeners Defined...\n")
        os.Exit(1)
    }

    this.Printf("Opening up for incoming connections")

    for i := range this.listeners {
        listener := this.listeners[i]

        this.Printf("Listening on %s (%s %s)", listener.Address(), listener.Secure(), listener.Protocol())
        go listener.Listen()
    }
}

func (this *Ircd) Me() string {
    return this.config.Ircd.ServerInfo.Name
}

func (this *Ircd) Description() string {
    return this.config.Ircd.ServerInfo.Description
}

func (this *Ircd) SetMotdFile(path string) {
    this.motd_file = path

    this.LoadMotd()
}

func (this *Ircd) LoadMotd() {
    content, error := ioutil.ReadFile(this.motd_file)

    if error != nil {
        this.Printf("Unable to load MOTD file: %s", error)
        return
    }

    this.motd_content = strings.Split(string(content), "\n", -1)
}
<commit_msg>Naming.<commit_after>/* vim: set sw=4 sts=4 et foldmethod=syntax : */

/*
 * Copyright (c) 2011 Alexander Færøy <ahf@0x90.dk>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
    "crypto/rand"
    "crypto/tls"
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "strconv"
    "strings"
    "time"
)
type Ircd struct {
    *log.Logger
    listeners []Listener
    config *ConfigurationFile

    motdFile string
    motdContent []string

    clientRegistry *ClientRegistry
}

func NewIrcd() *Ircd {
    ircd := new(Ircd)
    ircd.Logger = log.New(os.Stderr, "", log.Ldate | log.Ltime)
    ircd.listeners = make([]Listener, 0)
    ircd.clientRegistry = NewClientRegistry()

    return ircd
}

func (this *Ircd) SetConfigurationFile(config *ConfigurationFile) {
    this.config = config

    for i := range this.config.Ircd.Listeners {
        listener := this.config.Ircd.Listeners[i]
        hostport := listener.Host + ":" + strconv.Itoa(listener.Port)
        protocol := ProtocolFromString(listener.Type)

        if protocol == nil {
            this.Printf("Unknown protocol type: %s\n", listener.Type)
            continue
        }

        if listener.Tls {
            this.addSecureListener(*protocol, hostport)
        } else {
            this.addListener(*protocol, hostport)
        }
    }
}

func (this *Ircd) addCommonListener(p Protocol, address string, config *tls.Config) {
    var listener Listener

    switch p {
        case TCP: listener = NewTCPListener(this, address, config)
        case WebSocket: listener = NewWebSocketListener(this, address, config)
        default: panic("Unhandled Protocol.")
    }

    if listener != nil {
        this.listeners = append(this.listeners, listener)
    }
}

func (this *Ircd) addListener(protocol Protocol, address string) {
    this.addCommonListener(protocol, address, nil)
}

func (this *Ircd) addSecureListener(protocol Protocol, address string) {
    cert := this.config.Ircd.ServerInfo.Tls.Certificate
    key := this.config.Ircd.ServerInfo.Tls.Key
    errorMessage := fmt.Sprintf("Unable to add secure listener for %s", address)

    if cert == "" {
        this.Printf("%s: %s", errorMessage, "Empty TLS certificate in configuration file.")
        return
    }

    if key == "" {
        this.Printf("%s: %s", errorMessage, "Empty TLS key in configuration file.")
        return
    }

    certificate, error := tls.LoadX509KeyPair(cert, key)

    if error != nil {
        this.Printf("Error Loading Certificate: %s", error)
        return
    }

    config := &tls.Config{
        Rand: rand.Reader,
        Time: time.Seconds,
    }

    config.Certificates = make([]tls.Certificate, 1)
    config.Certificates[0] = certificate

    if protocol == WebSocket {
        config.NextProtos = []string{"http/1.1"}
    }

    this.addCommonListener(protocol, address, config)
}

func (this *Ircd) Run() {
    if len(this.listeners) == 0 {
        fmt.Printf("Error: No Listeners Defined...\n")
        os.Exit(1)
    }

    this.Printf("Opening up for incoming connections")

    for i := range this.listeners {
        listener := this.listeners[i]

        this.Printf("Listening on %s (%s %s)", listener.Address(), listener.Secure(), listener.Protocol())
        go listener.Listen()
    }
}

func (this *Ircd) Me() string {
    return this.config.Ircd.ServerInfo.Name
}

func (this *Ircd) Description() string {
    return this.config.Ircd.ServerInfo.Description
}

func (this *Ircd) SetMotdFile(path string) {
    this.motdFile = path

    this.LoadMotd()
}

func (this *Ircd) LoadMotd() {
    content, error := ioutil.ReadFile(this.motdFile)

    if error != nil {
        this.Printf("Unable to load MOTD file: %s", error)
        return
    }

    this.motdContent = strings.Split(string(content), "\n", -1)
}
<|endoftext|>


<commit_before>// Dinit is a mini init replacement useful for use inside Docker containers.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	verbose              bool
	port, sleep          int
	namespace, subsystem string
)

func main() {
	flag.IntVar(&port, "port", envInt("DINIT_PORT", 0), "port to export metricss for prometheus (DINIT_PORT)")
	flag.IntVar(&sleep, "sleep", envInt("DINIT_SLEEP", 5), "how many seconds to sleep before force killing programs (DINIT_SLEEP)")
	flag.StringVar(&namespace, "namespace", envString("DINIT_NAMESPACE", ""), "namespace to use for prometheus (DINIT_NAMESPACE)")
	flag.StringVar(&subsystem, "subsystem", envString("DINIT_SUBSYSTEM", ""), "subsystem to use for prometheus (DINIT_SUBSYSTEM)")
	flag.BoolVar(&verbose, "verbose", envBool("DINIT_VERBOSE", false), "be more verbose (DINIT_VERBOSE)")

	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: dinit [OPTION]... PROGRAM [PROGRAM]...")
		fmt.Fprintln(os.Stderr, "Start PROGRAMs by passing the enviroment and reap any zombies.\n")
		flag.PrintDefaults()
	}
	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatal("dinit: need at least one program")
	}

	if port > 0 {
		metrics()
	}

	cmds := []*exec.Cmd{}
	done := make(chan bool)

	for _, arg := range flag.Args() {
		args := strings.Fields(arg) // Split on spaces and execute.
		cmd := exec.Command(args[0], args[1:]...)
		cmds = append(cmds, cmd)

		go func() {
			err := cmd.Start()
			if err != nil {
				log.Fatal(err)
			}
			logf("dinit: pid %d started: %v", cmd.Process.Pid, cmd.Args)

			err = cmd.Wait()
			if err != nil {
				logf("dinit: pid %d, finished with error: %s", cmd.Process.Pid, err)
			} else {
				logf("dinit: pid %d, finished: %v", cmd.Process.Pid, cmd.Args)
			}
			done <- true
		}()
	}

	ints := make(chan os.Signal)
	chld := make(chan os.Signal)
	signal.Notify(ints, syscall.SIGINT, syscall.SIGTERM)
	signal.Notify(chld, syscall.SIGCHLD)

	i := 0
Wait:
	for {
		select {
		case <-chld:
			go reaper()
		case <-done:
			i++
			if len(cmds) == i {
				reaper()
				break Wait
			}
		case sig := <-ints:
			// There is a race here, because the process could have died, we don't care.
			for _, cmd := range cmds {
				logf("dinit: signal %d sent to pid %d", sig, cmd.Process.Pid)
				cmd.Process.Signal(sig)
			}

			time.Sleep(time.Duration(sleep) * time.Second)

			kill := []*os.Process{}
			for _, cmd := range cmds {
				if p, err := os.FindProcess(cmd.Process.Pid); err != nil {
					kill = append(kill, p)
				}
			}
			for _, p := range kill {
				logf("dinit: SIGKILL sent to pid %d", p.Pid)
				p.Signal(syscall.SIGKILL)
			}
		}
	}
}

func reaper() {
	for {
		var wstatus syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &wstatus, 0, nil)
		if err != nil {
			return
		}
		logf("dinit: pid %d reaped", pid)
		zombies.Inc()
	}
}

func logf(format string, v ...interface{}) {
	if !verbose {
		return
	}
	log.Printf("dinit: " + format, v...)
}

func envBool(k string, d bool) bool {
	x := os.Getenv(k)
	switch strings.ToLower(x) {
	case "true":
		return true
	case "false":
		return false
	}
	return d

}

func envInt(k string, d int) int {
	x := os.Getenv(k)
	if x != "" {
		if x1, e := strconv.Atoi(x); e != nil {
			return x1
		}
	}
	return d
}

func envString(k, d string) string {
	x := os.Getenv(k)
	if x != "" {
		return x
	}
	return d
}
<commit_msg>Call reaper() with defer<commit_after>// Dinit is a mini init replacement useful for use inside Docker containers.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	verbose              bool
	port, sleep          int
	namespace, subsystem string
)

func main() {
	flag.IntVar(&port, "port", envInt("DINIT_PORT", 0), "port to export metricss for prometheus (DINIT_PORT)")
	flag.IntVar(&sleep, "sleep", envInt("DINIT_SLEEP", 5), "how many seconds to sleep before force killing programs (DINIT_SLEEP)")
	flag.StringVar(&namespace, "namespace", envString("DINIT_NAMESPACE", ""), "namespace to use for prometheus (DINIT_NAMESPACE)")
	flag.StringVar(&subsystem, "subsystem", envString("DINIT_SUBSYSTEM", ""), "subsystem to use for prometheus (DINIT_SUBSYSTEM)")
	flag.BoolVar(&verbose, "verbose", envBool("DINIT_VERBOSE", false), "be more verbose (DINIT_VERBOSE)")

	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: dinit [OPTION]... PROGRAM [PROGRAM]...")
		fmt.Fprintln(os.Stderr, "Start PROGRAMs by passing the enviroment and reap any zombies.\n")
		flag.PrintDefaults()
	}
	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatal("dinit: need at least one program")
	}

	if port > 0 {
		metrics()
	}

	cmds := []*exec.Cmd{}
	done := make(chan bool)

	for _, arg := range flag.Args() {
		args := strings.Fields(arg) // Split on spaces and execute.
		cmd := exec.Command(args[0], args[1:]...)
		cmds = append(cmds, cmd)

		go func() {
			err := cmd.Start()
			if err != nil {
				log.Fatal(err)
			}
			logf("dinit: pid %d started: %v", cmd.Process.Pid, cmd.Args)

			err = cmd.Wait()
			if err != nil {
				logf("dinit: pid %d, finished with error: %s", cmd.Process.Pid, err)
			} else {
				logf("dinit: pid %d, finished: %v", cmd.Process.Pid, cmd.Args)
			}
			done <- true
		}()
	}

	ints := make(chan os.Signal)
	chld := make(chan os.Signal)
	signal.Notify(ints, syscall.SIGINT, syscall.SIGTERM)
	signal.Notify(chld, syscall.SIGCHLD)

	i := 0
	defer reaper()
Wait:
	for {
		select {
		case <-chld:
			go reaper()
		case <-done:
			i++
			if len(cmds) == i {
				break Wait
			}
		case sig := <-ints:
			// There is a race here, because the process could have died, we don't care.
			for _, cmd := range cmds {
				logf("dinit: signal %d sent to pid %d", sig, cmd.Process.Pid)
				cmd.Process.Signal(sig)
			}

			time.Sleep(time.Duration(sleep) * time.Second)

			kill := []*os.Process{}
			for _, cmd := range cmds {
				if p, err := os.FindProcess(cmd.Process.Pid); err != nil {
					kill = append(kill, p)
				}
			}
			for _, p := range kill {
				logf("dinit: SIGKILL sent to pid %d", p.Pid)
				p.Signal(syscall.SIGKILL)
			}
		}
	}
}

func reaper() {
	for {
		var wstatus syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &wstatus, 0, nil)
		if err != nil {
			return
		}
		logf("dinit: pid %d reaped", pid)
		zombies.Inc()
	}
}

func logf(format string, v ...interface{}) {
	if !verbose {
		return
	}
	log.Printf("dinit: " + format, v...)
}

func envBool(k string, d bool) bool {
	x := os.Getenv(k)
	switch strings.ToLower(x) {
	case "true":
		return true
	case "false":
		return false
	}
	return d

}

func envInt(k string, d int) int {
	x := os.Getenv(k)
	if x != "" {
		if x1, e := strconv.Atoi(x); e != nil {
			return x1
		}
	}
	return d
}

func envString(k, d string) string {
	x := os.Getenv(k)
	if x != "" {
		return x
	}
	return d
}
<|endoftext|>


<commit_before>// Package script is a library facilitating the creation of programs that resemble
// bash scripts.
package script

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

// ProcessResult contains the results of a process execution be it successful or not.
type ProcessResult struct {
	Cmd          *exec.Cmd
	ProcessState *os.ProcessState
	ProcessError error
	stdoutBuffer *bytes.Buffer
	stderrBuffer *bytes.Buffer
}

// NewProcessResult creates a new empty ProcessResult
func NewProcessResult() *ProcessResult {
	p := &ProcessResult{}
	p.stdoutBuffer = bytes.NewBuffer(make([]byte, 0, 100))
	p.stderrBuffer = bytes.NewBuffer(make([]byte, 0, 100))
	return p
}

// Output returns a string representation of the output of the process denoted
// by this struct.
func (pr *ProcessResult) Output() string {
	return pr.stdoutBuffer.String()
}

// Error returns a string representation of the stderr output of the process denoted
// by this struct.
func (pr *ProcessResult) Error() string {
	return pr.stderrBuffer.String()
}

// Successful returns true iff the process denoted by this struct was run
// successfully. Success is defined as the exit code being set to 0.
func (pr *ProcessResult) Successful() bool {
	fmt.Println(pr.ExitCode())
	return pr.ExitCode() == 0
}

// StateString returns a string representation of the process denoted by
// this struct
func (pr *ProcessResult) StateString() string {
	state := pr.ProcessState
	return fmt.Sprintf("PID: %q, Exited: %t, Exit Code: %q, Success: %t, User Time: %q", state.Pid(), state.Exited(), pr.ExitCode(), state.Success(), state.UserTime())
}

// ExitCode returns the exit code of the command denoted by this struct
func (pr *ProcessResult) ExitCode() int {
	var waitStatus syscall.WaitStatus
	if exitError, ok := pr.ProcessError.(*exec.ExitError); ok {
		waitStatus = exitError.Sys().(syscall.WaitStatus)
	} else {
		waitStatus = pr.ProcessState.Sys().(syscall.WaitStatus)
	}
	return waitStatus.ExitStatus()
}

// CommandPath finds the full path of a binary given its name.
func (c *Context) CommandPath(name string) string {
	cmd := exec.Command("which", name)
	cmdOutput, err := cmd.Output()
	if err != nil {
		return ""
	}
	return strings.Trim(string(cmdOutput), "\n")
}

// CommandExists checks if a given binary exists in PATH.
func (c *Context) CommandExists(name string) bool {
	return c.CommandPath(name) != ""
}

// MustCommandExist ensures a given binary exists in PATH, otherwise panics.
func (c *Context) MustCommandExist(name string) {
	if !c.CommandExists(name) {
		panic(fmt.Errorf("Command %s is not available. Please make sure it is installed and accessible.", name))
	}
}

// ExecuteDebug executes a system command, stdout and stderr are piped
func (c *Context) ExecuteDebug(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(false, false, name, args...)
	return
}

// ExecuteSilent executes a  system command without outputting stdout (it is
// still captured and can be retrieved using LastOutput())
func (c *Context) ExecuteSilent(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(true, false, name, args...)
	return
}

// ExecuteFullySilent executes a system command without outputting stdout or
// stderr (both are still captured and can be retrieved using LastOutput() and
// LastError())
func (c *Context) ExecuteFullySilent(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(true, true, name, args...)
	return
}

// MustExecuteDebug ensures a system command to be executed, otherwise panics
func (c *Context) MustExecuteDebug(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.Execute(false, false, name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// MustExecuteSilent ensures a system command to be executed without outputting
// stdout, otherwise panics
func (c *Context) MustExecuteSilent(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.ExecuteSilent(name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// MustExecuteFullySilent ensures a system command to be executed without
// outputting stdout and stderr, otherwise panics
func (c *Context) MustExecuteFullySilent(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.ExecuteFullySilent(name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// Execute executes a system command with configurable stdout and stderr output
func (c *Context) Execute(stdoutSilent bool, stderrSilent bool, name string, args ...string) (pr *ProcessResult, err error) {
	cmd, pr := c.prepareCommand(stdoutSilent, stderrSilent, name, args...)

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	pr.ProcessError = err

	return pr, err
}

// ExecuteDetached executes the given command in this context in the background (detached). This means the script execution instantly continues.
func (c *Context) ExecuteDetached(name string, args ...string) (cmd *exec.Cmd, pr *ProcessResult, err error) {
	cmd, pr = c.prepareCommand(true, true, name, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}
	err = cmd.Start()
	return
}

func (c Context) prepareCommand(stdoutSilent bool, stderrSilent bool, name string, args ...string) (*exec.Cmd, *ProcessResult) {
	pr := NewProcessResult()

	cmd := exec.Command(name, args...)
	pr.Cmd = cmd
	pr.ProcessState = cmd.ProcessState

	cmd.Dir = c.workingDir
	cmd.Env = c.getFullEnv()

	if stderrSilent {
		cmd.Stderr = pr.stderrBuffer
	} else {
		cmd.Stderr = io.MultiWriter(os.Stderr, pr.stderrBuffer)
	}
	if stderrSilent {
		cmd.Stdout = pr.stdoutBuffer
	} else {
		cmd.Stdout = io.MultiWriter(os.Stdout, pr.stdoutBuffer)
	}
	return cmd, pr
}
<commit_msg>Improve documentation<commit_after>// Package script is a library facilitating the creation of programs that resemble
// bash scripts.
package script

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

// ProcessResult contains the results of a process execution be it successful or not.
type ProcessResult struct {
	Cmd          *exec.Cmd
	ProcessState *os.ProcessState
	ProcessError error
	stdoutBuffer *bytes.Buffer
	stderrBuffer *bytes.Buffer
}

// NewProcessResult creates a new empty ProcessResult
func NewProcessResult() *ProcessResult {
	p := &ProcessResult{}
	p.stdoutBuffer = bytes.NewBuffer(make([]byte, 0, 100))
	p.stderrBuffer = bytes.NewBuffer(make([]byte, 0, 100))
	return p
}

// Output returns a string representation of the output of the process denoted
// by this struct.
func (pr *ProcessResult) Output() string {
	return pr.stdoutBuffer.String()
}

// Error returns a string representation of the stderr output of the process denoted
// by this struct.
func (pr *ProcessResult) Error() string {
	return pr.stderrBuffer.String()
}

// Successful returns true iff the process denoted by this struct was run
// successfully. Success is defined as the exit code being set to 0.
func (pr *ProcessResult) Successful() bool {
	fmt.Println(pr.ExitCode())
	return pr.ExitCode() == 0
}

// StateString returns a string representation of the process denoted by
// this struct
func (pr *ProcessResult) StateString() string {
	state := pr.ProcessState
	return fmt.Sprintf("PID: %q, Exited: %t, Exit Code: %q, Success: %t, User Time: %q", state.Pid(), state.Exited(), pr.ExitCode(), state.Success(), state.UserTime())
}

// ExitCode returns the exit code of the command denoted by this struct
func (pr *ProcessResult) ExitCode() int {
	var waitStatus syscall.WaitStatus
	if exitError, ok := pr.ProcessError.(*exec.ExitError); ok {
		waitStatus = exitError.Sys().(syscall.WaitStatus)
	} else {
		waitStatus = pr.ProcessState.Sys().(syscall.WaitStatus)
	}
	return waitStatus.ExitStatus()
}

// CommandPath finds the full path of a binary given its name. This requires the wich command to be present in the system.
func (c *Context) CommandPath(name string) string {
	cmd := exec.Command("which", name)
	cmdOutput, err := cmd.Output()
	if err != nil {
		return ""
	}
	return strings.Trim(string(cmdOutput), "\n")
}

// CommandExists checks if a given binary exists in PATH.
func (c *Context) CommandExists(name string) bool {
	return c.CommandPath(name) != ""
}

// MustCommandExist ensures a given binary exists in PATH, otherwise panics.
func (c *Context) MustCommandExist(name string) {
	if !c.CommandExists(name) {
		panic(fmt.Errorf("Command %s is not available. Please make sure it is installed and accessible.", name))
	}
}

// ExecuteDebug executes a system command, stdout and stderr are piped
func (c *Context) ExecuteDebug(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(false, false, name, args...)
	return
}

// ExecuteSilent executes a  system command without outputting stdout (it is
// still captured and can be retrieved using the returned ProcessResult)
func (c *Context) ExecuteSilent(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(true, false, name, args...)
	return
}

// ExecuteFullySilent executes a system command without outputting stdout or
// stderr (both are still captured and can be retrieved using the returned ProcessResult)
func (c *Context) ExecuteFullySilent(name string, args ...string) (pr *ProcessResult, err error) {
	pr, err = c.Execute(true, true, name, args...)
	return
}

// MustExecuteDebug ensures a system command to be executed, otherwise panics
func (c *Context) MustExecuteDebug(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.Execute(false, false, name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// MustExecuteSilent ensures a system command to be executed without outputting
// stdout, otherwise panics
func (c *Context) MustExecuteSilent(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.ExecuteSilent(name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// MustExecuteFullySilent ensures a system command to be executed without
// outputting stdout and stderr, otherwise panics
func (c *Context) MustExecuteFullySilent(name string, args ...string) (pr *ProcessResult) {
	pr, err := c.ExecuteFullySilent(name, args...)
	if err != nil {
		panic(err)
	}
	return
}

// Execute executes a system command with configurable stdout and stderr output
func (c *Context) Execute(stdoutSilent bool, stderrSilent bool, name string, args ...string) (pr *ProcessResult, err error) {
	cmd, pr := c.prepareCommand(stdoutSilent, stderrSilent, name, args...)

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	pr.ProcessError = err

	return pr, err
}

// ExecuteDetached executes the given command in this context in the background (detached). This means the script execution instantly continues.
func (c *Context) ExecuteDetached(name string, args ...string) (cmd *exec.Cmd, pr *ProcessResult, err error) {
	cmd, pr = c.prepareCommand(true, true, name, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid: true,
	}
	err = cmd.Start()
	return
}

func (c Context) prepareCommand(stdoutSilent bool, stderrSilent bool, name string, args ...string) (*exec.Cmd, *ProcessResult) {
	pr := NewProcessResult()

	cmd := exec.Command(name, args...)
	pr.Cmd = cmd
	pr.ProcessState = cmd.ProcessState

	cmd.Dir = c.workingDir
	cmd.Env = c.getFullEnv()

	if stderrSilent {
		cmd.Stderr = pr.stderrBuffer
	} else {
		cmd.Stderr = io.MultiWriter(os.Stderr, pr.stderrBuffer)
	}
	if stderrSilent {
		cmd.Stdout = pr.stdoutBuffer
	} else {
		cmd.Stdout = io.MultiWriter(os.Stdout, pr.stdoutBuffer)
	}
	return cmd, pr
}
<|endoftext|>


<commit_before>// Package profile provides a simple way to manage runtime/pprof
// profiling of your Go application.
package profile

import (
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
)

type profile struct {
	// Quiet suppresses informational messages during profiling.
	Quiet bool

	// CPUProfile controls if cpu profiling will be enabled.
	CPUProfile bool

	// MemProfile controls if memory profiling will be enabled.
	MemProfile bool

	// BlockProfile controls if block (contention) profiling will
	// be enabled.
	// It defaults to false.
	BlockProfile bool

	// NoShutdownHook controls whether the profiling package should
	// hook SIGINT to write profiles cleanly.
	NoShutdownHook bool

	// MemProfileRate sent the rate for the memory profile
	memProfileRate int

	// ProfilePath controls the base path where various profiling
	// files are written. If blank, the base path will be generated
	// by ioutil.TempDir.
	path string

	closers []func()
}

// NoShutdownHook controls whether the profiling package should
// hook SIGINT to write profiles cleanly.
// Programs with more sophisticated signal handling should set
// this to true and ensure the Stop() function returned from Start()
// is called during shutdown.
func NoShutdownHook(p *profile) { p.NoShutdownHook = true }

// Quiet suppresses informational messages during profiling.
func Quiet(p *profile) { p.Quiet = true }

// Sets the profile path
func ProfilePath(path string) func(*profile) {
	return func(p *profile) {
		p.path = path
	}
}

func (p *profile) NoProfiles() {
	p.CPUProfile = false
	p.MemProfile = false
	p.BlockProfile = false
}

// CPUProfile controls if cpu profiling will be enabled. It disables any previous profiling settings.
func CPUProfile(p *profile) {
	p.NoProfiles()
	p.CPUProfile = true
}

// MemProfile controls if memory profiling will be enabled. It disables any previous profiling settings.
func MemProfile(p *profile) {
	p.NoProfiles()
	p.MemProfile = true
}

// BlockProfile controls if block (contention) profiling will be enabled. It disables any previous profiling settings.
func BlockProfile(p *profile) {
	p.NoProfiles()
	p.BlockProfile = true
}

// path resolves the profile's path or outputs to a temporary directory
func (p *profile) profilePath() (resolvedPath string, err error) {
	if p := p.path; p != "" {
		return p, os.MkdirAll(p, 0777)
	}

	return ioutil.TempDir("", "profile")
}

func (p *profile) Stop() {
	for _, c := range p.closers {
		c()
	}
}

func baseProfile() *profile {
	prof := &profile{memProfileRate: 4096}
	CPUProfile(prof)
	return prof
}

// Start starts a new profiling session.
// The caller should call the Stop method on the value returned
// to cleanly stop profiling.
func Start(options ...func(*profile)) interface {
	Stop()
} {
	prof := baseProfile()
	for _, option := range options {
		option(prof)
	}

	path, err := prof.profilePath()
	if err != nil {
		log.Fatalf("profile: could not create initial output directory: %v", err)
	}

	if prof.Quiet {
		log.SetOutput(ioutil.Discard)
	}

	switch {
	case prof.CPUProfile:
		fn := filepath.Join(path, "cpu.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
		}
		log.Printf("profile: cpu profiling enabled, %s", fn)
		pprof.StartCPUProfile(f)
		prof.closers = append(prof.closers, func() {
			pprof.StopCPUProfile()
			f.Close()
		})

	case prof.MemProfile:
		fn := filepath.Join(path, "mem.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
		}
		old := runtime.MemProfileRate
		runtime.MemProfileRate = prof.memProfileRate
		log.Printf("profile: memory profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("heap").WriteTo(f, 0)
			f.Close()
			runtime.MemProfileRate = old
		})

	case prof.BlockProfile:
		fn := filepath.Join(path, "block.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create block profile %q: %v", fn, err)
		}
		runtime.SetBlockProfileRate(1)
		log.Printf("profile: block profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("block").WriteTo(f, 0)
			f.Close()
			runtime.SetBlockProfileRate(0)
		})
	}

	if !prof.NoShutdownHook {
		go func() {
			c := make(chan os.Signal, 1)
			signal.Notify(c, os.Interrupt)
			<-c

			log.Println("profile: caught interrupt, stopping profiles")
			prof.Stop()

			os.Exit(0)
		}()
	}

	return prof
}
<commit_msg>Updated comment<commit_after>// Package profile provides a simple way to manage runtime/pprof
// profiling of your Go application.
package profile

import (
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
)

type profile struct {
	// Quiet suppresses informational messages during profiling.
	Quiet bool

	// CPUProfile controls if cpu profiling will be enabled.
	CPUProfile bool

	// MemProfile controls if memory profiling will be enabled.
	MemProfile bool

	// BlockProfile controls if block (contention) profiling will
	// be enabled.
	// It defaults to false.
	BlockProfile bool

	// NoShutdownHook controls whether the profiling package should
	// hook SIGINT to write profiles cleanly.
	NoShutdownHook bool

	// MemProfileRate sent the rate for the memory profile
	memProfileRate int

	// path holds the base path where various profiling files are  written.
	// If blank, the base path will be generated by ioutil.TempDir.
	path string

	closers []func()
}

// NoShutdownHook controls whether the profiling package should
// hook SIGINT to write profiles cleanly.
// Programs with more sophisticated signal handling should set
// this to true and ensure the Stop() function returned from Start()
// is called during shutdown.
func NoShutdownHook(p *profile) { p.NoShutdownHook = true }

// Quiet suppresses informational messages during profiling.
func Quiet(p *profile) { p.Quiet = true }

// Sets the profile path
func ProfilePath(path string) func(*profile) {
	return func(p *profile) {
		p.path = path
	}
}

func (p *profile) NoProfiles() {
	p.CPUProfile = false
	p.MemProfile = false
	p.BlockProfile = false
}

// CPUProfile controls if cpu profiling will be enabled. It disables any previous profiling settings.
func CPUProfile(p *profile) {
	p.NoProfiles()
	p.CPUProfile = true
}

// MemProfile controls if memory profiling will be enabled. It disables any previous profiling settings.
func MemProfile(p *profile) {
	p.NoProfiles()
	p.MemProfile = true
}

// BlockProfile controls if block (contention) profiling will be enabled. It disables any previous profiling settings.
func BlockProfile(p *profile) {
	p.NoProfiles()
	p.BlockProfile = true
}

// path resolves the profile's path or outputs to a temporary directory
func (p *profile) profilePath() (resolvedPath string, err error) {
	if p := p.path; p != "" {
		return p, os.MkdirAll(p, 0777)
	}

	return ioutil.TempDir("", "profile")
}

func (p *profile) Stop() {
	for _, c := range p.closers {
		c()
	}
}

func baseProfile() *profile {
	prof := &profile{memProfileRate: 4096}
	CPUProfile(prof)
	return prof
}

// Start starts a new profiling session.
// The caller should call the Stop method on the value returned
// to cleanly stop profiling.
func Start(options ...func(*profile)) interface {
	Stop()
} {
	prof := baseProfile()
	for _, option := range options {
		option(prof)
	}

	path, err := prof.profilePath()
	if err != nil {
		log.Fatalf("profile: could not create initial output directory: %v", err)
	}

	if prof.Quiet {
		log.SetOutput(ioutil.Discard)
	}

	switch {
	case prof.CPUProfile:
		fn := filepath.Join(path, "cpu.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
		}
		log.Printf("profile: cpu profiling enabled, %s", fn)
		pprof.StartCPUProfile(f)
		prof.closers = append(prof.closers, func() {
			pprof.StopCPUProfile()
			f.Close()
		})

	case prof.MemProfile:
		fn := filepath.Join(path, "mem.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
		}
		old := runtime.MemProfileRate
		runtime.MemProfileRate = prof.memProfileRate
		log.Printf("profile: memory profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("heap").WriteTo(f, 0)
			f.Close()
			runtime.MemProfileRate = old
		})

	case prof.BlockProfile:
		fn := filepath.Join(path, "block.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatalf("profile: could not create block profile %q: %v", fn, err)
		}
		runtime.SetBlockProfileRate(1)
		log.Printf("profile: block profiling enabled, %s", fn)
		prof.closers = append(prof.closers, func() {
			pprof.Lookup("block").WriteTo(f, 0)
			f.Close()
			runtime.SetBlockProfileRate(0)
		})
	}

	if !prof.NoShutdownHook {
		go func() {
			c := make(chan os.Signal, 1)
			signal.Notify(c, os.Interrupt)
			<-c

			log.Println("profile: caught interrupt, stopping profiles")
			prof.Stop()

			os.Exit(0)
		}()
	}

	return prof
}
<|endoftext|>