<commit_before>// Copyright 2013, Örjan Persson. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package logging implements a logging infrastructure for Go. It supports
// different logging backends like syslog, file and memory. Multiple backends
// can be utilized with different log levels per backend and logger.
package logging

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"sync/atomic"
	"time"
)

// Redactor is an interface for types that may contain sensitive information
// (like passwords), which shouldn't be printed to the log. The idea was found
// in relog as part of the vitness project.
type Redactor interface {
	Redacted() interface{}
}

// Redact returns a string of * having the same length as s.
func Redact(s string) string {
	return strings.Repeat("*", len(s))
}

var (
	// Sequence number is incremented and utilized for all log records created.
	sequenceNo uint64

	// timeNow is a customizable for testing purposes.
	timeNow = time.Now
)

// Record represents a log record and contains the timestamp when the record
// was created, an increasing id, filename and line and finally the actual
// formatted log line.
type Record struct {
	Id     uint64
	Time   time.Time
	Module string
	Level  Level

	// message is kept as a pointer to have shallow copies update this once
	// needed.
	message   *string
	args      []interface{}
	fmt       string
	formatter Formatter
	formatted string
}

func (r *Record) Formatted(calldepth int) string {
	if r.formatted == "" {
		var buf bytes.Buffer
		r.formatter.Format(calldepth+1, r, &buf)
		r.formatted = buf.String()
	}
	return r.formatted
}

func (r *Record) Message() string {
	if r.message == nil {
		// Redact the arguments that implements the Redactor interface
		for i, arg := range r.args {
			if redactor, ok := arg.(Redactor); ok == true {
				r.args[i] = redactor.Redacted()
			}
		}
		msg := fmt.Sprintf(r.fmt, r.args...)
		r.message = &msg
	}
	return *r.message
}

type Logger struct {
	Module      string
	backend     LeveledBackend
	haveBackend bool

	// ExtraCallDepth can be used to add additional call depth when getting the
	// calling function. This is normally used when wrapping a logger.
	ExtraCalldepth int
}

func (l *Logger) SetBackend(backend LeveledBackend) {
	l.backend = backend
	l.haveBackend = true
}

// TODO call NewLogger and remove MustGetLogger?
// GetLogger creates and returns a Logger object based on the module name.
func GetLogger(module string) (*Logger, error) {
	return &Logger{Module: module}, nil
}

// MustGetLogger is like GetLogger but panics if the logger can't be created.
// It simplifies safe initialization of a global logger for eg. a package.
func MustGetLogger(module string) *Logger {
	logger, err := GetLogger(module)
	if err != nil {
		panic("logger: " + module + ": " + err.Error())
	}
	return logger
}

// Reset restores the internal state of the logging library.
func Reset() {
	// TODO make a global Init() method to be less magic? or make it such that
	// if there's no backends at all configured, we could use some tricks to
	// automatically setup backends based if we have a TTY or not.
	sequenceNo = 0
	b := SetBackend(NewLogBackend(os.Stderr, "", log.LstdFlags))
	b.SetLevel(DEBUG, "")
	SetFormatter(DefaultFormatter)
	timeNow = time.Now
}

// IsEnabledFor returns true if the logger is enabled for the given level.
func (l *Logger) IsEnabledFor(level Level) bool {
	return defaultBackend.IsEnabledFor(level, l.Module)
}

func (l *Logger) log(lvl Level, format string, args ...interface{}) {
	if !l.IsEnabledFor(lvl) {
		return
	}

	// Create the logging record and pass it in to the backend
	record := &Record{
		Id:     atomic.AddUint64(&sequenceNo, 1),
		Time:   timeNow(),
		Module: l.Module,
		Level:  lvl,
		fmt:    format,
		args:   args,
	}

	// TODO use channels to fan out the records to all backends?
	// TODO in case of errors, do something (tricky)

	// calldepth=2 brings the stack up to the caller of the level
	// methods, Info(), Fatal(), etc.
	// ExtraCallDepth allows this to be extended further up the stack in case we
	// are wrapping these methods, eg. to expose them package level
	if l.haveBackend {
		l.backend.Log(lvl, 2+l.ExtraCalldepth, record)
		return
	}

	defaultBackend.Log(lvl, 2+l.ExtraCalldepth, record)
}

// Fatal is equivalent to l.Critical(fmt.Sprint()) followed by a call to os.Exit(1).
func (l *Logger) Fatal(args ...interface{}) {
	s := fmt.Sprint(args...)
	l.log(CRITICAL, "%s", s)
	os.Exit(1)
}

// Fatalf is equivalent to l.Critical followed by a call to os.Exit(1).
func (l *Logger) Fatalf(format string, args ...interface{}) {
	l.log(CRITICAL, format, args...)
	os.Exit(1)
}

// Panic is equivalent to l.Critical(fmt.Sprint()) followed by a call to panic().
func (l *Logger) Panic(args ...interface{}) {
	s := fmt.Sprint(args...)
	l.log(CRITICAL, "%s", s)
	panic(s)
}

// Panicf is equivalent to l.Critical followed by a call to panic().
func (l *Logger) Panicf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	l.log(CRITICAL, "%s", s)
	panic(s)
}

// Critical logs a message using CRITICAL as log level.
func (l *Logger) Critical(format string, args ...interface{}) {
	l.log(CRITICAL, format, args...)
}

// Error logs a message using ERROR as log level.
func (l *Logger) Error(format string, args ...interface{}) {
	l.log(ERROR, format, args...)
}

// Warning logs a message using WARNING as log level.
func (l *Logger) Warning(format string, args ...interface{}) {
	l.log(WARNING, format, args...)
}

// Notice logs a message using NOTICE as log level.
func (l *Logger) Notice(format string, args ...interface{}) {
	l.log(NOTICE, format, args...)
}

// Info logs a message using INFO as log level.
func (l *Logger) Info(format string, args ...interface{}) {
	l.log(INFO, format, args...)
}

// Debug logs a message using DEBUG as log level.
func (l *Logger) Debug(format string, args ...interface{}) {
	l.log(DEBUG, format, args...)
}

func init() {
	Reset()
}
<commit_msg>Added Is<Level>Enabled methods for logging interfaces compatibility<commit_after>// Copyright 2013, Örjan Persson. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package logging implements a logging infrastructure for Go. It supports
// different logging backends like syslog, file and memory. Multiple backends
// can be utilized with different log levels per backend and logger.
package logging

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"sync/atomic"
	"time"
)

// Redactor is an interface for types that may contain sensitive information
// (like passwords), which shouldn't be printed to the log. The idea was found
// in relog as part of the vitness project.
type Redactor interface {
	Redacted() interface{}
}

// Redact returns a string of * having the same length as s.
func Redact(s string) string {
	return strings.Repeat("*", len(s))
}

var (
	// Sequence number is incremented and utilized for all log records created.
	sequenceNo uint64

	// timeNow is a customizable for testing purposes.
	timeNow = time.Now
)

// Record represents a log record and contains the timestamp when the record
// was created, an increasing id, filename and line and finally the actual
// formatted log line.
type Record struct {
	Id     uint64
	Time   time.Time
	Module string
	Level  Level

	// message is kept as a pointer to have shallow copies update this once
	// needed.
	message   *string
	args      []interface{}
	fmt       string
	formatter Formatter
	formatted string
}

func (r *Record) Formatted(calldepth int) string {
	if r.formatted == "" {
		var buf bytes.Buffer
		r.formatter.Format(calldepth+1, r, &buf)
		r.formatted = buf.String()
	}
	return r.formatted
}

func (r *Record) Message() string {
	if r.message == nil {
		// Redact the arguments that implements the Redactor interface
		for i, arg := range r.args {
			if redactor, ok := arg.(Redactor); ok == true {
				r.args[i] = redactor.Redacted()
			}
		}
		msg := fmt.Sprintf(r.fmt, r.args...)
		r.message = &msg
	}
	return *r.message
}

type Logger struct {
	Module      string
	backend     LeveledBackend
	haveBackend bool

	// ExtraCallDepth can be used to add additional call depth when getting the
	// calling function. This is normally used when wrapping a logger.
	ExtraCalldepth int
}

func (l *Logger) SetBackend(backend LeveledBackend) {
	l.backend = backend
	l.haveBackend = true
}

// TODO call NewLogger and remove MustGetLogger?
// GetLogger creates and returns a Logger object based on the module name.
func GetLogger(module string) (*Logger, error) {
	return &Logger{Module: module}, nil
}

// MustGetLogger is like GetLogger but panics if the logger can't be created.
// It simplifies safe initialization of a global logger for eg. a package.
func MustGetLogger(module string) *Logger {
	logger, err := GetLogger(module)
	if err != nil {
		panic("logger: " + module + ": " + err.Error())
	}
	return logger
}

// Reset restores the internal state of the logging library.
func Reset() {
	// TODO make a global Init() method to be less magic? or make it such that
	// if there's no backends at all configured, we could use some tricks to
	// automatically setup backends based if we have a TTY or not.
	sequenceNo = 0
	b := SetBackend(NewLogBackend(os.Stderr, "", log.LstdFlags))
	b.SetLevel(DEBUG, "")
	SetFormatter(DefaultFormatter)
	timeNow = time.Now
}

// IsEnabledFor returns true if the logger is enabled for the given level.
func (l *Logger) IsEnabledFor(level Level) bool {
	return defaultBackend.IsEnabledFor(level, l.Module)
}

func (l *Logger) log(lvl Level, format string, args ...interface{}) {
	if !l.IsEnabledFor(lvl) {
		return
	}

	// Create the logging record and pass it in to the backend
	record := &Record{
		Id:     atomic.AddUint64(&sequenceNo, 1),
		Time:   timeNow(),
		Module: l.Module,
		Level:  lvl,
		fmt:    format,
		args:   args,
	}

	// TODO use channels to fan out the records to all backends?
	// TODO in case of errors, do something (tricky)

	// calldepth=2 brings the stack up to the caller of the level
	// methods, Info(), Fatal(), etc.
	// ExtraCallDepth allows this to be extended further up the stack in case we
	// are wrapping these methods, eg. to expose them package level
	if l.haveBackend {
		l.backend.Log(lvl, 2+l.ExtraCalldepth, record)
		return
	}

	defaultBackend.Log(lvl, 2+l.ExtraCalldepth, record)
}

// Fatal is equivalent to l.Critical(fmt.Sprint()) followed by a call to os.Exit(1).
func (l *Logger) Fatal(args ...interface{}) {
	s := fmt.Sprint(args...)
	l.log(CRITICAL, "%s", s)
	os.Exit(1)
}

// Fatalf is equivalent to l.Critical followed by a call to os.Exit(1).
func (l *Logger) Fatalf(format string, args ...interface{}) {
	l.log(CRITICAL, format, args...)
	os.Exit(1)
}

// Panic is equivalent to l.Critical(fmt.Sprint()) followed by a call to panic().
func (l *Logger) Panic(args ...interface{}) {
	s := fmt.Sprint(args...)
	l.log(CRITICAL, "%s", s)
	panic(s)
}

// Panicf is equivalent to l.Critical followed by a call to panic().
func (l *Logger) Panicf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	l.log(CRITICAL, "%s", s)
	panic(s)
}

// Critical logs a message using CRITICAL as log level.
func (l *Logger) Critical(format string, args ...interface{}) {
	l.log(CRITICAL, format, args...)
}

// Error logs a message using ERROR as log level.
func (l *Logger) Error(format string, args ...interface{}) {
	l.log(ERROR, format, args...)
}

// Warning logs a message using WARNING as log level.
func (l *Logger) Warning(format string, args ...interface{}) {
	l.log(WARNING, format, args...)
}

// Notice logs a message using NOTICE as log level.
func (l *Logger) Notice(format string, args ...interface{}) {
	l.log(NOTICE, format, args...)
}

// Info logs a message using INFO as log level.
func (l *Logger) Info(format string, args ...interface{}) {
	l.log(INFO, format, args...)
}

// Debug logs a message using DEBUG as log level.
func (l *Logger) Debug(format string, args ...interface{}) {
	l.log(DEBUG, format, args...)
}

// IsFatalEnabled returns true if the logger is enabled for fatal messages.
func (l *Logger) IsFatalEnabled() bool {
    return defaultBackend.IsEnabledFor(CRITICAL, l.Module)
}

// IsPanicEnabled returns true if the logger is enabled for panic messages.
func (l *Logger) IsPanicEnabled() bool {
    return defaultBackend.IsEnabledFor(CRITICAL, l.Module)
}

// IsCriticalEnabled returns true if the logger is enabled for critical messages.
func (l *Logger) IsCriticalEnabled() bool {
    return defaultBackend.IsEnabledFor(CRITICAL, l.Module)
}

// IsErrorEnabled returns true if the logger is enabled for error messages.
func (l *Logger) IsErrorEnabled() bool {
    return defaultBackend.IsEnabledFor(ERROR, l.Module)
}

// IsWarningEnabled returns true if the logger is enabled for warning messages.
func (l *Logger) IsWarningEnabled() bool {
    return defaultBackend.IsEnabledFor(WARNING, l.Module)
}

// IsNoticeEnabled returns true if the logger is enabled for notice messages.
func (l *Logger) IsNoticeEnabled() bool {
    return defaultBackend.IsEnabledFor(NOTICE, l.Module)
}

// IsInfoEnabled returns true if the logger is enabled for info messages.
func (l *Logger) IsInfoEnabled() bool {
    return defaultBackend.IsEnabledFor(INFO, l.Module)
}

// IsDebugEnabled returns true if the logger is enabled for debug messages.
func (l *Logger) IsDebugEnabled() bool {
    return defaultBackend.IsEnabledFor(DEBUG, l.Module)
}

func init() {
	Reset()
}
<|endoftext|>


<commit_before>package operator

import (
	"archive/tar"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
)

// File represents an archive file.
type File struct {
	Path     string
	Rel      string
	FileInfo os.FileInfo
}

// String returns the path of a file.
func (f *File) String() string {
	return f.Path
}

// Tape represents an archive.
type Tape []*File

// Copy writes a tar archive of all members.
func (t Tape) Copy(w io.WriteCloser) error {
	archive := tar.NewWriter(w)
	for _, member := range t {
		file, err := os.Open(member.Path)
		if err != nil {
			// File might have been deleted, we can ignore it.
			if isNotExist(err) {
				continue
			}
			return err
		}
		info, err := os.Lstat(member.Path)
		if err != nil {
			return err
		}
		link, err := filepath.EvalSymlinks(member.Path)
		if err != nil {
			return err
		}
		header, err := tar.FileInfoHeader(info, link)
		if err != nil {
			return err
		}
		header.Name = member.Rel
		if err = archive.WriteHeader(header); err != nil {
			return err
		}
		if !info.Mode().IsRegular() {
			continue
		}
		if _, err = io.Copy(archive, file); err != nil {
			if err != tar.ErrWriteTooLong {
				return err
			}
		}
		if err := file.Close(); err != nil {
			return err
		}
	}
	return archive.Close()
}

func walk(cluster string) (files []*File, err error) {
	err = filepath.Walk(cluster, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			// An error occured, stop processing
			return err
		}
		if strings.Contains(path, "pg_xlog") && !info.IsDir() {
			// We don't want to archive WAL files but we want the pg_xlog directory
			return nil
		}
		if strings.Contains(path, "pg_log") && !info.IsDir() {
			// We don't want to archive log files but we want the pg_log directory
			return nil
		}
		if info.Name() == "postgresql.conf" || info.Name() == "postmaster.pid" {
			// Ignore configuration and pid files
			return nil
		}
		rel, err := filepath.Rel(cluster, path)
		if err != nil {
			return err
		}
		files = append(files, &File{
			Path:     path,
			Rel:      rel,
			FileInfo: info,
		})
		return nil
	})
	return files, err
}

// Archive creates an archive for the given directory.
func Archive(cluster string) (Tape, error) {
	var archive []*File
	files, err := walk(cluster)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		archive = append(archive, file)
	}
	return archive, nil
}

// Extract extracts the archive to the given directory.
func Extract(cluster string, archive io.ReadCloser) error {
	tr := tar.NewReader(archive)
	for {
		header, err := tr.Next()
		if err != nil {
			if err == io.EOF {
				// End of archive
				break
			}
			return err
		}
		filename := filepath.Join(cluster, header.Name)
		info := header.FileInfo()
		if info.IsDir() {
			os.MkdirAll(filename, info.Mode())
			continue
		}
		if isSymlink(info) {
			os.Symlink(header.Linkname, filename)
			continue
		}
		file, err := createFile(filename, info.Mode())
		if err != nil {
			return err
		}
		if _, err = io.Copy(file, tr); err != nil {
			return err
		}
	}
	return nil
}

func createFile(name string, mode os.FileMode) (*os.File, error) {
	if err := os.MkdirAll(filepath.Dir(name), 0700); err != nil {
		return nil, err
	}
	return os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, mode)
}

func isSymlink(fi os.FileInfo) bool {
	return fi.Mode()&os.ModeSymlink == os.ModeSymlink
}

func isNotExist(err error) bool {
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		err = pe.Err
	case *os.LinkError:
		err = pe.Err
	}
	return err == syscall.ENOENT || err == os.ErrNotExist
}
<commit_msg>Ensure we close files<commit_after>package operator

import (
	"archive/tar"
	"io"
	"os"
	"path/filepath"
	"strings"
	"syscall"
)

// File represents an archive file.
type File struct {
	Path     string
	Rel      string
	FileInfo os.FileInfo
}

// String returns the path of a file.
func (f *File) String() string {
	return f.Path
}

// Tape represents an archive.
type Tape []*File

// Copy writes a tar archive of all members.
func (t Tape) Copy(w io.WriteCloser) error {
	archive := tar.NewWriter(w)
	for _, member := range t {
		file, err := os.Open(member.Path)
		if err != nil {
			// File might have been deleted, we can ignore it.
			if isNotExist(err) {
				continue
			}
			return err
		}
		defer file.Close()
		info, err := os.Lstat(member.Path)
		if err != nil {
			return err
		}
		link, err := filepath.EvalSymlinks(member.Path)
		if err != nil {
			return err
		}
		header, err := tar.FileInfoHeader(info, link)
		if err != nil {
			return err
		}
		header.Name = member.Rel
		if err = archive.WriteHeader(header); err != nil {
			return err
		}
		if !info.Mode().IsRegular() {
			continue
		}
		if _, err = io.Copy(archive, file); err != nil {
			if err != tar.ErrWriteTooLong {
				return err
			}
		}
	}
	return archive.Close()
}

func walk(cluster string) (files []*File, err error) {
	err = filepath.Walk(cluster, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			// An error occured, stop processing
			return err
		}
		if strings.Contains(path, "pg_xlog") && !info.IsDir() {
			// We don't want to archive WAL files but we want the pg_xlog directory
			return nil
		}
		if strings.Contains(path, "pg_log") && !info.IsDir() {
			// We don't want to archive log files but we want the pg_log directory
			return nil
		}
		if info.Name() == "postgresql.conf" || info.Name() == "postmaster.pid" {
			// Ignore configuration and pid files
			return nil
		}
		rel, err := filepath.Rel(cluster, path)
		if err != nil {
			return err
		}
		files = append(files, &File{
			Path:     path,
			Rel:      rel,
			FileInfo: info,
		})
		return nil
	})
	return files, err
}

// Archive creates an archive for the given directory.
func Archive(cluster string) (Tape, error) {
	var archive []*File
	files, err := walk(cluster)
	if err != nil {
		return nil, err
	}
	for _, file := range files {
		archive = append(archive, file)
	}
	return archive, nil
}

// Extract extracts the archive to the given directory.
func Extract(cluster string, archive io.ReadCloser) error {
	tr := tar.NewReader(archive)
	for {
		header, err := tr.Next()
		if err != nil {
			if err == io.EOF {
				// End of archive
				break
			}
			return err
		}
		filename := filepath.Join(cluster, header.Name)
		info := header.FileInfo()
		if info.IsDir() {
			os.MkdirAll(filename, info.Mode())
			continue
		}
		if isSymlink(info) {
			os.Symlink(header.Linkname, filename)
			continue
		}
		file, err := createFile(filename, info.Mode())
		if err != nil {
			return err
		}
		defer file.Close()
		if _, err = io.Copy(file, tr); err != nil {
			return err
		}
	}
	return nil
}

func createFile(name string, mode os.FileMode) (*os.File, error) {
	if err := os.MkdirAll(filepath.Dir(name), 0700); err != nil {
		return nil, err
	}
	return os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, mode)
}

func isSymlink(fi os.FileInfo) bool {
	return fi.Mode()&os.ModeSymlink == os.ModeSymlink
}

func isNotExist(err error) bool {
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		err = pe.Err
	case *os.LinkError:
		err = pe.Err
	}
	return err == syscall.ENOENT || err == os.ErrNotExist
}
<|endoftext|>


<commit_before>package wav

import (
	"encoding/binary"
	"fmt"
	"io"
	"time"

	"github.com/faiface/pixel/audio"
	"github.com/pkg/errors"
)

type ReadSeekCloser interface {
	io.Reader
	io.Seeker
	io.Closer
}

func Decode(rsc ReadSeekCloser) (s audio.StreamSeekCloser, err error) {
	var d decoder
	d.rsc = rsc
	defer func() { // hacky way to always close rsc if an error occured
		if err != nil {
			d.rsc.Close()
		}
	}()
	herr := binary.Read(rsc, binary.LittleEndian, &d.h)
	if herr != nil {
		return nil, errors.Wrap(herr, "wav")
	}
	if string(d.h.RiffMark[:]) != "RIFF" {
		return nil, errors.New("wav: missing RIFF at the beginning")
	}
	if string(d.h.WaveMark[:]) != "WAVE" {
		return nil, errors.New("wav: unsupported file type")
	}
	if string(d.h.FmtMark[:]) != "fmt " {
		return nil, errors.New("wav: missing format chunk marker")
	}
	if string(d.h.DataMark[:]) != "data" {
		return nil, errors.New("wav: missing data chunk marker")
	}
	if d.h.FormatType != 1 {
		return nil, errors.New("wav: unsupported format type")
	}
	if d.h.NumChans <= 0 {
		return nil, errors.New("wav: invalid number of channels (less than 1)")
	}
	if d.h.BitsPerSample != 8 && d.h.BitsPerSample != 16 {
		return nil, errors.New("wav: unsupported number of bits per sample, 8 or 16 are supported")
	}
	return &d, nil
}

type header struct {
	RiffMark      [4]byte
	FileSize      int32
	WaveMark      [4]byte
	FmtMark       [4]byte
	FormatSize    int32
	FormatType    int16
	NumChans      int16
	SampleRate    int32
	ByteRate      int32
	BytesPerFrame int16
	BitsPerSample int16
	DataMark      [4]byte
	DataSize      int32
}

type decoder struct {
	rsc ReadSeekCloser
	h   header
	pos int32
	err error
}

func (s *decoder) Err() error {
	return s.err
}

func (s *decoder) Duration() time.Duration {
	numBytes := time.Duration(s.h.DataSize)
	perFrame := time.Duration(s.h.BytesPerFrame)
	sampRate := time.Duration(s.h.SampleRate)
	return numBytes / perFrame * time.Second / sampRate
}

func (s *decoder) Position() time.Duration {
	frameIndex := time.Duration(s.pos / int32(s.h.BytesPerFrame))
	frameTime := time.Second / time.Duration(s.h.SampleRate)
	return frameIndex * frameTime
}

func (s *decoder) Seek(d time.Duration) error {
	if d < 0 || s.Duration() < d {
		return fmt.Errorf("wav: seek duration %v out of range [%v, %v]", d, 0, s.Duration())
	}
	frame := int32(d / (time.Second / time.Duration(s.h.SampleRate)))
	pos := frame * int32(s.h.BytesPerFrame)
	_, err := s.rsc.Seek(int64(pos)+44, io.SeekStart) // 44 is the size of the header
	if err != nil {
		return errors.Wrap(err, "wav: seek error")
	}
	s.pos = pos
	return nil
}

func (s *decoder) Stream(samples [][2]float64) (n int, ok bool) {
	if s.err != nil || s.pos >= s.h.DataSize {
		return 0, false
	}
	var frameWidth int
	switch {
	case s.h.BitsPerSample == 8 && s.h.NumChans == 1:
		frameWidth = 1
	case s.h.BitsPerSample == 8 && s.h.NumChans >= 2:
		frameWidth = int(s.h.NumChans)
	case s.h.BitsPerSample == 16 && s.h.NumChans == 1:
		frameWidth = 2
	case s.h.BitsPerSample == 16 && s.h.NumChans >= 2:
		frameWidth = int(s.h.NumChans) * 2
	}
	p := make([]byte, len(samples)*frameWidth)
	n, err := s.rsc.Read(p)
	if err != nil {
		s.err = err
	}
	switch {
	case s.h.BitsPerSample == 8 && s.h.NumChans == 1:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			val := float64(p[i])/(1<<8-1)*2 - 1
			samples[j][0] = val
			samples[j][1] = val
		}
	case s.h.BitsPerSample == 8 && s.h.NumChans >= 2:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			samples[j][0] = float64(p[i+0])/(1<<8-1)*2 - 1
			samples[j][1] = float64(p[i+1])/(1<<8-1)*2 - 1
		}
	case s.h.BitsPerSample == 16 && s.h.NumChans == 1:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			val := float64(int16(p[i+0])+int16(p[i+1])*(1<<8)) / (1<<15 - 1)
			samples[j][0] = val
			samples[j][1] = val
		}
	case s.h.BitsPerSample == 16 && s.h.NumChans >= 2:
		for i, j := 0, 0; i <= n-frameWidth; i, j = i+frameWidth, j+1 {
			samples[j][0] = float64(int16(p[i+0])+int16(p[i+1])*(1<<8)) / (1<<15 - 1)
			samples[j][1] = float64(int16(p[i+2])+int16(p[i+3])*(1<<8)) / (1<<15 - 1)
		}
	}
	return n / frameWidth, true
}

func (s *decoder) Close() error {
	err := s.rsc.Close()
	if err != nil {
		return errors.Wrap(err, "wav")
	}
	return nil
}
<commit_msg>audio: wav: add doc comments<commit_after>// Package wav implements audio data decoding in WAVE format through an audio.StreamSeekCloser.
package wav

import (
	"encoding/binary"
	"fmt"
	"io"
	"time"

	"github.com/faiface/pixel/audio"
	"github.com/pkg/errors"
)

// ReadSeekCloser is a union of io.Reader, io.Seeker and io.Closer.
type ReadSeekCloser interface {
	io.Reader
	io.Seeker
	io.Closer
}

// Decode takes a ReadSeekCloser containing audio data in WAVE format and returns a
// StreamSeekCloser, which streams that audio.
//
// Do not close the supplied ReadSeekCloser, instead, use the Close method of the returned
// StreamSeekCloser when you want to release the resources.
func Decode(rsc ReadSeekCloser) (s audio.StreamSeekCloser, err error) {
	var d decoder
	d.rsc = rsc
	defer func() { // hacky way to always close rsc if an error occured
		if err != nil {
			d.rsc.Close()
		}
	}()
	herr := binary.Read(rsc, binary.LittleEndian, &d.h)
	if herr != nil {
		return nil, errors.Wrap(herr, "wav")
	}
	if string(d.h.RiffMark[:]) != "RIFF" {
		return nil, errors.New("wav: missing RIFF at the beginning")
	}
	if string(d.h.WaveMark[:]) != "WAVE" {
		return nil, errors.New("wav: unsupported file type")
	}
	if string(d.h.FmtMark[:]) != "fmt " {
		return nil, errors.New("wav: missing format chunk marker")
	}
	if string(d.h.DataMark[:]) != "data" {
		return nil, errors.New("wav: missing data chunk marker")
	}
	if d.h.FormatType != 1 {
		return nil, errors.New("wav: unsupported format type")
	}
	if d.h.NumChans <= 0 {
		return nil, errors.New("wav: invalid number of channels (less than 1)")
	}
	if d.h.BitsPerSample != 8 && d.h.BitsPerSample != 16 {
		return nil, errors.New("wav: unsupported number of bits per sample, 8 or 16 are supported")
	}
	return &d, nil
}

type header struct {
	RiffMark      [4]byte
	FileSize      int32
	WaveMark      [4]byte
	FmtMark       [4]byte
	FormatSize    int32
	FormatType    int16
	NumChans      int16
	SampleRate    int32
	ByteRate      int32
	BytesPerFrame int16
	BitsPerSample int16
	DataMark      [4]byte
	DataSize      int32
}

type decoder struct {
	rsc ReadSeekCloser
	h   header
	pos int32
	err error
}

func (s *decoder) Err() error {
	return s.err
}

func (s *decoder) Duration() time.Duration {
	numBytes := time.Duration(s.h.DataSize)
	perFrame := time.Duration(s.h.BytesPerFrame)
	sampRate := time.Duration(s.h.SampleRate)
	return numBytes / perFrame * time.Second / sampRate
}

func (s *decoder) Position() time.Duration {
	frameIndex := time.Duration(s.pos / int32(s.h.BytesPerFrame))
	frameTime := time.Second / time.Duration(s.h.SampleRate)
	return frameIndex * frameTime
}

func (s *decoder) Seek(d time.Duration) error {
	if d < 0 || s.Duration() < d {
		return fmt.Errorf("wav: seek duration %v out of range [%v, %v]", d, 0, s.Duration())
	}
	frame := int32(d / (time.Second / time.Duration(s.h.SampleRate)))
	pos := frame * int32(s.h.BytesPerFrame)
	_, err := s.rsc.Seek(int64(pos)+44, io.SeekStart) // 44 is the size of the header
	if err != nil {
		return errors.Wrap(err, "wav: seek error")
	}
	s.pos = pos
	return nil
}

func (s *decoder) Stream(samples [][2]float64) (n int, ok bool) {
	if s.err != nil || s.pos >= s.h.DataSize {
		return 0, false
	}
	var frameWidth int
	switch {
	case s.h.BitsPerSample == 8 && s.h.NumChans == 1:
		frameWidth = 1
	case s.h.BitsPerSample == 8 && s.h.NumChans >= 2:
		frameWidth = int(s.h.NumChans)
	case s.h.BitsPerSample == 16 && s.h.NumChans == 1:
		frameWidth = 2
	case s.h.BitsPerSample == 16 && s.h.NumChans >= 2:
		frameWidth = int(s.h.NumChans) * 2
	}
	p := make([]byte, len(samples)*frameWidth)
	n, err := s.rsc.Read(p)
	if err != nil {
		s.err = err
	}
	switch {
	case s.h.BitsPerSample == 8 && s.h.NumChans == 1:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			val := float64(p[i])/(1<<8-1)*2 - 1
			samples[j][0] = val
			samples[j][1] = val
		}
	case s.h.BitsPerSample == 8 && s.h.NumChans >= 2:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			samples[j][0] = float64(p[i+0])/(1<<8-1)*2 - 1
			samples[j][1] = float64(p[i+1])/(1<<8-1)*2 - 1
		}
	case s.h.BitsPerSample == 16 && s.h.NumChans == 1:
		for i, j := 0, 0; i < n-frameWidth; i, j = i+frameWidth, j+1 {
			val := float64(int16(p[i+0])+int16(p[i+1])*(1<<8)) / (1<<15 - 1)
			samples[j][0] = val
			samples[j][1] = val
		}
	case s.h.BitsPerSample == 16 && s.h.NumChans >= 2:
		for i, j := 0, 0; i <= n-frameWidth; i, j = i+frameWidth, j+1 {
			samples[j][0] = float64(int16(p[i+0])+int16(p[i+1])*(1<<8)) / (1<<15 - 1)
			samples[j][1] = float64(int16(p[i+2])+int16(p[i+3])*(1<<8)) / (1<<15 - 1)
		}
	}
	return n / frameWidth, true
}

func (s *decoder) Close() error {
	err := s.rsc.Close()
	if err != nil {
		return errors.Wrap(err, "wav")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const (
	LUNCHY_VERSION = "0.1.3"
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fileCopy(src string, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}

	defer s.Close()

	d, err := os.Create(dst)
	if err != nil {
		return err
	}

	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}

	return d.Close()
}

func findPlists(path string) []string {
	result := []string{}
	files, err := ioutil.ReadDir(path)

	if err != nil {
		return result
	}

	for _, file := range files {
		if !file.IsDir() {
			if (filepath.Ext(file.Name())) == ".plist" {
				name := strings.Replace(file.Name(), ".plist", "", -1)
				result = append(result, name)
			}
		}
	}

	return result
}

func getPlists() []string {
	path := fmt.Sprintf("%s/Library/LaunchAgents", os.Getenv("HOME"))
	files := findPlists(path)

	return files
}

func getPlist(name string) string {
	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			return plist
		}
	}

	return ""
}

func sliceIncludes(slice []string, match string) bool {
	for _, val := range slice {
		if val == match {
			return true
		}
	}

	return false
}

func printUsage() {
	fmt.Printf("Lunchy %s, the friendly launchctl wrapper\n", LUNCHY_VERSION)
	fmt.Println("Usage: lunchy [start|stop|restart|list|status|install|show|edit|remove|scan] [options]")
}

func printList() {
	for _, file := range getPlists() {
		fmt.Println(file)
	}
}

func printStatus(args []string) {
	out, err := exec.Command("launchctl", "list").Output()

	if err != nil {
		fatal("failed to get process list")
	}

	pattern := ""

	if len(args) == 3 {
		pattern = args[2]
	}

	installed := getPlists()
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")

	for _, line := range lines {
		chunks := strings.Split(line, "\t")
		clean_line := strings.Replace(line, "\t", " ", -1)

		if len(pattern) > 0 {
			if strings.Index(chunks[2], pattern) != -1 {
				if sliceIncludes(installed, chunks[2]) {
					fmt.Println(clean_line)
				}
			}
		} else {
			if sliceIncludes(installed, chunks[2]) {
				fmt.Println(clean_line)
			}
		}
	}
}

func exitWithInvalidArgs(args []string, msg string) {
	if len(args) < 3 {
		fmt.Println(msg)
		os.Exit(1)
	}
}

func startDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			startDaemon(plist)
		}
	}
}

func startDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "load", path).Output()

	if err != nil {
		fmt.Println("failed to start", name)
		return
	}

	fmt.Println("started", name)
}

func stopDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
		}
	}
}

func stopDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "unload", path).Output()

	if err != nil {
		fmt.Println("failed to stop", name)
		return
	}

	fmt.Println("stopped", name)
}

func restartDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
			startDaemon(plist)
		}
	}
}

func showPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			printPlistContent(plist)
			return
		}
	}
}

func printPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	contents, err := ioutil.ReadFile(path)

	if err != nil {
		fatal("unable to read plist")
	}

	fmt.Printf(string(contents))
}

func editPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			editPlistContent(plist)
			return
		}
	}
}

func editPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	editor := os.Getenv("EDITOR")

	if len(editor) == 0 {
		fatal("EDITOR environment variable is not set")
	}

	cmd := exec.Command(editor, path)

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	cmd.Start()
	cmd.Wait()
}

func installPlist(args []string) {
	exitWithInvalidArgs(args, "path required")

	path := args[2]

	if !fileExists(path) {
		fatal("source file does not exist")
	}

	info, _ := os.Stat(path)
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")
	new_path := fmt.Sprintf("%s/%s", base_path, info.Name())

	if fileExists(new_path) && os.Remove(new_path) != nil {
		fatal("unable to delete existing plist")
	}

	if fileCopy(path, new_path) != nil {
		fatal("failed to copy file")
	}

	fmt.Println(path, "installed to", base_path)
}

func removePlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			path := fmt.Sprintf("%s/%s.plist", base_path, plist)

			if os.Remove(path) == nil {
				fmt.Println("removed", path)
			} else {
				fmt.Println("failed to remove", path)
			}
		}
	}
}

func scanPath(args []string) {
	exitWithInvalidArgs(args, "path required")

	out, err := exec.Command("find", args[2], "-name", "*.plist").Output()

	if err != nil {
		fatal("failed to scan files")
	}

	fmt.Println(strings.TrimSpace(string(out)))
}

func fatal(message string) {
	fmt.Println(message)
	os.Exit(1)
}

func main() {
	args := os.Args

	if len(args) == 1 {
		printUsage()
		os.Exit(1)
	}

	switch args[1] {
	default:
		printUsage()
		os.Exit(1)
	case "help":
		printUsage()
		return
	case "list", "ls":
		printList()
		return
	case "status", "ps":
		printStatus(args)
		return
	case "start":
		startDaemons(args)
		return
	case "stop":
		stopDaemons(args)
		return
	case "restart":
		restartDaemons(args)
		return
	case "show":
		showPlist(args)
		return
	case "edit":
		editPlist(args)
		return
	case "install":
		installPlist(args)
		return
	case "remove", "rm":
		removePlist(args)
		return
	case "scan":
		scanPath(args)
		return
	}
}
<commit_msg>Version bump: 0.1.4<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const (
	LUNCHY_VERSION = "0.1.4"
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fileCopy(src string, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}

	defer s.Close()

	d, err := os.Create(dst)
	if err != nil {
		return err
	}

	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}

	return d.Close()
}

func findPlists(path string) []string {
	result := []string{}
	files, err := ioutil.ReadDir(path)

	if err != nil {
		return result
	}

	for _, file := range files {
		if !file.IsDir() {
			if (filepath.Ext(file.Name())) == ".plist" {
				name := strings.Replace(file.Name(), ".plist", "", -1)
				result = append(result, name)
			}
		}
	}

	return result
}

func getPlists() []string {
	path := fmt.Sprintf("%s/Library/LaunchAgents", os.Getenv("HOME"))
	files := findPlists(path)

	return files
}

func getPlist(name string) string {
	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			return plist
		}
	}

	return ""
}

func sliceIncludes(slice []string, match string) bool {
	for _, val := range slice {
		if val == match {
			return true
		}
	}

	return false
}

func printUsage() {
	fmt.Printf("Lunchy %s, the friendly launchctl wrapper\n", LUNCHY_VERSION)
	fmt.Println("Usage: lunchy [start|stop|restart|list|status|install|show|edit|remove|scan] [options]")
}

func printList() {
	for _, file := range getPlists() {
		fmt.Println(file)
	}
}

func printStatus(args []string) {
	out, err := exec.Command("launchctl", "list").Output()

	if err != nil {
		fatal("failed to get process list")
	}

	pattern := ""

	if len(args) == 3 {
		pattern = args[2]
	}

	installed := getPlists()
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")

	for _, line := range lines {
		chunks := strings.Split(line, "\t")
		clean_line := strings.Replace(line, "\t", " ", -1)

		if len(pattern) > 0 {
			if strings.Index(chunks[2], pattern) != -1 {
				if sliceIncludes(installed, chunks[2]) {
					fmt.Println(clean_line)
				}
			}
		} else {
			if sliceIncludes(installed, chunks[2]) {
				fmt.Println(clean_line)
			}
		}
	}
}

func exitWithInvalidArgs(args []string, msg string) {
	if len(args) < 3 {
		fmt.Println(msg)
		os.Exit(1)
	}
}

func startDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			startDaemon(plist)
		}
	}
}

func startDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "load", path).Output()

	if err != nil {
		fmt.Println("failed to start", name)
		return
	}

	fmt.Println("started", name)
}

func stopDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
		}
	}
}

func stopDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "unload", path).Output()

	if err != nil {
		fmt.Println("failed to stop", name)
		return
	}

	fmt.Println("stopped", name)
}

func restartDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
			startDaemon(plist)
		}
	}
}

func showPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			printPlistContent(plist)
			return
		}
	}
}

func printPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	contents, err := ioutil.ReadFile(path)

	if err != nil {
		fatal("unable to read plist")
	}

	fmt.Printf(string(contents))
}

func editPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			editPlistContent(plist)
			return
		}
	}
}

func editPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	editor := os.Getenv("EDITOR")

	if len(editor) == 0 {
		fatal("EDITOR environment variable is not set")
	}

	cmd := exec.Command(editor, path)

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	cmd.Start()
	cmd.Wait()
}

func installPlist(args []string) {
	exitWithInvalidArgs(args, "path required")

	path := args[2]

	if !fileExists(path) {
		fatal("source file does not exist")
	}

	info, _ := os.Stat(path)
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")
	new_path := fmt.Sprintf("%s/%s", base_path, info.Name())

	if fileExists(new_path) && os.Remove(new_path) != nil {
		fatal("unable to delete existing plist")
	}

	if fileCopy(path, new_path) != nil {
		fatal("failed to copy file")
	}

	fmt.Println(path, "installed to", base_path)
}

func removePlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			path := fmt.Sprintf("%s/%s.plist", base_path, plist)

			if os.Remove(path) == nil {
				fmt.Println("removed", path)
			} else {
				fmt.Println("failed to remove", path)
			}
		}
	}
}

func scanPath(args []string) {
	exitWithInvalidArgs(args, "path required")

	out, err := exec.Command("find", args[2], "-name", "*.plist").Output()

	if err != nil {
		fatal("failed to scan files")
	}

	fmt.Println(strings.TrimSpace(string(out)))
}

func fatal(message string) {
	fmt.Println(message)
	os.Exit(1)
}

func main() {
	args := os.Args

	if len(args) == 1 {
		printUsage()
		os.Exit(1)
	}

	switch args[1] {
	default:
		printUsage()
		os.Exit(1)
	case "help":
		printUsage()
		return
	case "list", "ls":
		printList()
		return
	case "status", "ps":
		printStatus(args)
		return
	case "start":
		startDaemons(args)
		return
	case "stop":
		stopDaemons(args)
		return
	case "restart":
		restartDaemons(args)
		return
	case "show":
		showPlist(args)
		return
	case "edit":
		editPlist(args)
		return
	case "install":
		installPlist(args)
		return
	case "remove", "rm":
		removePlist(args)
		return
	case "scan":
		scanPath(args)
		return
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"log"
	"github.com/mailhog/data"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// MongoDB represents MongoDB backed storage backend
type MongoDB struct {
	Session    *mgo.Session
	Collection *mgo.Collection
}

// CreateMongoDB creates a MongoDB backed storage backend
func CreateMongoDB(uri, db, coll string) *MongoDB {
	log.Printf("Connecting to MongoDB: %s\n", uri)
	session, err := mgo.Dial(uri)
	if err != nil {
		log.Printf("Error connecting to MongoDB: %s", err)
		return nil
	}
	err = session.DB(db).C(coll).EnsureIndexKey("created")
	if err != nil {
		log.Printf("Failed creating index: %s", err)
		return nil
	}
	return &MongoDB{
		Session:    session,
		Collection: session.DB(db).C(coll),
	}
}

// Store stores a message in MongoDB and returns its storage ID
func (mongo *MongoDB) Store(m *data.Message) (string, error) {
	err := mongo.Collection.Insert(m)
	if err != nil {
		log.Printf("Error inserting message: %s", err)
		return "", err
	}
	return string(m.ID), nil
}

// Count returns the number of stored messages
func (mongo *MongoDB) Count() int {
	c, _ := mongo.Collection.Count()
	return c
}

// Search finds messages matching the query
func (mongo *MongoDB) Search(kind, query string, start, limit int) (*data.Messages, int, error) {
	messages := &data.Messages{}
	var count = 0
	var field = "raw.data"
	switch kind {
		case "to":
			field = "raw.to"
		case "from":
			field = "raw.from"
	}
	err := mongo.Collection.Find(bson.M{field: bson.RegEx{Pattern: query, Options: "i"}}).Skip(start).Limit(limit).Sort("-created").Select(bson.M{
		"id":              1,
		"_id":             1,
		"from":            1,
		"to":              1,
		"content.headers": 1,
		"content.size":    1,
		"created":         1,
		"raw":             1,
	}).All(messages)
	if err != nil {
		log.Printf("Error loading messages: %s", err)
		return nil, 0, err
	}
	count, _ = mongo.Collection.Find(bson.M{field: bson.RegEx{Pattern: query, Options: "i"}}).Count()

	return messages, count, nil
}

// List returns a list of messages by index
func (mongo *MongoDB) List(start int, limit int) (*data.Messages, error) {
	messages := &data.Messages{}
	err := mongo.Collection.Find(bson.M{}).Skip(start).Limit(limit).Sort("-created").Select(bson.M{
		"id":              1,
		"_id":             1,
		"from":            1,
		"to":              1,
		"content.headers": 1,
		"content.size":    1,
		"created":         1,
		"raw":             1,
	}).All(messages)
	if err != nil {
		log.Printf("Error loading messages: %s", err)
		return nil, err
	}
	return messages, nil
}

// DeleteOne deletes an individual message by storage ID
func (mongo *MongoDB) DeleteOne(id string) error {
	_, err := mongo.Collection.RemoveAll(bson.M{"id": id})
	return err
}

// DeleteAll deletes all messages stored in MongoDB
func (mongo *MongoDB) DeleteAll() error {
	_, err := mongo.Collection.RemoveAll(bson.M{})
	return err
}

// Load loads an individual message by storage ID
func (mongo *MongoDB) Load(id string) (*data.Message, error) {
	result := &data.Message{}
	err := mongo.Collection.Find(bson.M{"id": id}).One(&result)
	if err != nil {
		log.Printf("Error loading message: %s", err)
		return nil, err
	}
	return result, nil
}
<commit_msg>update import path for mgo<commit_after>package storage

import (
	"log"
	"github.com/mailhog/data"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// MongoDB represents MongoDB backed storage backend
type MongoDB struct {
	Session    *mgo.Session
	Collection *mgo.Collection
}

// CreateMongoDB creates a MongoDB backed storage backend
func CreateMongoDB(uri, db, coll string) *MongoDB {
	log.Printf("Connecting to MongoDB: %s\n", uri)
	session, err := mgo.Dial(uri)
	if err != nil {
		log.Printf("Error connecting to MongoDB: %s", err)
		return nil
	}
	err = session.DB(db).C(coll).EnsureIndexKey("created")
	if err != nil {
		log.Printf("Failed creating index: %s", err)
		return nil
	}
	return &MongoDB{
		Session:    session,
		Collection: session.DB(db).C(coll),
	}
}

// Store stores a message in MongoDB and returns its storage ID
func (mongo *MongoDB) Store(m *data.Message) (string, error) {
	err := mongo.Collection.Insert(m)
	if err != nil {
		log.Printf("Error inserting message: %s", err)
		return "", err
	}
	return string(m.ID), nil
}

// Count returns the number of stored messages
func (mongo *MongoDB) Count() int {
	c, _ := mongo.Collection.Count()
	return c
}

// Search finds messages matching the query
func (mongo *MongoDB) Search(kind, query string, start, limit int) (*data.Messages, int, error) {
	messages := &data.Messages{}
	var count = 0
	var field = "raw.data"
	switch kind {
		case "to":
			field = "raw.to"
		case "from":
			field = "raw.from"
	}
	err := mongo.Collection.Find(bson.M{field: bson.RegEx{Pattern: query, Options: "i"}}).Skip(start).Limit(limit).Sort("-created").Select(bson.M{
		"id":              1,
		"_id":             1,
		"from":            1,
		"to":              1,
		"content.headers": 1,
		"content.size":    1,
		"created":         1,
		"raw":             1,
	}).All(messages)
	if err != nil {
		log.Printf("Error loading messages: %s", err)
		return nil, 0, err
	}
	count, _ = mongo.Collection.Find(bson.M{field: bson.RegEx{Pattern: query, Options: "i"}}).Count()

	return messages, count, nil
}

// List returns a list of messages by index
func (mongo *MongoDB) List(start int, limit int) (*data.Messages, error) {
	messages := &data.Messages{}
	err := mongo.Collection.Find(bson.M{}).Skip(start).Limit(limit).Sort("-created").Select(bson.M{
		"id":              1,
		"_id":             1,
		"from":            1,
		"to":              1,
		"content.headers": 1,
		"content.size":    1,
		"created":         1,
		"raw":             1,
	}).All(messages)
	if err != nil {
		log.Printf("Error loading messages: %s", err)
		return nil, err
	}
	return messages, nil
}

// DeleteOne deletes an individual message by storage ID
func (mongo *MongoDB) DeleteOne(id string) error {
	_, err := mongo.Collection.RemoveAll(bson.M{"id": id})
	return err
}

// DeleteAll deletes all messages stored in MongoDB
func (mongo *MongoDB) DeleteAll() error {
	_, err := mongo.Collection.RemoveAll(bson.M{})
	return err
}

// Load loads an individual message by storage ID
func (mongo *MongoDB) Load(id string) (*data.Message, error) {
	result := &data.Message{}
	err := mongo.Collection.Find(bson.M{"id": id}).One(&result)
	if err != nil {
		log.Printf("Error loading message: %s", err)
		return nil, err
	}
	return result, nil
}
<|endoftext|>


<commit_before>package managexe

import (
	"context"
	"fmt"

	"golang.org/x/sync/errgroup"
)

type ErrPanic struct {
	Err interface{}
}

func (ep ErrPanic) Error() string {
	return fmt.Sprint(ep.Err)
}

type Execer interface {
	Exec(context.Context) error
}

type Manager struct {
	ch      chan Execer
	errCh   chan error
	workerN int
}

func (m *Manager) NumTask() int {
	return len(m.ch)
}

func (m *Manager) ErrCh() <-chan error {
	return m.errCh
}

func (m *Manager) AddTask(task Execer) {
	m.ch <- task
}

func (m *Manager) Run(ctx context.Context) {
	eg, egctx := errgroup.WithContext(ctx)

	for i := 0; i < m.workerN; i++ {
		eg.Go(func() error {
			for task := range m.ch {
				select {
				case <-egctx.Done():
					return egctx.Err()
				default:
					func() {
						defer func() {
							if err := recover(); err != nil {
								m.errCh <- ErrPanic{err}
							}
						}()

						if err := task.Exec(ctx); err != nil {
							// TODO: error
							m.errCh <- err
						}
					}()
				}
			}

			return nil
		})
	}

	eg.Wait()
}

func NewManager(workerN, bufferN int) *Manager {
	return &Manager{
		ch:      make(chan Execer, bufferN),
		errCh:   make(chan error, bufferN),
		workerN: workerN,
	}
}
<commit_msg>雑に一時停止・再開を実装<commit_after>package managexe

import (
	"context"
	"fmt"
	"sync"

	"golang.org/x/sync/errgroup"
)

type ErrPanic struct {
	Err interface{}
}

func (ep ErrPanic) Error() string {
	return fmt.Sprint(ep.Err)
}

type Execer interface {
	Exec(context.Context) error
}

type Manager struct {
	sync.Mutex
	ch       chan Execer
	errCh    chan error
	pauseCh  chan struct{}
	isPaused bool
	workerN  int
}

func (m *Manager) NumTask() int {
	return len(m.ch)
}

func (m *Manager) ErrCh() <-chan error {
	return m.errCh
}

func (m *Manager) Pause() {
	m.Lock()
	defer m.Unlock()
	if m.isPaused {
		return
	}

	m.isPaused = true
	m.pauseCh = make(chan struct{})
}

func (m *Manager) Resume() {
	m.Lock()
	defer m.Unlock()
	if !m.isPaused {
		return
	}

	m.isPaused = false
	close(m.pauseCh)
}

func (m *Manager) AddTask(task Execer) {
	m.ch <- task
}

func (m *Manager) Run(ctx context.Context) {
	eg, egctx := errgroup.WithContext(ctx)

	close(m.pauseCh)
	m.isPaused = false

	for i := 0; i < m.workerN; i++ {
		eg.Go(func() error {
			for task := range m.ch {
				<-m.pauseCh

				select {
				case <-egctx.Done():
					return egctx.Err()
				default:
					func() {
						defer func() {
							if err := recover(); err != nil {
								m.errCh <- ErrPanic{err}
							}
						}()

						if err := task.Exec(ctx); err != nil {
							// TODO: error
							m.errCh <- err
						}
					}()
				}
			}

			return nil
		})
	}

	eg.Wait()
}

func NewManager(workerN, bufferN int) *Manager {
	return &Manager{
		ch:       make(chan Execer, bufferN),
		errCh:    make(chan error, bufferN),
		isPaused: true,
		workerN:  workerN,
	}
}
<|endoftext|>


<commit_before>package runc

import (
	"os/exec"
	"syscall"
)

var Monitor ProcessMonitor = &defaultMonitor{}

// ProcessMonitor is an interface for process monitoring
//
// It allows daemons using go-runc to have a SIGCHLD handler
// to handle exits without introducing races between the handler
// and go's exec.Cmd
// These methods should match the methods exposed by exec.Cmd to provide
// a consistent experience for the caller
type ProcessMonitor interface {
	Output(*exec.Cmd) ([]byte, error)
	CombinedOutput(*exec.Cmd) ([]byte, error)
	Run(*exec.Cmd) error
	Start(*exec.Cmd) error
	Wait(*exec.Cmd) (int, error)
}

type defaultMonitor struct {
}

func (m *defaultMonitor) Output(c *exec.Cmd) ([]byte, error) {
	return c.Output()
}

func (m *defaultMonitor) CombinedOutput(c *exec.Cmd) ([]byte, error) {
	return c.CombinedOutput()
}

func (m *defaultMonitor) Run(c *exec.Cmd) error {
	return c.Run()
}

func (m *defaultMonitor) Start(c *exec.Cmd) error {
	return c.Start()
}

func (m *defaultMonitor) Wait(c *exec.Cmd) (int, error) {
	status, err := c.Process.Wait()
	if err != nil {
		return -1, err
	}
	return status.Sys().(syscall.WaitStatus).ExitStatus(), nil
}
<commit_msg>Use cmd.Wait in monitor<commit_after>package runc

import (
	"os/exec"
	"syscall"
)

var Monitor ProcessMonitor = &defaultMonitor{}

// ProcessMonitor is an interface for process monitoring
//
// It allows daemons using go-runc to have a SIGCHLD handler
// to handle exits without introducing races between the handler
// and go's exec.Cmd
// These methods should match the methods exposed by exec.Cmd to provide
// a consistent experience for the caller
type ProcessMonitor interface {
	Output(*exec.Cmd) ([]byte, error)
	CombinedOutput(*exec.Cmd) ([]byte, error)
	Run(*exec.Cmd) error
	Start(*exec.Cmd) error
	Wait(*exec.Cmd) (int, error)
}

type defaultMonitor struct {
}

func (m *defaultMonitor) Output(c *exec.Cmd) ([]byte, error) {
	return c.Output()
}

func (m *defaultMonitor) CombinedOutput(c *exec.Cmd) ([]byte, error) {
	return c.CombinedOutput()
}

func (m *defaultMonitor) Run(c *exec.Cmd) error {
	return c.Run()
}

func (m *defaultMonitor) Start(c *exec.Cmd) error {
	return c.Start()
}

func (m *defaultMonitor) Wait(c *exec.Cmd) (int, error) {
	if err := c.Wait(); err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				return status.ExitStatus(), nil
			}
		}
		return -1, err
	}
	return 0, nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"github.com/samalba/dockerclient"
	"gopkg.in/BlueDragonX/simplelog.v1"
	"strings"
	"time"
)

// Monitor docker for service changes and emit events.
type ServiceMonitor struct {
	client       *dockerclient.DockerClient
	hostname     string
	tags         map[string]bool
	configVar    string
	tagsVar      string
	state        int32
	containers   map[string]bool
	services     map[string]*Service
	pollInterval time.Duration
	stop         chan bool
	log          *simplelog.Logger
}

// Create a new service monitor listening on the given URL. Look for service
// config in the Docker environment variable names configVar.
func NewServiceMonitor(url, hostname string, tags []string, configVar, tagsVar string, pollInterval time.Duration, log *simplelog.Logger) (mon *ServiceMonitor, err error) {
	mon = &ServiceMonitor{}
	mon.client, err = dockerclient.NewDockerClient(url, nil)
	mon.hostname = hostname
	mon.tags = make(map[string]bool)
	mon.configVar = configVar
	mon.tagsVar = tagsVar
	mon.state = Stopped
	mon.pollInterval = pollInterval
	mon.stop = make(chan bool)
	mon.log = log

	if tags != nil {
		for _, tag := range tags {
			if tag != "" {
				mon.tags[tag] = true
			}
		}
	}
	return
}

func (mon *ServiceMonitor) addContainer(serviceEvents chan ServiceEvent, containerId string) {
	errorFmt := "container %.12s: %s"
	var err error
	var containerInfo *dockerclient.ContainerInfo
	if containerInfo, err = mon.client.InspectContainer(containerId); err != nil {
		mon.log.Error(errorFmt, containerId, err)
		return
	}

	configEnv := ""
	tagsEnv := ""
	for _, envVar := range containerInfo.Config.Env {
		envName, envValue := parseEnv(envVar)
		if envName == mon.configVar {
			configEnv = envValue
		} else if envName == mon.tagsVar {
			tagsEnv = envValue
		}
	}

	if configEnv == "" {
		mon.log.Debug(errorFmt, containerId, "no services defined, skipping")
		return
	}

	tags := parseTags(tagsEnv)
	if len(mon.tags) > 0 {
		found := false
		for _, tag := range tags {
			if _, found = mon.tags[tag]; found {
				break
			}
		}
		if !found {
			mon.log.Debug(errorFmt, containerId, "not tagged, skipping")
			return
		}
	}

	configValues := strings.Split(configEnv, ",")
	for _, configValue := range configValues {
		svc := &Service{}
		if err = svc.loadConfig(configValue); err != nil {
			mon.log.Warn(errorFmt, containerId, err)
			return
		}
		if err = svc.loadInfo(containerInfo, mon.hostname); err != nil {
			mon.log.Warn(errorFmt, containerId, err)
			return
		}

		oldSvc, update := mon.services[svc.Hash()]
		if update {
			serviceEvents <- ServiceEvent{Heartbeat, svc}
		} else if update && *svc == *oldSvc {
			serviceEvents <- ServiceEvent{Update, svc}
		} else {
			serviceEvents <- ServiceEvent{Add, svc}
		}
		mon.services[svc.Hash()] = svc
	}
}

func (mon *ServiceMonitor) removeContainer(serviceEvents chan ServiceEvent, containerId string) {
	remove := []string{}
	for hash, svc := range mon.services {
		if svc.ContainerId == containerId {
			remove = append(remove, hash)
		}
	}

	for _, hash := range remove {
		serviceEvents <- ServiceEvent{Remove, mon.services[hash]}
		delete(mon.services, hash)
	}
}

func (mon *ServiceMonitor) poll(serviceEvents chan ServiceEvent) {
	var err error
	var containers []dockerclient.Container
	if containers, err = mon.client.ListContainers(false); err != nil {
		mon.log.Error("polling failed: %s", err)
		return
	}

	mon.log.Debug("polling for containers")

	containerIds := make(map[string]bool, len(containers))
	for _, container := range containers {
		mon.addContainer(serviceEvents, container.Id)
		containerIds[container.Id] = true
	}

	for id := range mon.containers {
		if _, ok := containerIds[id]; !ok {
			mon.removeContainer(serviceEvents, id)
		}
	}
	mon.containers = containerIds
}

func (mon *ServiceMonitor) Listen(serviceEvents chan ServiceEvent) error {
	if !stateListening(&mon.state) {
		return errors.New("already listening")
	}

	mon.containers = make(map[string]bool)
	mon.services = make(map[string]*Service)
	containerEvents := make(chan ContainerEvent, 1)

	cb := func(e *dockerclient.Event, args ...interface{}) {
		if e.Status == "start" {
			containerEvents <- ContainerEvent{Add, e.Id}
		} else if e.Status == "die" {
			containerEvents <- ContainerEvent{Remove, e.Id}
		}
	}

	mon.poll(serviceEvents)
	mon.client.StartMonitorEvents(cb)

Loop:
	for {
		select {
		case e := <-containerEvents:
			switch e.State {
			case Add:
				mon.addContainer(serviceEvents, e.ContainerId)
			case Remove:
				mon.removeContainer(serviceEvents, e.ContainerId)
			}
		case <-time.After(mon.pollInterval):
			mon.poll(serviceEvents)
		case <-mon.stop:
			break Loop
		}
	}

	mon.client.StopAllMonitorEvents()
	for _, service := range mon.services {
		serviceEvents <- ServiceEvent{Remove, service}
	}
	close(serviceEvents)

	stateStopped(&mon.state)
	return nil
}

func (mon *ServiceMonitor) Stop() error {
	if !stateStopping(&mon.state) {
		return errors.New("not listening")
	}
	mon.stop <- true
	return nil
}
<commit_msg>Fix changed DockerClient ListContainers call.<commit_after>package main

import (
	"errors"
	"github.com/samalba/dockerclient"
	"gopkg.in/BlueDragonX/simplelog.v1"
	"strings"
	"time"
)

// Monitor docker for service changes and emit events.
type ServiceMonitor struct {
	client       *dockerclient.DockerClient
	hostname     string
	tags         map[string]bool
	configVar    string
	tagsVar      string
	state        int32
	containers   map[string]bool
	services     map[string]*Service
	pollInterval time.Duration
	stop         chan bool
	log          *simplelog.Logger
}

// Create a new service monitor listening on the given URL. Look for service
// config in the Docker environment variable names configVar.
func NewServiceMonitor(url, hostname string, tags []string, configVar, tagsVar string, pollInterval time.Duration, log *simplelog.Logger) (mon *ServiceMonitor, err error) {
	mon = &ServiceMonitor{}
	mon.client, err = dockerclient.NewDockerClient(url, nil)
	mon.hostname = hostname
	mon.tags = make(map[string]bool)
	mon.configVar = configVar
	mon.tagsVar = tagsVar
	mon.state = Stopped
	mon.pollInterval = pollInterval
	mon.stop = make(chan bool)
	mon.log = log

	if tags != nil {
		for _, tag := range tags {
			if tag != "" {
				mon.tags[tag] = true
			}
		}
	}
	return
}

func (mon *ServiceMonitor) addContainer(serviceEvents chan ServiceEvent, containerId string) {
	errorFmt := "container %.12s: %s"
	var err error
	var containerInfo *dockerclient.ContainerInfo
	if containerInfo, err = mon.client.InspectContainer(containerId); err != nil {
		mon.log.Error(errorFmt, containerId, err)
		return
	}

	configEnv := ""
	tagsEnv := ""
	for _, envVar := range containerInfo.Config.Env {
		envName, envValue := parseEnv(envVar)
		if envName == mon.configVar {
			configEnv = envValue
		} else if envName == mon.tagsVar {
			tagsEnv = envValue
		}
	}

	if configEnv == "" {
		mon.log.Debug(errorFmt, containerId, "no services defined, skipping")
		return
	}

	tags := parseTags(tagsEnv)
	if len(mon.tags) > 0 {
		found := false
		for _, tag := range tags {
			if _, found = mon.tags[tag]; found {
				break
			}
		}
		if !found {
			mon.log.Debug(errorFmt, containerId, "not tagged, skipping")
			return
		}
	}

	configValues := strings.Split(configEnv, ",")
	for _, configValue := range configValues {
		svc := &Service{}
		if err = svc.loadConfig(configValue); err != nil {
			mon.log.Warn(errorFmt, containerId, err)
			return
		}
		if err = svc.loadInfo(containerInfo, mon.hostname); err != nil {
			mon.log.Warn(errorFmt, containerId, err)
			return
		}

		oldSvc, update := mon.services[svc.Hash()]
		if update {
			serviceEvents <- ServiceEvent{Heartbeat, svc}
		} else if update && *svc == *oldSvc {
			serviceEvents <- ServiceEvent{Update, svc}
		} else {
			serviceEvents <- ServiceEvent{Add, svc}
		}
		mon.services[svc.Hash()] = svc
	}
}

func (mon *ServiceMonitor) removeContainer(serviceEvents chan ServiceEvent, containerId string) {
	remove := []string{}
	for hash, svc := range mon.services {
		if svc.ContainerId == containerId {
			remove = append(remove, hash)
		}
	}

	for _, hash := range remove {
		serviceEvents <- ServiceEvent{Remove, mon.services[hash]}
		delete(mon.services, hash)
	}
}

func (mon *ServiceMonitor) poll(serviceEvents chan ServiceEvent) {
	var err error
	var containers []dockerclient.Container
	if containers, err = mon.client.ListContainers(false, false); err != nil {
		mon.log.Error("polling failed: %s", err)
		return
	}

	mon.log.Debug("polling for containers")

	containerIds := make(map[string]bool, len(containers))
	for _, container := range containers {
		mon.addContainer(serviceEvents, container.Id)
		containerIds[container.Id] = true
	}

	for id := range mon.containers {
		if _, ok := containerIds[id]; !ok {
			mon.removeContainer(serviceEvents, id)
		}
	}
	mon.containers = containerIds
}

func (mon *ServiceMonitor) Listen(serviceEvents chan ServiceEvent) error {
	if !stateListening(&mon.state) {
		return errors.New("already listening")
	}

	mon.containers = make(map[string]bool)
	mon.services = make(map[string]*Service)
	containerEvents := make(chan ContainerEvent, 1)

	cb := func(e *dockerclient.Event, args ...interface{}) {
		if e.Status == "start" {
			containerEvents <- ContainerEvent{Add, e.Id}
		} else if e.Status == "die" {
			containerEvents <- ContainerEvent{Remove, e.Id}
		}
	}

	mon.poll(serviceEvents)
	mon.client.StartMonitorEvents(cb)

Loop:
	for {
		select {
		case e := <-containerEvents:
			switch e.State {
			case Add:
				mon.addContainer(serviceEvents, e.ContainerId)
			case Remove:
				mon.removeContainer(serviceEvents, e.ContainerId)
			}
		case <-time.After(mon.pollInterval):
			mon.poll(serviceEvents)
		case <-mon.stop:
			break Loop
		}
	}

	mon.client.StopAllMonitorEvents()
	for _, service := range mon.services {
		serviceEvents <- ServiceEvent{Remove, service}
	}
	close(serviceEvents)

	stateStopped(&mon.state)
	return nil
}

func (mon *ServiceMonitor) Stop() error {
	if !stateStopping(&mon.state) {
		return errors.New("not listening")
	}
	mon.stop <- true
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package configfiles

import (
	"fmt"
	"path/filepath"
	"testing"

	"github.com/pkg/errors"

	apiequality "k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/runtime"
	kubeletconfigv1beta1 "k8s.io/kubelet/config/v1beta1"
	kubeletconfig "k8s.io/kubernetes/pkg/kubelet/apis/config"
	kubeletscheme "k8s.io/kubernetes/pkg/kubelet/apis/config/scheme"
	utilfiles "k8s.io/kubernetes/pkg/kubelet/kubeletconfig/util/files"
	utiltest "k8s.io/kubernetes/pkg/kubelet/kubeletconfig/util/test"
	utilfs "k8s.io/kubernetes/pkg/util/filesystem"
)

const configDir = "/test-config-dir"
const relativePath = "relative/path/test"
const kubeletFile = "kubelet"

func TestLoad(t *testing.T) {
	cases := []struct {
		desc      string
		file      *string
		expect    *kubeletconfig.KubeletConfiguration
		err       string
		strictErr bool
	}{
		// missing file
		{
			desc: "missing file",
			err:  "failed to read",
		},
		// empty file
		{
			desc: "empty file",
			file: newString(``),
			err:  "was empty",
		},
		// invalid format
		{
			desc: "invalid yaml",
			file: newString(`*`),
			err:  "failed to decode",
		},
		{
			desc: "invalid json",
			file: newString(`{*`),
			err:  "failed to decode",
		},
		// invalid object
		{
			desc: "missing kind",
			file: newString(`{"apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			err:  "failed to decode",
		},
		{
			desc: "missing version",
			file: newString(`{"kind":"KubeletConfiguration"}`),
			err:  "failed to decode",
		},
		{
			desc: "unregistered kind",
			file: newString(`{"kind":"BogusKind","apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			err:  "failed to decode",
		},
		{
			desc: "unregistered version",
			file: newString(`{"kind":"KubeletConfiguration","apiVersion":"bogusversion"}`),
			err:  "failed to decode",
		},

		// empty object with correct kind and version should result in the defaults for that kind and version
		{
			desc: "default from yaml",
			file: newString(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1`),
			expect: newConfig(t),
		},
		{
			desc:   "default from json",
			file:   newString(`{"kind":"KubeletConfiguration","apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			expect: newConfig(t),
		},

		// relative path
		{
			desc: "yaml, relative path is resolved",
			file: newString(fmt.Sprintf(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
staticPodPath: %s`, relativePath)),
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath)
				return kc
			}(),
		},
		{
			desc: "json, relative path is resolved",
			file: newString(fmt.Sprintf(`{"kind":"KubeletConfiguration","apiVersion":"kubelet.config.k8s.io/v1beta1","staticPodPath":"%s"}`, relativePath)),
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath)
				return kc
			}(),
		},
		{
			// This should fail from v1beta2+
			desc: "duplicate field",
			file: newString(fmt.Sprintf(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
staticPodPath: %s
staticPodPath: %s/foo`, relativePath, relativePath)),
			// err:       `key "staticPodPath" already set`,
			// strictErr: true,
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath, "foo")
				return kc
			}(),
		},
		{
			// This should fail from v1beta2+
			desc: "unknown field",
			file: newString(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
foo: bar`),
			// err:       "found unknown field: foo",
			// strictErr: true,
			expect: newConfig(t),
		},
	}

	for _, c := range cases {
		t.Run(c.desc, func(t *testing.T) {
			fs := utilfs.NewFakeFs()
			path := filepath.Join(configDir, kubeletFile)
			if c.file != nil {
				if err := addFile(fs, path, *c.file); err != nil {
					t.Fatalf("unexpected error: %v", err)
				}
			}
			loader, err := NewFsLoader(fs, path)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			kc, err := loader.Load()

			if c.strictErr && !runtime.IsStrictDecodingError(errors.Cause(err)) {
				t.Fatalf("got error: %v, want strict decoding error", err)
			}
			if utiltest.SkipRest(t, c.desc, err, c.err) {
				return
			}
			if !apiequality.Semantic.DeepEqual(c.expect, kc) {
				t.Fatalf("expect %#v but got %#v", *c.expect, *kc)
			}
		})
	}
}

func TestResolveRelativePaths(t *testing.T) {
	absolutePath := filepath.Join(configDir, "absolute")
	cases := []struct {
		desc   string
		path   string
		expect string
	}{
		{"empty path", "", ""},
		{"absolute path", absolutePath, absolutePath},
		{"relative path", relativePath, filepath.Join(configDir, relativePath)},
	}

	paths := kubeletconfig.KubeletConfigurationPathRefs(newConfig(t))
	if len(paths) == 0 {
		t.Fatalf("requires at least one path field to exist in the KubeletConfiguration type")
	}
	for _, c := range cases {
		t.Run(c.desc, func(t *testing.T) {
			// set the path, resolve it, and check if it resolved as we would expect
			*(paths[0]) = c.path
			resolveRelativePaths(paths, configDir)
			if *(paths[0]) != c.expect {
				t.Fatalf("expect %s but got %s", c.expect, *(paths[0]))
			}
		})
	}
}

func newString(s string) *string {
	return &s
}

func addFile(fs utilfs.Filesystem, path string, file string) error {
	if err := utilfiles.EnsureDir(fs, filepath.Dir(path)); err != nil {
		return err
	}
	return utilfiles.ReplaceFile(fs, path, []byte(file))
}

func newConfig(t *testing.T) *kubeletconfig.KubeletConfiguration {
	kubeletScheme, _, err := kubeletscheme.NewSchemeAndCodecs()
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// get the built-in default configuration
	external := &kubeletconfigv1beta1.KubeletConfiguration{}
	kubeletScheme.Default(external)
	kc := &kubeletconfig.KubeletConfiguration{}
	err = kubeletScheme.Convert(external, kc, nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	return kc
}
<commit_msg>Removed usage of github.com/pkg/errors<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package configfiles

import (
	"errors"
	"fmt"
	"path/filepath"
	"testing"

	apiequality "k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/runtime"
	kubeletconfigv1beta1 "k8s.io/kubelet/config/v1beta1"
	kubeletconfig "k8s.io/kubernetes/pkg/kubelet/apis/config"
	kubeletscheme "k8s.io/kubernetes/pkg/kubelet/apis/config/scheme"
	utilfiles "k8s.io/kubernetes/pkg/kubelet/kubeletconfig/util/files"
	utiltest "k8s.io/kubernetes/pkg/kubelet/kubeletconfig/util/test"
	utilfs "k8s.io/kubernetes/pkg/util/filesystem"
)

const configDir = "/test-config-dir"
const relativePath = "relative/path/test"
const kubeletFile = "kubelet"

func TestLoad(t *testing.T) {
	cases := []struct {
		desc      string
		file      *string
		expect    *kubeletconfig.KubeletConfiguration
		err       string
		strictErr bool
	}{
		// missing file
		{
			desc: "missing file",
			err:  "failed to read",
		},
		// empty file
		{
			desc: "empty file",
			file: newString(``),
			err:  "was empty",
		},
		// invalid format
		{
			desc: "invalid yaml",
			file: newString(`*`),
			err:  "failed to decode",
		},
		{
			desc: "invalid json",
			file: newString(`{*`),
			err:  "failed to decode",
		},
		// invalid object
		{
			desc: "missing kind",
			file: newString(`{"apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			err:  "failed to decode",
		},
		{
			desc: "missing version",
			file: newString(`{"kind":"KubeletConfiguration"}`),
			err:  "failed to decode",
		},
		{
			desc: "unregistered kind",
			file: newString(`{"kind":"BogusKind","apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			err:  "failed to decode",
		},
		{
			desc: "unregistered version",
			file: newString(`{"kind":"KubeletConfiguration","apiVersion":"bogusversion"}`),
			err:  "failed to decode",
		},

		// empty object with correct kind and version should result in the defaults for that kind and version
		{
			desc: "default from yaml",
			file: newString(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1`),
			expect: newConfig(t),
		},
		{
			desc:   "default from json",
			file:   newString(`{"kind":"KubeletConfiguration","apiVersion":"kubelet.config.k8s.io/v1beta1"}`),
			expect: newConfig(t),
		},

		// relative path
		{
			desc: "yaml, relative path is resolved",
			file: newString(fmt.Sprintf(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
staticPodPath: %s`, relativePath)),
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath)
				return kc
			}(),
		},
		{
			desc: "json, relative path is resolved",
			file: newString(fmt.Sprintf(`{"kind":"KubeletConfiguration","apiVersion":"kubelet.config.k8s.io/v1beta1","staticPodPath":"%s"}`, relativePath)),
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath)
				return kc
			}(),
		},
		{
			// This should fail from v1beta2+
			desc: "duplicate field",
			file: newString(fmt.Sprintf(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
staticPodPath: %s
staticPodPath: %s/foo`, relativePath, relativePath)),
			// err:       `key "staticPodPath" already set`,
			// strictErr: true,
			expect: func() *kubeletconfig.KubeletConfiguration {
				kc := newConfig(t)
				kc.StaticPodPath = filepath.Join(configDir, relativePath, "foo")
				return kc
			}(),
		},
		{
			// This should fail from v1beta2+
			desc: "unknown field",
			file: newString(`kind: KubeletConfiguration
apiVersion: kubelet.config.k8s.io/v1beta1
foo: bar`),
			// err:       "found unknown field: foo",
			// strictErr: true,
			expect: newConfig(t),
		},
	}

	for _, c := range cases {
		t.Run(c.desc, func(t *testing.T) {
			fs := utilfs.NewFakeFs()
			path := filepath.Join(configDir, kubeletFile)
			if c.file != nil {
				if err := addFile(fs, path, *c.file); err != nil {
					t.Fatalf("unexpected error: %v", err)
				}
			}
			loader, err := NewFsLoader(fs, path)
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			kc, err := loader.Load()

			if c.strictErr && !runtime.IsStrictDecodingError(errors.Unwrap(err)) {
				t.Fatalf("got error: %v, want strict decoding error", err)
			}
			if utiltest.SkipRest(t, c.desc, err, c.err) {
				return
			}
			if !apiequality.Semantic.DeepEqual(c.expect, kc) {
				t.Fatalf("expect %#v but got %#v", *c.expect, *kc)
			}
		})
	}
}

func TestResolveRelativePaths(t *testing.T) {
	absolutePath := filepath.Join(configDir, "absolute")
	cases := []struct {
		desc   string
		path   string
		expect string
	}{
		{"empty path", "", ""},
		{"absolute path", absolutePath, absolutePath},
		{"relative path", relativePath, filepath.Join(configDir, relativePath)},
	}

	paths := kubeletconfig.KubeletConfigurationPathRefs(newConfig(t))
	if len(paths) == 0 {
		t.Fatalf("requires at least one path field to exist in the KubeletConfiguration type")
	}
	for _, c := range cases {
		t.Run(c.desc, func(t *testing.T) {
			// set the path, resolve it, and check if it resolved as we would expect
			*(paths[0]) = c.path
			resolveRelativePaths(paths, configDir)
			if *(paths[0]) != c.expect {
				t.Fatalf("expect %s but got %s", c.expect, *(paths[0]))
			}
		})
	}
}

func newString(s string) *string {
	return &s
}

func addFile(fs utilfs.Filesystem, path string, file string) error {
	if err := utilfiles.EnsureDir(fs, filepath.Dir(path)); err != nil {
		return err
	}
	return utilfiles.ReplaceFile(fs, path, []byte(file))
}

func newConfig(t *testing.T) *kubeletconfig.KubeletConfiguration {
	kubeletScheme, _, err := kubeletscheme.NewSchemeAndCodecs()
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// get the built-in default configuration
	external := &kubeletconfigv1beta1.KubeletConfiguration{}
	kubeletScheme.Default(external)
	kc := &kubeletconfig.KubeletConfiguration{}
	err = kubeletScheme.Convert(external, kc, nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	return kc
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

func main() {
	var bindCommands []*exec.Cmd

	bindings := GetAllBindings()
	for _, binding := range bindings {
		bindCommands = append(bindCommands, CreateCommand(binding))
	}

	err := RunCommands(bindCommands)
	if err != nil {
		panic(err)
	}
}

type Credentials struct {
	Host string `json:"host"`
	Port int    `json:"port"`
	User string `json:"user"`
	Pass string `json:"password"`
}

type Binding struct {
	Credentials Credentials `json:"credentials"`
	Plan        string      `json:"plan"`
	Name        string      `json:"name"`
}

func GetAllBindings() []Binding {
	sshfsBindings := make(map[string][]Binding)
	jsonString := os.Getenv("VCAP_SERVICES")

	err := json.Unmarshal([]byte(jsonString), &sshfsBindings)
	if err != nil {
		panic("Unable to decode JSON in env var VCAP_SERVICES")
	}
	return sshfsBindings["sshfs"]
}

func CreateCommand(binding Binding) *exec.Cmd {
	cmd := exec.Command(
		"sshfs",
		fmt.Sprintf("%s@%s:", binding.Credentials.User, binding.Credentials.Host),
		"-p", strconv.Itoa(binding.Credentials.Port),
		"-o", "password_stdin",
		"-o", "StrictHostKeyChecking=false",
		binding.Name,
	)

	cmd.Stdin = strings.NewReader(binding.Credentials.Pass + "\n")

	return cmd
}

func RunCommands(commands []*exec.Cmd) error {
	for _, command := range commands {
		commandOutput, err := command.CombinedOutput()
		if err != nil {
			return errors.New(fmt.Sprintf(
				"Failed while running %s : %s : %s",
				command,
				commandOutput,
				err.Error()))
		}
	}

	return nil
}
<commit_msg>Adds some logging<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

func main() {
	var bindCommands []*exec.Cmd

	bindings := GetAllBindings()
	for _, binding := range bindings {
		bindCommands = append(bindCommands, CreateCommand(binding))
	}
	fmt.Printf("Found %d bindings to SSHFS service", len(bindings))

	if len(bindings) > 0 {
		err := RunCommands(bindCommands)
		if err != nil {
			panic(err)
		}
		fmt.Printf("Finished mounting SSHFS file systems")
	}
}

type Credentials struct {
	Host string `json:"host"`
	Port int    `json:"port"`
	User string `json:"user"`
	Pass string `json:"password"`
}

type Binding struct {
	Credentials Credentials `json:"credentials"`
	Plan        string      `json:"plan"`
	Name        string      `json:"name"`
}

func GetAllBindings() []Binding {
	sshfsBindings := make(map[string][]Binding)
	jsonString := os.Getenv("VCAP_SERVICES")

	err := json.Unmarshal([]byte(jsonString), &sshfsBindings)
	if err != nil {
		panic("Unable to decode JSON in env var VCAP_SERVICES")
	}
	return sshfsBindings["sshfs"]
}

func CreateCommand(binding Binding) *exec.Cmd {
	cmd := exec.Command(
		"sshfs",
		fmt.Sprintf("%s@%s:", binding.Credentials.User, binding.Credentials.Host),
		"-p", strconv.Itoa(binding.Credentials.Port),
		"-o", "password_stdin",
		"-o", "StrictHostKeyChecking=false",
		binding.Name,
	)

	cmd.Stdin = strings.NewReader(binding.Credentials.Pass + "\n")

	return cmd
}

func RunCommands(commands []*exec.Cmd) error {
	for _, command := range commands {
		commandOutput, err := command.CombinedOutput()
		if err != nil {
			return errors.New(fmt.Sprintf(
				"Failed while running %s : %s : %s",
				command,
				commandOutput,
				err.Error()))
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package sql

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"strings"

	"github.com/marcuswestin/fun-go/errs"
	"github.com/marcuswestin/fun-go/random"
)

type ShardSet struct {
	username        string
	password        string
	host            string
	port            int
	dbNamePrefix    string
	numShards       int
	maxShards       int
	maxConns        int
	shards          []*Shard
	beginEndHandler func() (func(), error)
	metricsHandler  func(query string, shardName string) func()
	log             *log.Logger
}

func WithBeginEndHandler(handler func() (func(), error)) func(*ShardSet) {
	return func(s *ShardSet) {
		s.beginEndHandler = handler
	}
}
func WithMetricsHandler(handler func(query string, shardName string) func()) func(*ShardSet) {
	return func(s *ShardSet) {
		s.metricsHandler = handler
	}
}
func WithLogger(logger *log.Logger) func(*ShardSet) {
	return func(s *ShardSet) {
		s.log = logger
	}
}
func NewShardSet(username string, password string, host string, port int, dbNamePrefix string, numShards int, maxShards int, maxConns int, options ...func(*ShardSet)) *ShardSet {
	s := &ShardSet{
		username:     username,
		password:     password,
		host:         host,
		port:         port,
		dbNamePrefix: dbNamePrefix,
		numShards:    numShards,
		maxShards:    maxShards,
		maxConns:     maxConns,
		log:          log.New(ioutil.Discard, "", 0),
	}
	for _, option := range options {
		option(s)
	}
	return s
}

func (s *ShardSet) Connect() (err errs.Err) {
	s.shards = make([]*Shard, s.numShards)
	for i := 0; i < s.numShards; i++ {
		err = s.addShard(i)
		if err != nil {
			return
		}
	}
	return
}

func (s *ShardSet) Shard(id int64) *Shard {
	if id == 0 {
		panic("Bad shard index id 0")
	}
	shardIndex := ((id - 1) % int64(s.maxShards)) // 1->0, 2->1, 3->2 ..., 65000->65000, 65001->0, 65002->1
	return s.shards[shardIndex]
}

func (s *ShardSet) All() []*Shard {
	all := make([]*Shard, len(s.shards))
	for i, shard := range s.shards {
		all[i] = shard
	}
	return all
}

func (s *ShardSet) RandomShard() *Shard {
	return s.shards[random.Between(0, len(s.shards))]
}

func (s *ShardSet) addShard(i int) (err errs.Err) {
	autoIncrementOffset := i + 1
	dbName := fmt.Sprint(s.dbNamePrefix, autoIncrementOffset)
	s.shards[i], err = newShard(s, dbName, autoIncrementOffset)
	if err != nil {
		return
	}
	return
}

func newShard(s *ShardSet, dbName string, autoIncrementOffset int) (*Shard, errs.Err) {
	connVars := ConnVariables{
		"autocommit":               "true",
		"clientFoundRows":          "true",
		"charset":                  "utf8mb4",
		"collation":                "utf8_unicode_ci",
		"auto_increment_increment": strconv.Itoa(s.maxShards),
		"auto_increment_offset":    strconv.Itoa(autoIncrementOffset),
		"sql_mode":                 "STRICT_ALL_TABLES",
	}

	db, err := dbOpener(s.username, s.password, dbName, s.host, s.port, connVars)
	if err != nil {
		return nil, err
	}

	db.SetMaxOpenConns(s.maxConns)
	// db.SetMaxIdleConns(n)
	stdErr := db.Ping()
	if stdErr != nil {
		return nil, errs.Wrap(stdErr, nil)
	}
	return &Shard{DBName: dbName, db: db, sqlConn: db, BeginEndHandler: s.beginEndHandler, MetricsHandler: s.metricsHandler, log: s.log}, nil
}

func SetOpener(opener Opener) {
	if dbOpener != nil {
		panic("Opener already set - did you import two driver adapters?")
	}
	dbOpener = opener
}

type ConnVariables map[string]string

func (connVars ConnVariables) Join(sep string) string {
	kvps := make([]string, len(connVars))
	i := 0
	for param, val := range connVars {
		kvps[i] = param + "=" + val
		i += 1
	}
	return strings.Join(kvps, sep)
}

type Opener func(username, password, dbName, host string, port int, connVars ConnVariables) (*sql.DB, errs.Err)

var dbOpener Opener
<commit_msg>CB-892: pass down tls option to mysql driver (#7)<commit_after>package sql

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"strings"

	"github.com/marcuswestin/fun-go/errs"
	"github.com/marcuswestin/fun-go/random"
)

type ShardSet struct {
	username string
	password string
	host     string
	port     int
	// tls has to be a string b/c we need to be able to pass in `skip-verify`
	// more info: https://github.com/go-sql-driver/mysql#tls
	tls             string
	dbNamePrefix    string
	numShards       int
	maxShards       int
	maxConns        int
	shards          []*Shard
	beginEndHandler func() (func(), error)
	metricsHandler  func(query string, shardName string) func()
	log             *log.Logger
}

func WithBeginEndHandler(handler func() (func(), error)) func(*ShardSet) {
	return func(s *ShardSet) {
		s.beginEndHandler = handler
	}
}
func WithMetricsHandler(handler func(query string, shardName string) func()) func(*ShardSet) {
	return func(s *ShardSet) {
		s.metricsHandler = handler
	}
}
func WithLogger(logger *log.Logger) func(*ShardSet) {
	return func(s *ShardSet) {
		s.log = logger
	}
}

func WithTLS(tls string) func(*ShardSet) {
	return func(s *ShardSet) {
		s.tls = tls
	}
}

func NewShardSet(username string, password string, host string, port int, dbNamePrefix string, numShards int, maxShards int, maxConns int, options ...func(*ShardSet)) *ShardSet {
	s := &ShardSet{
		username:     username,
		password:     password,
		host:         host,
		port:         port,
		dbNamePrefix: dbNamePrefix,
		numShards:    numShards,
		maxShards:    maxShards,
		maxConns:     maxConns,
		log:          log.New(ioutil.Discard, "", 0),
	}
	for _, option := range options {
		option(s)
	}
	return s
}

func (s *ShardSet) Connect() (err errs.Err) {
	s.shards = make([]*Shard, s.numShards)
	for i := 0; i < s.numShards; i++ {
		err = s.addShard(i)
		if err != nil {
			return
		}
	}
	return
}

func (s *ShardSet) Shard(id int64) *Shard {
	if id == 0 {
		panic("Bad shard index id 0")
	}
	shardIndex := ((id - 1) % int64(s.maxShards)) // 1->0, 2->1, 3->2 ..., 65000->65000, 65001->0, 65002->1
	return s.shards[shardIndex]
}

func (s *ShardSet) All() []*Shard {
	all := make([]*Shard, len(s.shards))
	for i, shard := range s.shards {
		all[i] = shard
	}
	return all
}

func (s *ShardSet) RandomShard() *Shard {
	return s.shards[random.Between(0, len(s.shards))]
}

func (s *ShardSet) addShard(i int) (err errs.Err) {
	autoIncrementOffset := i + 1
	dbName := fmt.Sprint(s.dbNamePrefix, autoIncrementOffset)
	s.shards[i], err = newShard(s, dbName, autoIncrementOffset)
	if err != nil {
		return
	}
	return
}

func newShard(s *ShardSet, dbName string, autoIncrementOffset int) (*Shard, errs.Err) {
	tls := s.tls
	if tls == "" {
		tls = "false"
	}
	connVars := ConnVariables{
		"autocommit":               "true",
		"clientFoundRows":          "true",
		"charset":                  "utf8mb4",
		"collation":                "utf8_unicode_ci",
		"auto_increment_increment": strconv.Itoa(s.maxShards),
		"auto_increment_offset":    strconv.Itoa(autoIncrementOffset),
		"sql_mode":                 "STRICT_ALL_TABLES",
		"tls":                      tls,
	}

	db, err := dbOpener(s.username, s.password, dbName, s.host, s.port, connVars)
	if err != nil {
		return nil, err
	}

	db.SetMaxOpenConns(s.maxConns)
	// db.SetMaxIdleConns(n)
	stdErr := db.Ping()
	if stdErr != nil {
		return nil, errs.Wrap(stdErr, nil)
	}
	return &Shard{DBName: dbName, db: db, sqlConn: db, BeginEndHandler: s.beginEndHandler, MetricsHandler: s.metricsHandler, log: s.log}, nil
}

func SetOpener(opener Opener) {
	if dbOpener != nil {
		panic("Opener already set - did you import two driver adapters?")
	}
	dbOpener = opener
}

type ConnVariables map[string]string

func (connVars ConnVariables) Join(sep string) string {
	kvps := make([]string, len(connVars))
	i := 0
	for param, val := range connVars {
		kvps[i] = param + "=" + val
		i += 1
	}
	return strings.Join(kvps, sep)
}

type Opener func(username, password, dbName, host string, port int, connVars ConnVariables) (*sql.DB, errs.Err)

var dbOpener Opener
<|endoftext|>